Skip to content

Instantly share code, notes, and snippets.

@adharris
Created November 28, 2012 19:52
Show Gist options
  • Save adharris/4163702 to your computer and use it in GitHub Desktop.
Save adharris/4163702 to your computer and use it in GitHub Desktop.
PostgreSQL demo of Array types using Golang
package main
import (
"database/sql"
"errors"
"fmt"
_ "github.com/bmizerany/pq"
"os"
"regexp"
"strings"
)
func main() {
db := dbConnect()
makeTestTables(db)
defer db.Close()
defer cleanup(db)
// Insert Some Data
db.Exec(`INSERT INTO array_test VALUES ('{"String1", "String2"}')`)
// arrays can be selected as strings...
dataString := selectAsString(db)
fmt.Println("SELECT as String:", dataString)
// Or by using array functions...
dataUnnest := selectUsingUnnest(db)
fmt.Println("SELECT using Unnest:", dataUnnest)
// Or by defining a scan type and parsing the return value
dataSlice := selectAsSlice(db)
fmt.Println("SELECT by parsing:", dataSlice)
// Arrays can be updated by replacing the entire array:
newArray := []interface{}{"String1", "String3", "String4", "String5"}
updateArray(db, newArray)
dataSlice = selectAsSlice(db)
fmt.Println("UPDATE entire array", dataSlice)
// or by appending / prepending value(s):
AppendToArray(db, "String6")
dataSlice = selectAsSlice(db)
fmt.Println("UPDATE with append:", dataSlice)
// or by replacing individual values:
ReplaceInArray(db, 2, "NULL")
dataSlice = selectAsSlice(db)
fmt.Println("UPDATE with replace:", dataSlice)
// Deleting by index requires slicing and is inefficient:
DeleteFromArray(db, 3)
dataSlice = selectAsSlice(db)
fmt.Println("UPDATE deleting index:", dataSlice)
}
// Arrays are serialized to strings {value, value...} by the database.
// these strings selected, updated and inserted like any string
func selectAsString(db *sql.DB) string {
row := db.QueryRow("SELECT data FROM array_test")
var asString string
err := row.Scan(&asString)
if err != nil {
panic(err)
}
return asString
}
// The UNNEST function expands an array into multiple rows. Each row
// can then be scanned individually.
func selectUsingUnnest(db *sql.DB) []string {
results := make([]string, 0)
rows, err := db.Query("SELECT UNNEST(data) FROM array_test")
if err != nil {
panic(err)
}
var scanString string
for rows.Next() {
rows.Scan(&scanString)
results = append(results, scanString)
}
return results
}
// By defining a wrapper type around a slice which implements
// sql.Scanner, we can scan the array directly into the type.
func selectAsSlice(db *sql.DB) StringSlice {
row := db.QueryRow("SELECT data FROM array_test")
var asSlice StringSlice
err := row.Scan(&asSlice)
if err != nil {
panic(err)
}
return asSlice
}
// Update an array by replacing the whole array with new values.
// This _could_ be done by serializing the StringSlice type using
// sql.driver.Valuer, but then we would have to validate the type
// of each value manually and format it for insert by hand. Instead,
// the ARRAY[...] format allows us to use query parameters to construct
// the array, ie ARRAY[$1, $2, $3], which then allows the database
// driver to coerce the variables into the right format for us.
func updateArray(db *sql.DB, array []interface{}) {
params := make([]string, 0, len(array))
for i := range array {
params = append(params, fmt.Sprintf("$%v", i+1))
}
query := fmt.Sprintf("UPDATE array_test SET data = ARRAY[%s]", strings.Join(params, ", "))
db.Exec(query, array...)
}
// The ARRAY_APPEND and ARRAY_PREPEND functions can be used to add single
// values to arrays. ARRAY_CAT combines two arrays. The || operator can
// do the same thing:
// SET data = data || <value>
// SET data = data || ARRAY[<value1>, <value2>]
func AppendToArray(db *sql.DB, value string) {
_, err := db.Exec("UPDATE array_test SET data = ARRAY_APPEND(data, $1)", value)
if err != nil {
panic(err)
}
}
// Arrays are 1-indexed. Individual elements can be used in expressions,
// updated, or selected by indexing the array.
func ReplaceInArray(db *sql.DB, index int, newValue string) {
_, err := db.Exec("UPDATE array_test SET data[$1] = $2", index, newValue)
if err != nil {
panic(err)
}
}
// Arrays support slice indexing:
// ARRAY['a', 'b', 'c'][1:2] == ARRAY['a', 'b']
// The ARRAY_UPPER function gets the length of an array for a specified dimension
// Deleting a value from an array amounts to slicing the array into two parts
// and combining them back together.
func DeleteFromArray(db *sql.DB, i int) {
_, err := db.Exec("UPDATE array_test SET data = array_cat(data[0:$1], data[$2:ARRAY_UPPER(data, 1) + 1])", i-1, i+1)
if err != nil {
panic(err)
}
}
type StringSlice []string
// Implements sql.Scanner for the String slice type
// Scanners take the database value (in this case as a byte slice)
// and sets the value of the type. Here we cast to a string and
// do a regexp based parse
func (s *StringSlice) Scan(src interface{}) error {
asBytes, ok := src.([]byte)
if !ok {
return error(errors.New("Scan source was not []bytes"))
}
asString := string(asBytes)
parsed := parseArray(asString)
(*s) = StringSlice(parsed)
return nil
}
// PARSING ARRAYS
// SEE http://www.postgresql.org/docs/9.1/static/arrays.html#ARRAYS-IO
// Arrays are output within {} and a delimiter, which is a comma for most
// postgres types (; for box)
//
// Individual values are surrounded by quotes:
// The array output routine will put double quotes around element values if
// they are empty strings, contain curly braces, delimiter characters,
// double quotes, backslashes, or white space, or match the word NULL.
// Double quotes and backslashes embedded in element values will be
// backslash-escaped. For numeric data types it is safe to assume that double
// quotes will never appear, but for textual data types one should be prepared
// to cope with either the presence or absence of quotes.
// construct a regexp to extract values:
var (
// unquoted array values must not contain: (" , \ { } whitespace NULL)
// and must be at least one char
unquotedChar = `[^",\\{}\s(NULL)]`
unquotedValue = fmt.Sprintf("(%s)+", unquotedChar)
// quoted array values are surrounded by double quotes, can be any
// character except " or \, which must be backslash escaped:
quotedChar = `[^"\\]|\\"|\\\\`
quotedValue = fmt.Sprintf("\"(%s)*\"", quotedChar)
// an array value may be either quoted or unquoted:
arrayValue = fmt.Sprintf("(?P<value>(%s|%s))", unquotedValue, quotedValue)
// Array values are separated with a comma IF there is more than one value:
arrayExp = regexp.MustCompile(fmt.Sprintf("((%s)(,)?)", arrayValue))
valueIndex int
)
// Find the index of the 'value' named expression
func init() {
for i, subexp := range arrayExp.SubexpNames() {
if subexp == "value" {
valueIndex = i
break
}
}
}
// Parse the output string from the array type.
// Regex used: (((?P<value>(([^",\\{}\s(NULL)])+|"([^"\\]|\\"|\\\\)*")))(,)?)
func parseArray(array string) []string {
results := make([]string, 0)
matches := arrayExp.FindAllStringSubmatch(array, -1)
for _, match := range matches {
s := match[valueIndex]
// the string _might_ be wrapped in quotes, so trim them:
s = strings.Trim(s, "\"")
results = append(results, s)
}
return results
}
// DB HELPERs
func dbConnect() *sql.DB {
datname := os.Getenv("PGDATABASE")
sslmode := os.Getenv("PGSSLMODE")
if datname == "" {
os.Setenv("PGDATABASE", "pqgotest")
}
if sslmode == "" {
os.Setenv("PGSSLMODE", "disable")
}
conn, err := sql.Open("postgres", "")
if err != nil {
panic(err)
}
return conn
}
// Create a table with an array type
// Can also use the syntax CREATE TABLE array_test (data varchar ARRAY)
func makeTestTables(db *sql.DB) {
_, err := db.Exec("CREATE TABLE array_test (data varchar[])")
if err != nil {
panic(err)
}
}
func cleanup(db *sql.DB) {
db.Exec("DROP TABLE array_test")
}
@penland365
Copy link

Man, nothing but thanks for this gist. It's going to save me almost an entire day.

@jaredfolkins
Copy link

Thank you for this!

@vinsonexx
Copy link

Thank you, it's useful!

@sescobb27
Copy link

Oh men, you are my hero, i spent few days looking for something like this, i'm going to test this for calling postgres functions from golang, thanks

@marksievers
Copy link

+1 Gold.

@mbivert
Copy link

mbivert commented Apr 10, 2014

indeed, very nice!

@collinglass
Copy link

+1

@satran
Copy link

satran commented Oct 10, 2014

Line 186 probably needs a fix. It was breaking on characters N, U and L.
Works fine with

unquotedChar  = `[^",\\{}\s]|NULL`

@iostrovok
Copy link

There is bad work any complex case for []text array.

@peterbe
Copy link

peterbe commented Oct 21, 2014

What about from StringSlice back to a string that can be inserted in SQL?

@kyleboyle
Copy link

Here is an alternative approach that only supports single dimension string arrays. It uses the csv parser lib to do the parsing after a bit of preprocessing. And saving to postgres is as easy as creating a string of the proper format:

// for replacing escaped quotes except if it is preceded by a literal backslash
//  eg "\\" should translate to a quoted element whose value is \

var quoteEscapeRegex = regexp.MustCompile(`([^\\]([\\]{2})*)\\"`)

// Scan convert to a slice of strings
// http://www.postgresql.org/docs/9.1/static/arrays.html#ARRAYS-IO
func (s *StringSlice) Scan(src interface{}) error {
    asBytes, ok := src.([]byte)
    if !ok {
        return error(errors.New("Scan source was not []bytes"))
    }
    str := string(asBytes)

    // change quote escapes for csv parser
    str = quoteEscapeRegex.ReplaceAllString(str, `$1""`)
    str = strings.Replace(str, `\\`, `\`, -1)
    // remove braces
    str = str[1 : len(str)-1]
    csvReader := csv.NewReader(strings.NewReader(str))

    slice, err := csvReader.Read()

    if err != nil {
        return err
    }

    (*s) = StringSlice(slice)

    return nil
}

func (s StringSlice) Value() (driver.Value, error) {
    // string escapes.
    // \ => \\\
    // " => \"
    for i, elem := range s {
        s[i] = `"` + strings.Replace(strings.Replace(elem, `\`, `\\\`, -1), `"`, `\"`, -1) + `"`
    }
    return "{" + strings.Join(s, ",") + "}", nil
}

Tests:

func TestStringSliceScan(t *testing.T) {
    var slice StringSlice

    err := slice.Scan([]byte(`{"12",45,"abc,\\\"d\\ef\\\\"}`))

    if err != nil {
        t.Errorf("Could not scan array, %v", err)
        return
    }

    if slice[0] != "12" || slice[2] != `abc,\"d\ef\\` {
        t.Errorf("Did not get expected slice contents")
    }
}

func TestStringSliceDbValue(t *testing.T) {
    slice := StringSlice([]string{`as"f\df`, "43", "}adsf"})

    val, err := slice.Value()
    if err != nil {
        t.Errorf("Could not convert to db string")
    }

    if str, ok := val.(string); ok {
        if `{"as\"f\\\df","43","}adsf"}` != str {
            t.Errorf("db value expecting %s got %s", `{"as\"f\\\df","43","}adsf"}`, str)
        }
    } else {
        t.Errorf("Could not convert %v to string for comparison", val)
    }
}

@rafroehlich2
Copy link

Well this was an awesome find. Thanks for saving me a large amount of time!

@tarikjn
Copy link

tarikjn commented Feb 22, 2015

An easier approach that works with all kind of nested arrays/composite types, if you only need to retrieve values from Postgres, is simply to use native Postgres functions to_json/array_to_json/row_to_json in your SELECT query and then json unmarshall the resulting string.

@nwoolls
Copy link

nwoolls commented Feb 26, 2015

@peterbe

type StringSlice []string

func (s StringSlice) Value() (driver.Value, error) {
    if len(s) == 0 {
        return nil, nil
    }

    var buffer bytes.Buffer

    buffer.WriteString("{")
    last := len(s) - 1
    for i, val := range s {
        buffer.WriteString(strconv.Quote(val))
        if i != last {
            buffer.WriteString(",")
        }
    }
    buffer.WriteString("}")

    return string(buffer.Bytes()), nil
}

@jsdir
Copy link

jsdir commented Apr 23, 2015

👍

@davidmz
Copy link

davidmz commented Jul 6, 2015

@kyleboyle, your Scan needs a branch for a '{}' case.

@teejayvanslyke
Copy link

Thanks! Quite helpful.

@mikeatlas-r7
Copy link

mikeatlas-r7 commented Feb 28, 2018

@kyleboyle @davidmz here's the Scan supporting the empty array case:

func (s *StringSlice) Scan(src interface{}) error {
	var str string
	if asBytes, ok := src.([]byte); ok {
		str = string(asBytes)
	} else if asStr, ok := src.(string); ok {
		str = string(asStr)
	} else {
		return error(errors.New(
			"StringSlice Scan source was not []bytes or string"))
	}

	// change quote escapes for csv parser
	str = quoteEscapeRegex.ReplaceAllString(str, `$1""`)
	str = strings.Replace(str, `\\`, `\`, -1)

	// remove braces
	str = str[1 : len(str)-1]

	// if the result was empty, just return empty slice
	if len(str) == 0 {
		(*s) = StringSlice([]string{})
		return nil
	}

	// otherwise, parse the list of values.
	csvReader := csv.NewReader(strings.NewReader(str))
	slice, err := csvReader.Read()
	if err != nil {
		return err
	}
	(*s) = StringSlice(slice)

	return nil
}

@vituchon
Copy link

Thanks a lot for sharing!! Let me show an example of working with arrays using pg.Array function, that leverages a little bit the mapping beteween go slices and psql arrays.


package main

import (
	"database/sql"
	"fmt"

	"log"

	"github.com/lib/pq"
)

const (
	host     = "localhost"
	port     = 5431
	user     = "user"
	password = "secret :P" // actually would be difficult to guess 
	dbname   = "testdb"
)

type SampleRow struct {
	Id           int
	MultiplesIds []int64
}

func ensureSampleTableExists(db *sql.DB) {
	_, err := db.Exec(`
	CREATE TABLE IF NOT EXISTS sample ( id serial NOT NULL,
	  multiples_ids bigint[],
	  CONSTRAINT sample_pk PRIMARY KEY (id)
	)`)
	if err != nil {
		panic(err)
	}
}

func main() {
	psqlInfo := fmt.Sprintf("host=%s port=%d user=%s "+
		"password=%s dbname=%s sslmode=disable",
		host, port, user, password, dbname)
	db, err := sql.Open("postgres", psqlInfo)
	if err != nil {
		panic(err)
	}
	defer db.Close()
	ensureSampleTableExists(db)

	/*var x []int64 // must be int64 to work when reading... compatibility issues i guess...
	db.QueryRow("SELECT ARRAY[235, 401]").Scan(pq.Array(&x))
	log.Println(x)*/
	insertRow(db, SampleRow{MultiplesIds: []int64{2, 3}})
	samples, err := getRows(db)
	if err != nil {
		log.Panicln(err)
	} else {
		log.Println(samples)
	}
}

func insertRow(Db *sql.DB, s SampleRow) (SampleRow, error) {
	err := Db.QueryRow(`
		INSERT INTO sample(multiples_ids)
	    VALUES ($1) RETURNING id`, pq.Array(s.MultiplesIds)).Scan(&s.Id)
	return s, err
}

func getRows(Db *sql.DB) ([]SampleRow, error) {
	rows, err := Db.Query(`
		SELECT id, multiples_ids
		FROM sample`)
	if err != nil {
		if err == sql.ErrNoRows {
			return nil, nil
		}
		return nil, err
	}
	defer rows.Close()
	samples := []SampleRow{}
	for rows.Next() {
		var sample SampleRow
		err = rows.Scan(&sample.Id, pq.Array(&sample.MultiplesIds))
		if err != nil {
			log.Panicln("Error while scanning feched results", err)
			return nil, err
		}
		samples = append(samples, sample)
	}

	return samples, nil
}


@Entrio
Copy link

Entrio commented Mar 10, 2021

How about an integer slice? (not int64)

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment