jochen
d8b4d20886
letta crash-loops on 'type "vector" does not exist': pgvector is not a trusted extension, so only the provider's superuser can create it, and the provisioner never did. A contribution may now name extensions; the provider creates each (IF NOT EXISTS, available ones only) in the consumer's database on every pass. Go per the standing rule for a TypeScript module that changes. letta asks for vector.
616 lines
20 KiB
Go
616 lines
20 KiB
Go
package main
|
|
|
|
// postgres's admin client — postgres's own code, living in the module (novox/hq ADR 0039). Both this
|
|
// module's tools and its provisioner use it, and nothing outside postgres does.
|
|
//
|
|
// SQL goes over the wire protocol (pgconn), in the simple query protocol: a statement is sent as it
|
|
// was given, the way `psql -c` sent it when this module was TypeScript and could take no driver. One
|
|
// boundary, Dialer, and every method is built on it — which is also what the tests replace.
|
|
|
|
import (
|
|
"context"
|
|
"errors"
|
|
"fmt"
|
|
"net"
|
|
"net/url"
|
|
"os"
|
|
"regexp"
|
|
"strconv"
|
|
"strings"
|
|
"sync"
|
|
"time"
|
|
"unicode/utf8"
|
|
|
|
"github.com/jackc/pgx/v5/pgconn"
|
|
)
|
|
|
|
// Reader is the login a caller's statement runs as (novox/hq issue 193). It may read every table and
|
|
// change nothing: `pg_read_all_data` and no other grant, and every transaction it opens is read-only
|
|
// by the server's own setting. A statement cannot climb out of a login the way it can out of a
|
|
// transaction wrapped around it as text: `COMMIT; DROP …` ended the old wrapper and ran the rest as
|
|
// the superuser, and even one read-only statement as a superuser can run a program on the server.
|
|
const Reader = "mesh_store_reader"
|
|
|
|
// readerOptions make the reader's session read-only from its first statement, before the role's own
|
|
// setting is read (PGOPTIONS, when this module shelled out to psql).
|
|
var readerOptions = map[string]string{
|
|
"default_transaction_read_only": "on",
|
|
"statement_timeout": "60s",
|
|
}
|
|
|
|
// maxResultBytes bounds what one read-only query may hand back, as psql's 16 MiB buffer did.
|
|
const maxResultBytes = 16 << 20
|
|
|
|
// Result is one statement's answer: the command tag's verb and the rows, each a column's text or nil.
|
|
type Result struct {
|
|
Command string
|
|
Fields []string
|
|
Rows [][]*string
|
|
}
|
|
|
|
// Maps is the rows keyed by their columns, as the tools return them.
|
|
func (r Result) Maps() []map[string]any {
|
|
out := make([]map[string]any, 0, len(r.Rows))
|
|
for _, row := range r.Rows {
|
|
m := map[string]any{}
|
|
for i, name := range r.Fields {
|
|
if i < len(row) && row[i] != nil {
|
|
m[name] = *row[i]
|
|
} else {
|
|
m[name] = nil
|
|
}
|
|
}
|
|
out = append(out, m)
|
|
}
|
|
return out
|
|
}
|
|
|
|
// Session is one open connection: Run sends SQL in the simple protocol and answers every statement's
|
|
// result, in order.
|
|
type Session interface {
|
|
Run(ctx context.Context, sql string) ([]Result, error)
|
|
Close()
|
|
}
|
|
|
|
// Login is who a session connects as, to which database, with which session settings.
|
|
type Login struct {
|
|
Database string
|
|
User string
|
|
Password string
|
|
Options map[string]string
|
|
}
|
|
|
|
// Dialer opens a session. The real one is pgconn; a test's records what it was asked.
|
|
type Dialer func(ctx context.Context, l Login) (Session, error)
|
|
|
|
// Conn is where the server is and who administers it.
|
|
type Conn struct {
|
|
Host string
|
|
Port int
|
|
User string
|
|
Password string
|
|
SSLMode string
|
|
// ReaderPassword is the read-only login's password, which the mesh mints for this module
|
|
// (`own-secrets.reader`). Empty when the mesh has not delivered it: then a caller's statement is
|
|
// refused, never run as the admin (novox/hq issue 193).
|
|
ReaderPassword string
|
|
}
|
|
|
|
// Client is postgres's admin client.
|
|
type Client struct {
|
|
conn Conn
|
|
dial Dialer
|
|
|
|
readerMu sync.Mutex
|
|
readerReady bool
|
|
}
|
|
|
|
// NewClient is a client over a dialer; nil dials the server for real.
|
|
func NewClient(conn Conn, dial Dialer) *Client {
|
|
if dial == nil {
|
|
dial = pgDialer(conn)
|
|
}
|
|
return &Client{conn: conn, dial: dial}
|
|
}
|
|
|
|
var unfilled = regexp.MustCompile(`^\$\{[^}]*\}$`)
|
|
|
|
// ClientFromEnv builds the client from the module's words. MESH_POSTGRES_* first (the documented
|
|
// names), then the MESH_PROVISION_* keys the manifest sets. Fails without a host and an admin password.
|
|
func ClientFromEnv(env func(string) string) (*Client, error) {
|
|
var u *url.URL
|
|
if raw := env("MESH_PROVISION_POSTGRES"); raw != "" {
|
|
if parsed, err := url.Parse(raw); err == nil {
|
|
u = parsed
|
|
}
|
|
}
|
|
host := env("MESH_POSTGRES_HOST")
|
|
if host == "" && u != nil {
|
|
host = u.Hostname()
|
|
}
|
|
// MESH_PROVISION_POSTGRES_PORT is the seat's twin (mesh-controller's internal/envfile.Placed
|
|
// pattern): which port this machine actually put mesh-store at, when that differs from the
|
|
// connection string's. Empty, or a placeholder the mesh never filled, adds nothing.
|
|
seatPort := strings.TrimSpace(env("MESH_PROVISION_POSTGRES_PORT"))
|
|
if unfilled.MatchString(seatPort) {
|
|
seatPort = ""
|
|
}
|
|
portSource := firstOf(env("MESH_POSTGRES_PORT"), seatPort)
|
|
if portSource == "" && u != nil {
|
|
portSource = u.Port()
|
|
}
|
|
port, err := strconv.Atoi(portSource)
|
|
if err != nil || port == 0 {
|
|
port = 5432
|
|
}
|
|
user := env("MESH_POSTGRES_USER")
|
|
if user == "" && u != nil && u.User != nil {
|
|
user = u.User.Username()
|
|
}
|
|
if user == "" {
|
|
user = "postgres"
|
|
}
|
|
password := firstOf(env("MESH_POSTGRES_PASSWORD"), readSecretFile(env("MESH_PROVISION_PASSWORD_FILE")))
|
|
if host == "" || password == "" {
|
|
return nil, errors.New("postgres host or admin password is not set — postgres's own code cannot reach the server")
|
|
}
|
|
sslmode := "prefer"
|
|
if u != nil && u.Query().Get("sslmode") != "" {
|
|
sslmode = u.Query().Get("sslmode")
|
|
}
|
|
reader := firstOf(env("MESH_POSTGRES_READER_PASSWORD"), readSecretFile(env("MESH_POSTGRES_READER_PASSWORD_FILE")))
|
|
return NewClient(Conn{Host: host, Port: port, User: user, Password: password, SSLMode: sslmode, ReaderPassword: reader}, nil), nil
|
|
}
|
|
|
|
func firstOf(values ...string) string {
|
|
for _, v := range values {
|
|
if v != "" {
|
|
return v
|
|
}
|
|
}
|
|
return ""
|
|
}
|
|
|
|
func readSecretFile(path string) string {
|
|
if path == "" {
|
|
return ""
|
|
}
|
|
b, err := os.ReadFile(path)
|
|
if err != nil {
|
|
return ""
|
|
}
|
|
return strings.TrimSpace(string(b))
|
|
}
|
|
|
|
// ---- the boundary -------------------------------------------------------------------------------
|
|
|
|
// as runs SQL as one login against one database and answers the last statement's result.
|
|
func (c *Client) as(ctx context.Context, l Login, sql string) (Result, error) {
|
|
s, err := c.dial(ctx, l)
|
|
if err != nil {
|
|
return Result{}, err
|
|
}
|
|
defer s.Close()
|
|
results, err := s.Run(ctx, sql)
|
|
if err != nil {
|
|
return Result{}, err
|
|
}
|
|
if len(results) == 0 {
|
|
return Result{}, nil
|
|
}
|
|
// The last statement that answered rows, or the last statement: what psql -c printed last.
|
|
for i := len(results) - 1; i >= 0; i-- {
|
|
if len(results[i].Fields) > 0 {
|
|
return results[i], nil
|
|
}
|
|
}
|
|
return results[len(results)-1], nil
|
|
}
|
|
|
|
// Query runs SQL as the admin against a database ("postgres" when empty).
|
|
func (c *Client) Query(ctx context.Context, database, sql string) (Result, error) {
|
|
if database == "" {
|
|
database = "postgres"
|
|
}
|
|
return c.as(ctx, Login{Database: database, User: c.conn.User, Password: c.conn.Password}, sql)
|
|
}
|
|
|
|
func (c *Client) exists(ctx context.Context, sql string) (bool, error) {
|
|
r, err := c.Query(ctx, "", sql)
|
|
return len(r.Rows) > 0, err
|
|
}
|
|
|
|
// ---- what the provisioner does -----------------------------------------------------------------
|
|
|
|
// RoleStatement is the DDL that makes or re-sets a consumer's login. VALID UNTIL 'infinity': an
|
|
// expired password is refused like a wrong one, so the check the provisioner runs would report it
|
|
// lost, and only clearing the expiry makes applying it again work.
|
|
func RoleStatement(exists bool, role, password string) string {
|
|
verb := "CREATE"
|
|
if exists {
|
|
verb = "ALTER"
|
|
}
|
|
return fmt.Sprintf("%s ROLE %s WITH LOGIN PASSWORD %s VALID UNTIL 'infinity'", verb, Ident(role), Literal(password))
|
|
}
|
|
|
|
// CreateDatabaseAndRole makes a login role and a database it owns, idempotently. A database that
|
|
// exists is never recreated, and nothing here drops anything.
|
|
func (c *Client) CreateDatabaseAndRole(ctx context.Context, database, role, password string) error {
|
|
has, err := c.exists(ctx, "SELECT 1 FROM pg_roles WHERE rolname = "+Literal(role))
|
|
if err != nil {
|
|
return err
|
|
}
|
|
if _, err := c.Query(ctx, "", RoleStatement(has, role, password)); err != nil {
|
|
return err
|
|
}
|
|
has, err = c.exists(ctx, "SELECT 1 FROM pg_database WHERE datname = "+Literal(database))
|
|
if err != nil {
|
|
return err
|
|
}
|
|
if !has {
|
|
if _, err := c.Query(ctx, "", fmt.Sprintf("CREATE DATABASE %s OWNER %s", Ident(database), Ident(role))); err != nil {
|
|
return err
|
|
}
|
|
}
|
|
_, err = c.Query(ctx, "", fmt.Sprintf("GRANT ALL PRIVILEGES ON DATABASE %s TO %s", Ident(database), Ident(role)))
|
|
return err
|
|
}
|
|
|
|
// Extensions reads a contribution's `extensions`: absent is none; otherwise a list of names.
|
|
func Extensions(values map[string]any) ([]string, error) {
|
|
raw, ok := values["extensions"]
|
|
if !ok || raw == nil {
|
|
return nil, nil
|
|
}
|
|
list, ok := raw.([]any)
|
|
if !ok {
|
|
return nil, fmt.Errorf("extensions must be a list of extension names, not %T", raw)
|
|
}
|
|
seen := map[string]bool{}
|
|
var out []string
|
|
for _, v := range list {
|
|
name, ok := v.(string)
|
|
name = strings.TrimSpace(name)
|
|
if !ok || name == "" {
|
|
return nil, fmt.Errorf("extensions must be a list of extension names; %v is not one", v)
|
|
}
|
|
if !seen[name] {
|
|
seen[name] = true
|
|
out = append(out, name)
|
|
}
|
|
}
|
|
return out, nil
|
|
}
|
|
|
|
// ExtensionStatement is the DDL that installs one extension in the database it is run in. IF NOT
|
|
// EXISTS: run on every pass, and a second run changes nothing. There is no statement that removes one.
|
|
func ExtensionStatement(name string) string {
|
|
return "CREATE EXTENSION IF NOT EXISTS " + Ident(name)
|
|
}
|
|
|
|
// Unavailable is the names asked for that the server does not offer.
|
|
func Unavailable(want []string, available map[string]bool) []string {
|
|
var missing []string
|
|
for _, name := range want {
|
|
if !available[name] {
|
|
missing = append(missing, name)
|
|
}
|
|
}
|
|
return missing
|
|
}
|
|
|
|
// EnsureExtensions installs each named extension in the database, as the admin, connected to that
|
|
// database — most extensions (pgvector among them) are not trusted, so the consumer that owns the
|
|
// database cannot install them itself. Only names the server lists in pg_available_extensions are
|
|
// asked for; any other is refused, by name, before anything runs. An extension is never dropped.
|
|
func (c *Client) EnsureExtensions(ctx context.Context, database string, want []string) error {
|
|
if len(want) == 0 {
|
|
return nil
|
|
}
|
|
r, err := c.Query(ctx, "", "SELECT name FROM pg_available_extensions")
|
|
if err != nil {
|
|
return err
|
|
}
|
|
available := map[string]bool{}
|
|
for _, row := range r.Rows {
|
|
if len(row) > 0 && row[0] != nil {
|
|
available[*row[0]] = true
|
|
}
|
|
}
|
|
if missing := Unavailable(want, available); len(missing) > 0 {
|
|
return fmt.Errorf("database %s asks for extension(s) %s, which this server does not offer "+
|
|
"(not in pg_available_extensions); refused, nothing installed", database, strings.Join(quoted(missing), ", "))
|
|
}
|
|
for _, name := range want {
|
|
if _, err := c.Query(ctx, database, ExtensionStatement(name)); err != nil {
|
|
return fmt.Errorf("extension %q in %s: %w", name, database, err)
|
|
}
|
|
}
|
|
return nil
|
|
}
|
|
|
|
func quoted(names []string) []string {
|
|
out := make([]string, len(names))
|
|
for i, n := range names {
|
|
out[i] = strconv.Quote(n)
|
|
}
|
|
return out
|
|
}
|
|
|
|
// lostCodes are the server's ways of saying a login, its password or its database is wrong or gone:
|
|
// invalid_password, invalid_authorization_specification (no such role, not permitted to log in),
|
|
// invalid_catalog_name (no such database), insufficient_privilege (no CONNECT).
|
|
var lostCodes = map[string]bool{"28P01": true, "28000": true, "3D000": true, "42501": true}
|
|
|
|
// IsLost says whether an error is the server saying the credential is wrong or gone.
|
|
func IsLost(err error) bool {
|
|
var pg *pgconn.PgError
|
|
return errors.As(err, &pg) && lostCodes[pg.Code]
|
|
}
|
|
|
|
// Holds says whether `role` can log in to `database` with exactly `password` and finds every wanted
|
|
// extension installed there: the consumer's own view, checked by connecting as it. Read-only. false
|
|
// only when the server says so; an unreachable server is an error, because being unable to ask is
|
|
// not evidence of loss (novox/hq issue 120).
|
|
func (c *Client) Holds(ctx context.Context, database, role, password string, extensions []string) (bool, error) {
|
|
ctx, cancel := context.WithTimeout(ctx, 20*time.Second)
|
|
defer cancel()
|
|
r, err := c.as(ctx, Login{Database: database, User: role, Password: password}, "SELECT extname FROM pg_extension")
|
|
if err != nil {
|
|
if IsLost(err) {
|
|
return false, nil
|
|
}
|
|
return false, err
|
|
}
|
|
installed := map[string]bool{}
|
|
for _, row := range r.Rows {
|
|
if len(row) > 0 && row[0] != nil {
|
|
installed[*row[0]] = true
|
|
}
|
|
}
|
|
return len(Unavailable(extensions, installed)) == 0, nil
|
|
}
|
|
|
|
// LockRole withdraws a consumer without destroying anything (novox/hq issue 241): its login can no
|
|
// longer log in and its open connections are ended, and its database stays exactly as it was, under
|
|
// its own name. A consumer that comes back is given the same database — create sets LOGIN again.
|
|
func (c *Client) LockRole(ctx context.Context, role string) error {
|
|
has, err := c.exists(ctx, "SELECT 1 FROM pg_roles WHERE rolname = "+Literal(role))
|
|
if err != nil || !has {
|
|
return err
|
|
}
|
|
if _, err := c.Query(ctx, "", fmt.Sprintf("ALTER ROLE %s NOLOGIN", Ident(role))); err != nil {
|
|
return err
|
|
}
|
|
_, err = c.Query(ctx, "", "SELECT pg_terminate_backend(pid) FROM pg_stat_activity WHERE usename = "+
|
|
Literal(role)+" AND pid <> pg_backend_pid()")
|
|
return err
|
|
}
|
|
|
|
// ---- what the tools do --------------------------------------------------------------------------
|
|
|
|
// RetireDatabase takes a database out of service on purpose: renamed aside to
|
|
// `<name>_deleted_<date>` and its owner locked. Never a drop — the data stays on the server under the
|
|
// new name until a person removes it by hand. Answers the name it now has.
|
|
func (c *Client) RetireDatabase(ctx context.Context, database string, now time.Time) (string, error) {
|
|
found, err := c.Query(ctx, "", "SELECT pg_get_userbyid(datdba) AS owner FROM pg_database WHERE datname = "+Literal(database))
|
|
if err != nil {
|
|
return "", err
|
|
}
|
|
if len(found.Rows) == 0 {
|
|
return "", fmt.Errorf("no database named %s", database)
|
|
}
|
|
owner := ""
|
|
if row := found.Rows[0]; len(row) > 0 && row[0] != nil {
|
|
owner = *row[0]
|
|
}
|
|
aside := RetiredName(database, now)
|
|
taken, err := c.exists(ctx, "SELECT 1 FROM pg_database WHERE datname = "+Literal(aside))
|
|
if err != nil {
|
|
return "", err
|
|
}
|
|
if taken {
|
|
return "", fmt.Errorf("%s already exists; retire it by hand first", aside)
|
|
}
|
|
if owner != "" && owner != "postgres" {
|
|
if _, err := c.Query(ctx, "", fmt.Sprintf("ALTER ROLE %s NOLOGIN", Ident(owner))); err != nil {
|
|
return "", err
|
|
}
|
|
}
|
|
if _, err := c.Query(ctx, "", "SELECT pg_terminate_backend(pid) FROM pg_stat_activity WHERE datname = "+
|
|
Literal(database)+" AND pid <> pg_backend_pid()"); err != nil {
|
|
return "", err
|
|
}
|
|
if _, err := c.Query(ctx, "", fmt.Sprintf("ALTER DATABASE %s RENAME TO %s", Ident(database), Ident(aside))); err != nil {
|
|
return "", err
|
|
}
|
|
return aside, nil
|
|
}
|
|
|
|
// Database is one row of the listing.
|
|
type Database struct {
|
|
Name string `json:"name"`
|
|
SizeBytes int64 `json:"sizeBytes"`
|
|
}
|
|
|
|
// ListDatabases is the non-template databases with their size.
|
|
func (c *Client) ListDatabases(ctx context.Context) ([]Database, error) {
|
|
r, err := c.Query(ctx, "", "SELECT datname, pg_database_size(datname) AS size FROM pg_database WHERE datistemplate = false ORDER BY datname")
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
out := []Database{}
|
|
for _, row := range r.Rows {
|
|
if len(row) < 2 || row[0] == nil {
|
|
continue
|
|
}
|
|
d := Database{Name: *row[0]}
|
|
if row[1] != nil {
|
|
d.SizeBytes, _ = strconv.ParseInt(*row[1], 10, 64)
|
|
}
|
|
out = append(out, d)
|
|
}
|
|
return out, nil
|
|
}
|
|
|
|
// ReaderStatements are the DDL that make the read-only login, every attribute stated so an existing
|
|
// role someone widened is narrowed again.
|
|
func ReaderStatements(exists bool, password string) []string {
|
|
verb := "CREATE"
|
|
if exists {
|
|
verb = "ALTER"
|
|
}
|
|
return []string{
|
|
fmt.Sprintf("%s ROLE %s WITH LOGIN NOSUPERUSER NOCREATEDB NOCREATEROLE NOREPLICATION "+
|
|
"NOBYPASSRLS INHERIT PASSWORD %s VALID UNTIL 'infinity'", verb, Ident(Reader), Literal(password)),
|
|
"GRANT pg_read_all_data TO " + Ident(Reader),
|
|
"ALTER ROLE " + Ident(Reader) + " SET default_transaction_read_only = on",
|
|
"ALTER ROLE " + Ident(Reader) + " SET statement_timeout = '60s'",
|
|
}
|
|
}
|
|
|
|
var errReaderMissing = errors.New("the read-only login's password was not delivered (own-secrets.reader, " +
|
|
"MESH_POSTGRES_READER_PASSWORD_FILE), so the statement is refused rather than run as the admin (novox/hq issue 193)")
|
|
|
|
// EnsureReader makes the read-only login, idempotently, with the password the mesh minted — as the
|
|
// admin, because only the admin can make a role. Once per process; a failure is asked again next call.
|
|
func (c *Client) EnsureReader(ctx context.Context) error {
|
|
if c.conn.ReaderPassword == "" {
|
|
return errReaderMissing
|
|
}
|
|
c.readerMu.Lock()
|
|
defer c.readerMu.Unlock()
|
|
if c.readerReady {
|
|
return nil
|
|
}
|
|
has, err := c.exists(ctx, "SELECT 1 FROM pg_roles WHERE rolname = "+Literal(Reader))
|
|
if err != nil {
|
|
return err
|
|
}
|
|
for _, sql := range ReaderStatements(has, c.conn.ReaderPassword) {
|
|
if _, err := c.Query(ctx, "", sql); err != nil {
|
|
return err
|
|
}
|
|
}
|
|
c.readerReady = true
|
|
return nil
|
|
}
|
|
|
|
var firstWord = regexp.MustCompile(`^\s*([A-Za-z]+)`)
|
|
|
|
// ReadOnlyQuery runs a caller's statement against a named database as the read-only login, for the
|
|
// postgres_query tool and the store seat's `query` verb (novox/hq ADR 0159, issue 193).
|
|
//
|
|
// **Read-only by the login, not by text around the statement.** The statement is sent as it was
|
|
// given, as the reader, whose role can write nothing and whose session the server makes read-only.
|
|
// Never as the admin: without the reader's password the call is refused.
|
|
func (c *Client) ReadOnlyQuery(ctx context.Context, database, sql string) (Result, error) {
|
|
if err := c.EnsureReader(ctx); err != nil {
|
|
return Result{}, err
|
|
}
|
|
r, err := c.as(ctx, Login{Database: database, User: Reader, Password: c.conn.ReaderPassword, Options: readerOptions}, sql)
|
|
if err != nil {
|
|
return Result{}, err
|
|
}
|
|
size := 0
|
|
for _, row := range r.Rows {
|
|
for _, v := range row {
|
|
if v != nil {
|
|
size += len(*v)
|
|
}
|
|
}
|
|
}
|
|
if size > maxResultBytes {
|
|
return Result{}, fmt.Errorf("the result is larger than %d MiB; narrow the query", maxResultBytes>>20)
|
|
}
|
|
r.Command = ""
|
|
if m := firstWord.FindStringSubmatch(sql); m != nil {
|
|
r.Command = strings.ToUpper(m[1])
|
|
}
|
|
return r, nil
|
|
}
|
|
|
|
// ---- quoting and names -------------------------------------------------------------------------
|
|
|
|
// Ident quotes a SQL identifier: double quotes, internal ones doubled.
|
|
func Ident(id string) string { return `"` + strings.ReplaceAll(id, `"`, `""`) + `"` }
|
|
|
|
// Literal quotes a SQL string literal: single quotes, internal ones doubled (standard_conforming_strings).
|
|
func Literal(v string) string { return "'" + strings.ReplaceAll(v, "'", "''") + "'" }
|
|
|
|
// RetiredName is `<name>_deleted_<yyyymmdd>`, within postgres's 63 bytes.
|
|
func RetiredName(database string, now time.Time) string {
|
|
suffix := "_deleted_" + now.UTC().Format("20060102")
|
|
keep := 63 - len(suffix)
|
|
if len(database) > keep {
|
|
database = database[:keep]
|
|
for !utf8.ValidString(database) {
|
|
database = database[:len(database)-1]
|
|
}
|
|
}
|
|
return database + suffix
|
|
}
|
|
|
|
// ---- the real dialer ---------------------------------------------------------------------------
|
|
|
|
type pgSession struct{ c *pgconn.PgConn }
|
|
|
|
func (s pgSession) Run(ctx context.Context, sql string) ([]Result, error) {
|
|
results, err := s.c.Exec(ctx, sql).ReadAll()
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
out := make([]Result, 0, len(results))
|
|
for _, r := range results {
|
|
if r.Err != nil {
|
|
return nil, r.Err
|
|
}
|
|
res := Result{Command: strings.SplitN(r.CommandTag.String(), " ", 2)[0]}
|
|
for _, f := range r.FieldDescriptions {
|
|
res.Fields = append(res.Fields, f.Name)
|
|
}
|
|
for _, row := range r.Rows {
|
|
cells := make([]*string, len(row))
|
|
for i, v := range row {
|
|
if v != nil {
|
|
s := string(v)
|
|
cells[i] = &s
|
|
}
|
|
}
|
|
res.Rows = append(res.Rows, cells)
|
|
}
|
|
out = append(out, res)
|
|
}
|
|
return out, nil
|
|
}
|
|
|
|
func (s pgSession) Close() {
|
|
ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
|
|
defer cancel()
|
|
_ = s.c.Close(ctx)
|
|
}
|
|
|
|
func pgDialer(conn Conn) Dialer {
|
|
return func(ctx context.Context, l Login) (Session, error) {
|
|
u := url.URL{
|
|
Scheme: "postgres",
|
|
User: url.UserPassword(l.User, l.Password),
|
|
Host: net.JoinHostPort(conn.Host, strconv.Itoa(conn.Port)),
|
|
Path: "/" + l.Database,
|
|
RawQuery: url.Values{"sslmode": {firstOf(conn.SSLMode, "prefer")}, "connect_timeout": {"10"}}.Encode(),
|
|
}
|
|
cfg, err := pgconn.ParseConfig(u.String())
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
for k, v := range l.Options {
|
|
cfg.RuntimeParams[k] = v
|
|
}
|
|
c, err := pgconn.ConnectConfig(ctx, cfg)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
return pgSession{c: c}, nil
|
|
}
|
|
}
|