feat: Managed databases with any accessible server

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,
 			})