Last active
December 25, 2015 04:49
-
-
Save fd/6920174 to your computer and use it in GitHub Desktop.
This is a wrapper for the `sql.Driver` in `github.com/lib/pq` which allows you to set the schema search path for each connection in the connection pool.
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
package postgres | |
import ( | |
"database/sql" | |
"database/sql/driver" | |
"fmt" | |
"github.com/lib/pq" | |
"strings" | |
) | |
type drv struct{} | |
func init() { | |
sql.Register("pq-schema", &drv{}) | |
} | |
func (drv *drv) Open(data_source string) (driver.Conn, error) { | |
conn, err := pq.Open(data_source) | |
if err != nil { | |
return nil, err | |
} | |
data_source, schemas, err := parseOpts(data_source) | |
if err != nil { | |
conn.Close() | |
return nil, err | |
} | |
if schemas == "" { | |
return conn, nil | |
} | |
execer, ok := conn.(driver.Execer) | |
if !ok { | |
conn.Close() | |
return nil, fmt.Errorf("Expected driver to be an driver.Execer") | |
} | |
_, err = execer.Exec("SET search_path TO "+schemas+";", nil) | |
if err != nil { | |
return nil, err | |
} | |
return conn, nil | |
} | |
func parseOpts(in string) (out, schemas string, err error) { | |
if len(in) == 0 { | |
return | |
} | |
ops := []string{} | |
ps := strings.Split(in, " ") | |
for _, p := range ps { | |
kv := strings.Split(p, "=") | |
if len(kv) < 2 { | |
return "", "", fmt.Errorf("invalid option: %q", p) | |
} | |
if kv[0] == "search_path" { | |
schemas = kv[1] | |
} else { | |
ops = append(ops, p) | |
} | |
} | |
out = strings.Join(ops, " ") | |
schemas = strings.Replace(schemas, ";", "", -1) | |
return | |
} |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment