feat(createdb): Create ephemeral databases (#2894)

* feat(createdb): Create ephemeral databases and return a database URI

Author: kyleconroy

internal/cmd/cmd.go

@@ -25,6 +25,7 @@ import (
 )
 
 func init() {
+	createDBCmd.Flags().StringP("queryset", "", "", "name of the queryset to use")
 	uploadCmd.Flags().BoolP("dry-run", "", false, "dump upload request (default: false)")
 	initCmd.Flags().BoolP("v1", "", false, "generate v1 config yaml file")
 	initCmd.Flags().BoolP("v2", "", true, "generate v2 config yaml file")
@@ -41,6 +42,7 @@ func Do(args []string, stdin io.Reader, stdout io.Writer, stderr io.Writer) int
 	rootCmd.PersistentFlags().Bool("no-database", false, "disable database connections (default: false)")
 
 	rootCmd.AddCommand(checkCmd)
+	rootCmd.AddCommand(createDBCmd)
 	rootCmd.AddCommand(diffCmd)
 	rootCmd.AddCommand(genCmd)
 	rootCmd.AddCommand(initCmd)

internal/cmd/createdb.go

@@ -0,0 +1,101 @@
+package cmd
+
+import (
+	"context"
+	"fmt"
+	"os"
+	"runtime/trace"
+
+	"github.com/spf13/cobra"
+	"github.com/sqlc-dev/sqlc/internal/config"
+	"github.com/sqlc-dev/sqlc/internal/migrations"
+	"github.com/sqlc-dev/sqlc/internal/quickdb"
+	pb "github.com/sqlc-dev/sqlc/internal/quickdb/v1"
+	"github.com/sqlc-dev/sqlc/internal/sql/sqlpath"
+)
+
+var createDBCmd = &cobra.Command{
+	Use:   "createdb",
+	Short: "Create an ephemeral database",
+	Args:  cobra.NoArgs,
+	RunE: func(cmd *cobra.Command, args []string) error {
+		defer trace.StartRegion(cmd.Context(), "createdb").End()
+		stderr := cmd.ErrOrStderr()
+		dir, filename := getConfigPath(stderr, cmd.Flag("file"))
+		querySetName, err := cmd.Flags().GetString("queryset")
+		if err != nil {
+			return err
+		}
+		err = CreateDB(cmd.Context(), dir, filename, querySetName, &Options{
+			Env:    ParseEnv(cmd),
+			Stderr: stderr,
+		})
+		if err != nil {
+			fmt.Fprintln(stderr, err.Error())
+			os.Exit(1)
+		}
+		return nil
+	},
+}
+
+func CreateDB(ctx context.Context, dir, filename, querySetName string, o *Options) error {
+	_, conf, err := o.ReadConfig(dir, filename)
+	if err != nil {
+		return err
+	}
+	// Find the first queryset with a managed database
+	var queryset *config.SQL
+	var count int
+	for _, sql := range conf.SQL {
+		sql := sql
+		if querySetName != "" && sql.Name != querySetName {
+			continue
+		}
+		if sql.Database != nil && sql.Database.Managed {
+			queryset = &sql
+			count += 1
+		}
+	}
+	if queryset == nil && querySetName != "" {
+		return fmt.Errorf("no queryset found with name %q", querySetName)
+	}
+	if queryset == nil {
+		return fmt.Errorf("no querysets configured to use a managed database")
+	}
+	if count > 1 {
+		return fmt.Errorf("multiple querysets configured to use managed databases")
+	}
+	if queryset.Engine != config.EnginePostgreSQL {
+		return fmt.Errorf("managed databases currently only support PostgreSQL")
+	}
+
+	var ddl []string
+	files, err := sqlpath.Glob(queryset.Schema)
+	if err != nil {
+		return err
+	}
+	for _, schema := range files {
+		contents, err := os.ReadFile(schema)
+		if err != nil {
+			return fmt.Errorf("read file: %w", err)
+		}
+		ddl = append(ddl, migrations.RemoveRollbackStatements(string(contents)))
+	}
+
+	client, err := quickdb.NewClientFromConfig(conf.Cloud)
+	if err != nil {
+		return fmt.Errorf("client error: %w", err)
+	}
+
+	resp, err := client.CreateEphemeralDatabase(ctx, &pb.CreateEphemeralDatabaseRequest{
+		Engine:     "postgresql",
+		Region:     quickdb.GetClosestRegion(),
+		Migrations: ddl,
+	})
+	if err != nil {
+		return fmt.Errorf("managed: create database: %w", err)
+	}
+	fmt.Fprintln(os.Stderr, "WARNING: This database will be removed in two minutes")
+	fmt.Println(resp.Uri)
+	return nil
+}

internal/config/config.go

@@ -103,6 +103,7 @@ type GenGo struct {
 }
 
 type SQL struct {
+	Name                 string    `json:"name" yaml:"name"`
 	Engine               Engine    `json:"engine,omitempty" yaml:"engine"`
 	Schema               Paths     `json:"schema" yaml:"schema"`
 	Queries              Paths     `json:"queries" yaml:"queries"`

internal/config/v_one.go

@@ -143,6 +143,7 @@ func (c *V1GenerateSettings) Translate() Config {
 			pkg.StrictOrderBy = &defaultValue
 		}
 		conf.SQL = append(conf.SQL, SQL{
+			Name:     pkg.Name,
 			Engine:   pkg.Engine,
 			Database: pkg.Database,
 			Schema:   pkg.Schema,

internal/config/v_two.json

@@ -31,6 +31,9 @@
                     "engine"
                 ],
                 "properties": {
+                    "name": {
+                        "type": "string"
+                    },
                     "engine": {
                         "enum": [
                             "postgresql",
@@ -439,4 +442,4 @@
             }
         }
     }
-}
+}