feat: Managed databases with any accessible server (#3421)

Managed databases only work using database servers hosted by sqlc Cloud. I'm going to add support for using managed databases with any database server you can access.

Database servers will be configured in a separate servers collection in the configuration file.

Each entry will have a mandatory uri field and an optional name. I may add an optional engine field if the URI scheme isn't enough to infer the engine type.

When using a database server not hosted by sqlc Cloud, sqlc will automatically create databases as needed based on the schema for the associated query set. All operations against these databases will be read-only, so they can be created once and re-used.

Author: kyleconroy

Makefile

@@ -9,6 +9,9 @@ install:
 test:
 	go test ./...
 
+test-managed:
+	MYSQL_SERVER_URI="invalid" POSTGRESQL_SERVER_URI="postgres://postgres:mysecretpassword@localhost:5432/postgres" go test -v ./...
+
 vet:
 	go vet ./...
 

internal/cmd/options.go

@@ -7,12 +7,14 @@ import (
 )
 
 type Options struct {
-	Env          Env
-	Stderr       io.Writer
-	MutateConfig func(*config.Config)
+	Env    Env
+	Stderr io.Writer
 	// TODO: Move these to a command-specific struct
 	Tags    []string
 	Against string
+
+	// Testing only
+	MutateConfig func(*config.Config)
 }
 
 func (o *Options) ReadConfig(dir, filename string) (string, *config.Config, error) {

internal/compiler/engine.go

@@ -6,13 +6,12 @@ import (
 
 	"github.com/sqlc-dev/sqlc/internal/analyzer"
 	"github.com/sqlc-dev/sqlc/internal/config"
+	"github.com/sqlc-dev/sqlc/internal/dbmanager"
 	"github.com/sqlc-dev/sqlc/internal/engine/dolphin"
 	"github.com/sqlc-dev/sqlc/internal/engine/postgresql"
 	pganalyze "github.com/sqlc-dev/sqlc/internal/engine/postgresql/analyzer"
 	"github.com/sqlc-dev/sqlc/internal/engine/sqlite"
 	"github.com/sqlc-dev/sqlc/internal/opts"
-	"github.com/sqlc-dev/sqlc/internal/quickdb"
-	pb "github.com/sqlc-dev/sqlc/internal/quickdb/v1"
 	"github.com/sqlc-dev/sqlc/internal/sql/catalog"
 )
 
@@ -23,7 +22,7 @@ type Compiler struct {
 	parser   Parser
 	result   *Result
 	analyzer analyzer.Analyzer
-	client   pb.QuickClient
+	client   dbmanager.Client
 
 	schema []string
 }
@@ -32,10 +31,7 @@ func NewCompiler(conf config.SQL, combo config.CombinedSettings) (*Compiler, err
 	c := &Compiler{conf: conf, combo: combo}
 
 	if conf.Database != nil && conf.Database.Managed {
-		client, err := quickdb.NewClientFromConfig(combo.Global.Cloud)
-		if err != nil {
-			return nil, fmt.Errorf("client error: %w", err)
-		}
+		client := dbmanager.NewClient(combo.Global.Servers)
 		c.client = client
 	}
 
@@ -89,4 +85,7 @@ func (c *Compiler) Close(ctx context.Context) {
 	if c.analyzer != nil {
 		c.analyzer.Close(ctx)
 	}
+	if c.client != nil {
+		c.client.Close(ctx)
+	}
 }

internal/config/config.go

@@ -59,13 +59,20 @@ const (
 type Config struct {
 	Version   string               `json:"version" yaml:"version"`
 	Cloud     Cloud                `json:"cloud" yaml:"cloud"`
+	Servers   []Server             `json:"servers" yaml:"servers"`
 	SQL       []SQL                `json:"sql" yaml:"sql"`
 	Overrides Overrides            `json:"overrides,omitempty" yaml:"overrides"`
 	Plugins   []Plugin             `json:"plugins" yaml:"plugins"`
 	Rules     []Rule               `json:"rules" yaml:"rules"`
 	Options   map[string]yaml.Node `json:"options" yaml:"options"`
 }
 
+type Server struct {
+	Name   string `json:"name,omitempty" yaml:"name"`
+	Engine Engine `json:"engine,omitempty" yaml:"engine"`
+	URI    string `json:"uri" yaml:"uri"`
+}
+
 type Database struct {
 	URI     string `json:"uri" yaml:"uri"`
 	Managed bool   `json:"managed" yaml:"managed"`

internal/dbmanager/client.go

@@ -0,0 +1,134 @@
+package dbmanager
+
+import (
+	"context"
+	"fmt"
+	"hash/fnv"
+	"io"
+	"net/url"
+	"strings"
+
+	"github.com/jackc/pgx/v5"
+	"golang.org/x/sync/singleflight"
+
+	"github.com/sqlc-dev/sqlc/internal/config"
+	"github.com/sqlc-dev/sqlc/internal/pgx/poolcache"
+	"github.com/sqlc-dev/sqlc/internal/shfmt"
+)
+
+type CreateDatabaseRequest struct {
+	Engine     string
+	Migrations []string
+}
+
+type CreateDatabaseResponse struct {
+	Uri string
+}
+
+type Client interface {
+	CreateDatabase(context.Context, *CreateDatabaseRequest) (*CreateDatabaseResponse, error)
+	Close(context.Context)
+}
+
+var flight singleflight.Group
+
+type ManagedClient struct {
+	cache    *poolcache.Cache
+	replacer *shfmt.Replacer
+	servers  []config.Server
+}
+
+func dbid(migrations []string) string {
+	h := fnv.New64()
+	for _, query := range migrations {
+		io.WriteString(h, query)
+	}
+	return fmt.Sprintf("%x", h.Sum(nil))
+}
+
+func (m *ManagedClient) CreateDatabase(ctx context.Context, req *CreateDatabaseRequest) (*CreateDatabaseResponse, error) {
+	hash := dbid(req.Migrations)
+	name := fmt.Sprintf("sqlc_managed_%s", hash)
+
+	var base string
+	for _, server := range m.servers {
+		if server.Engine == config.EnginePostgreSQL {
+			base = server.URI
+			break
+		}
+	}
+
+	if strings.TrimSpace(base) == "" {
+		return nil, fmt.Errorf("no PostgreSQL database server found")
+	}
+
+	serverUri := m.replacer.Replace(base)
+	pool, err := m.cache.Open(ctx, serverUri)
+	if err != nil {
+		return nil, err
+	}
+
+	uri, err := url.Parse(serverUri)
+	if err != nil {
+		return nil, err
+	}
+	uri.Path = name
+
+	key := uri.String()
+	_, err, _ = flight.Do(key, func() (interface{}, error) {
+		// TODO: Use a parameterized query
+		row := pool.QueryRow(ctx,
+			fmt.Sprintf(`SELECT datname FROM pg_database WHERE datname = '%s'`, name))
+
+		var datname string
+		if err := row.Scan(&datname); err == nil {
+			return nil, nil
+		}
+
+		if _, err := pool.Exec(ctx, fmt.Sprintf(`CREATE DATABASE "%s"`, name)); err != nil {
+			return nil, err
+		}
+
+		conn, err := pgx.Connect(ctx, uri.String())
+		if err != nil {
+			return nil, fmt.Errorf("connect %s: %s", name, err)
+		}
+		defer conn.Close(ctx)
+
+		var migrationErr error
+		for _, q := range req.Migrations {
+			if len(strings.TrimSpace(q)) == 0 {
+				continue
+			}
+			if _, err := conn.Exec(ctx, q); err != nil {
+				migrationErr = fmt.Errorf("%s: %s", q, err)
+				break
+			}
+		}
+
+		if migrationErr != nil {
+			pool.Exec(ctx, fmt.Sprintf(`DROP DATABASE "%s" IF EXISTS WITH (FORCE)`, name))
+			return nil, migrationErr
+		}
+
+		return nil, nil
+	})
+
+	if err != nil {
+		return nil, err
+	}
+
+	return &CreateDatabaseResponse{Uri: key}, err
+}
+
+func (m *ManagedClient) Close(ctx context.Context) {
+	m.cache.Close()
+}
+
+func NewClient(servers []config.Server) *ManagedClient {
+	return &ManagedClient{
+		cache:    poolcache.New(),
+		servers:  servers,
+		replacer: shfmt.NewReplacer(nil),
+	}
+}

internal/endtoend/endtoend_test.go

@@ -120,21 +120,28 @@ func TestReplay(t *testing.T) {
 		"managed-db": {
 			Mutate: func(t *testing.T, path string) func(*config.Config) {
 				return func(c *config.Config) {
+					c.Servers = []config.Server{
+						{
+							Name:   "postgres",
+							Engine: config.EnginePostgreSQL,
+							URI:    local.PostgreSQLServer(),
+						},
+
+						{
+							Name:   "mysql",
+							Engine: config.EngineMySQL,
+							URI:    local.MySQLServer(),
+						},
+					}
 					for i := range c.SQL {
-						files := []string{}
-						for _, s := range c.SQL[i].Schema {
-							files = append(files, filepath.Join(path, s))
-						}
 						switch c.SQL[i].Engine {
 						case config.EnginePostgreSQL:
-							uri := local.ReadOnlyPostgreSQL(t, files)
 							c.SQL[i].Database = &config.Database{
-								URI: uri,
+								Managed: true,
 							}
 						case config.EngineMySQL:
-							uri := local.MySQL(t, files)
 							c.SQL[i].Database = &config.Database{
-								URI: uri,
+								Managed: true,
 							}
 						default:
 							// pass
@@ -165,8 +172,6 @@ func TestReplay(t *testing.T) {
 		for _, replay := range FindTests(t, "testdata", name) {
 			tc := replay
 			t.Run(filepath.Join(name, tc.Name), func(t *testing.T) {
-				t.Parallel()
-
 				var stderr bytes.Buffer
 				var output map[string]string
 				var err error

internal/endtoend/testdata/pg_timezone_names/schema.sql

@@ -0,0 +1 @@
+SELECT 1;

internal/endtoend/testdata/pg_timezone_names/sqlc.json

@@ -3,7 +3,7 @@
   "sql": [
     {
       "engine": "postgresql",
-      "schema": "query.sql",
+      "schema": "schema.sql",
       "queries": "query.sql",
       "gen": {
         "go": {
@@ -15,7 +15,7 @@
     },
     {
       "engine": "postgresql",
-      "schema": "query.sql",
+      "schema": "schema.sql",
       "queries": "query.sql",
       "gen": {
         "go": {
@@ -27,7 +27,7 @@
     },
     {
       "engine": "postgresql",
-      "schema": "query.sql",
+      "schema": "schema.sql",
       "queries": "query.sql",
       "gen": {
         "go": {

internal/engine/postgresql/analyzer/analyze.go

@@ -13,8 +13,8 @@ import (
 
 	core "github.com/sqlc-dev/sqlc/internal/analysis"
 	"github.com/sqlc-dev/sqlc/internal/config"
+	"github.com/sqlc-dev/sqlc/internal/dbmanager"
 	"github.com/sqlc-dev/sqlc/internal/opts"
-	pb "github.com/sqlc-dev/sqlc/internal/quickdb/v1"
 	"github.com/sqlc-dev/sqlc/internal/shfmt"
 	"github.com/sqlc-dev/sqlc/internal/sql/ast"
 	"github.com/sqlc-dev/sqlc/internal/sql/named"
@@ -23,7 +23,7 @@ import (
 
 type Analyzer struct {
 	db       config.Database
-	client   pb.QuickClient
+	client   dbmanager.Client
 	pool     *pgxpool.Pool
 	dbg      opts.Debug
 	replacer *shfmt.Replacer
@@ -32,7 +32,7 @@ type Analyzer struct {
 	tables   sync.Map
 }
 
-func New(client pb.QuickClient, db config.Database) *Analyzer {
+func New(client dbmanager.Client, db config.Database) *Analyzer {
 	return &Analyzer{
 		db:       db,
 		dbg:      opts.DebugFromEnv(),
@@ -201,7 +201,7 @@ func (a *Analyzer) Analyze(ctx context.Context, n ast.Node, query string, migrat
 			if a.client == nil {
 				return nil, fmt.Errorf("client is nil")
 			}
-			edb, err := a.client.CreateEphemeralDatabase(ctx, &pb.CreateEphemeralDatabaseRequest{
+			edb, err := a.client.CreateDatabase(ctx, &dbmanager.CreateDatabaseRequest{
 				Engine:     "postgresql",
 				Migrations: migrations,
 			})