fix: robustly strip psql meta commands

internal/cmd/createdb.go

@@ -10,8 +10,7 @@ import (
 	"github.com/spf13/cobra"
 	"github.com/sqlc-dev/sqlc/internal/config"
 	"github.com/sqlc-dev/sqlc/internal/dbmanager"
-	"github.com/sqlc-dev/sqlc/internal/migrations"
-	"github.com/sqlc-dev/sqlc/internal/sql/sqlpath"
+	"github.com/sqlc-dev/sqlc/internal/schemautil"
 )
 
 var createDBCmd = &cobra.Command{
@@ -75,17 +74,12 @@ func CreateDB(ctx context.Context, dir, filename, querySetName string, o *Option
 	}
 
 	var ddl []string
-	files, err := sqlpath.Glob(queryset.Schema)
+	ddl, err = schemautil.LoadSchemasForApply(queryset.Schema, string(queryset.Engine), func(warning string) {
+		fmt.Fprintln(o.Stderr, warning)
+	})
 	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)))
-	}
 
 	now := time.Now().UTC().UnixNano()
 	client := dbmanager.NewClient(conf.Servers)

internal/cmd/generate.go

@@ -269,6 +269,9 @@ func parse(ctx context.Context, name, dir string, sql config.SQL, combo config.C
 		return nil, true
 	}
 	if err := c.ParseCatalog(sql.Schema); err != nil {
+		for _, warning := range c.Warnings() {
+			fmt.Fprintln(stderr, warning)
+		}
 		fmt.Fprintf(stderr, "# package %s\n", name)
 		if parserErr, ok := err.(*multierr.Error); ok {
 			for _, fileErr := range parserErr.Errs() {
@@ -279,6 +282,9 @@ func parse(ctx context.Context, name, dir string, sql config.SQL, combo config.C
 		}
 		return nil, true
 	}
+	for _, warning := range c.Warnings() {
+		fmt.Fprintln(stderr, warning)
+	}
 	if debugDumpCatalog.Value() == "1" {
 		debug.Dump(c.Catalog())
 	}

internal/cmd/process.go

@@ -121,12 +121,12 @@ func processQuerySets(ctx context.Context, rp ResultProcessor, conf *config.Conf
 	if err := grp.Wait(); err != nil {
 		return err
 	}
-	if errored {
-		for i, _ := range stderrs {
-			if _, err := io.Copy(stderr, &stderrs[i]); err != nil {
-				return err
-			}
+	for i := range stderrs {
+		if _, err := io.Copy(stderr, &stderrs[i]); err != nil {
+			return err
 		}
+	}
+	if errored {
 		return fmt.Errorf("errored")
 	}
 	return nil

internal/cmd/verify.go

@@ -14,11 +14,10 @@ import (
 
 	"github.com/sqlc-dev/sqlc/internal/config"
 	"github.com/sqlc-dev/sqlc/internal/dbmanager"
-	"github.com/sqlc-dev/sqlc/internal/migrations"
 	"github.com/sqlc-dev/sqlc/internal/plugin"
 	"github.com/sqlc-dev/sqlc/internal/quickdb"
 	pb "github.com/sqlc-dev/sqlc/internal/quickdb/v1"
-	"github.com/sqlc-dev/sqlc/internal/sql/sqlpath"
+	"github.com/sqlc-dev/sqlc/internal/schemautil"
 )
 
 func init() {
@@ -93,17 +92,12 @@ func Verify(ctx context.Context, dir, filename string, opts *Options) error {
 
 			// Read the schema files into memory, removing rollback statements
 			var ddl []string
-			files, err := sqlpath.Glob(current.Schema)
+			ddl, err = schemautil.LoadSchemasForApply(current.Schema, string(current.Engine), func(warning string) {
+				fmt.Fprintln(stderr, warning)
+			})
 			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)))
-			}
 
 			var codegen plugin.GenerateRequest
 			if err := proto.Unmarshal(qs.CodegenRequest.Contents, &codegen); err != nil {

internal/cmd/vet.go

@@ -27,12 +27,11 @@ import (
 	"github.com/sqlc-dev/sqlc/internal/config"
 	"github.com/sqlc-dev/sqlc/internal/dbmanager"
 	"github.com/sqlc-dev/sqlc/internal/debug"
-	"github.com/sqlc-dev/sqlc/internal/migrations"
 	"github.com/sqlc-dev/sqlc/internal/opts"
 	"github.com/sqlc-dev/sqlc/internal/plugin"
 	"github.com/sqlc-dev/sqlc/internal/quickdb"
+	"github.com/sqlc-dev/sqlc/internal/schemautil"
 	"github.com/sqlc-dev/sqlc/internal/shfmt"
-	"github.com/sqlc-dev/sqlc/internal/sql/sqlpath"
 	"github.com/sqlc-dev/sqlc/internal/sqlcdebug"
 	"github.com/sqlc-dev/sqlc/internal/vet"
 )
@@ -429,18 +428,12 @@ func (c *checker) fetchDatabaseUri(ctx context.Context, s config.SQL) (string, f
 		c.Client = dbmanager.NewClient(c.Conf.Servers)
 	})
 
-	var ddl []string
-	files, err := sqlpath.Glob(s.Schema)
+	ddl, err := schemautil.LoadSchemasForApply(s.Schema, string(s.Engine), func(warning string) {
+		fmt.Fprintln(c.Stderr, warning)
+	})
 	if err != nil {
 		return "", cleanup, err
 	}
-	for _, schema := range files {
-		contents, err := os.ReadFile(schema)
-		if err != nil {
-			return "", cleanup, fmt.Errorf("read file: %w", err)
-		}
-		ddl = append(ddl, migrations.RemoveRollbackStatements(string(contents)))
-	}
 
 	resp, err := c.Client.CreateDatabase(ctx, &dbmanager.CreateDatabaseRequest{
 		Engine:     string(s.Engine),
@@ -545,18 +538,15 @@ func (c *checker) checkSQL(ctx context.Context, s config.SQL) error {
 			defer db.Close()
 			// For in-memory SQLite databases, apply migrations
 			if isInMemorySQLite(dburl) {
-				files, err := sqlpath.Glob(s.Schema)
+				ddl, err := schemautil.LoadSchemasForApply(s.Schema, string(s.Engine), func(warning string) {
+					fmt.Fprintln(c.Stderr, warning)
+				})
 				if err != nil {
 					return fmt.Errorf("schema: %w", err)
 				}
-				for _, schema := range files {
-					contents, err := os.ReadFile(schema)
-					if err != nil {
-						return fmt.Errorf("read schema file: %w", err)
-					}
-					ddl := migrations.RemoveRollbackStatements(string(contents))
-					if _, err := db.ExecContext(ctx, ddl); err != nil {
-						return fmt.Errorf("apply schema %s: %w", schema, err)
+				for _, stmt := range ddl {
+					if _, err := db.ExecContext(ctx, stmt); err != nil {
+						return fmt.Errorf("apply schema: %w", err)
 					}
 				}
 			}

internal/compiler/compile.go

@@ -52,10 +52,25 @@ func (c *Compiler) parseCatalog(schemas []string) error {
 			merr.Add(path, "", 0, err)
 			continue
 		}
-		contents := migrations.RemoveRollbackStatements(string(blob))
-		contents = migrations.RemovePsqlMetaCommands(contents)
+		contents, warnings, err := migrations.PreprocessSchema(string(blob), string(c.conf.Engine))
+		if err != nil {
+			merr.Add(path, string(blob), 0, err)
+			continue
+		}
+		var applyContents string
+		if c.usesManagedAnalyzer() {
+			applyContents, _, err = migrations.PreprocessSchemaForApply(string(blob), string(c.conf.Engine))
+			if err != nil {
+				merr.Add(path, string(blob), 0, err)
+				continue
+			}
+		}
+		c.warns = append(c.warns, warnings...)
 		files = append(files, schemaFile{name: path, contents: contents})
 		c.schema = append(c.schema, contents)
+		if c.usesManagedAnalyzer() {
+			c.applySchema = append(c.applySchema, applyContents)
+		}
 	}
 
 	if c.coreAnalysis {

internal/compiler/compile_test.go

@@ -0,0 +1,43 @@
+package compiler
+
+import (
+	"os"
+	"path/filepath"
+	"strings"
+	"testing"
+
+	"github.com/sqlc-dev/sqlc/internal/config"
+	"github.com/sqlc-dev/sqlc/internal/multierr"
+	"github.com/sqlc-dev/sqlc/internal/opts"
+)
+
+func TestParseCatalogManagedAnalyzerRejectsSemanticPsqlCommandsForApply(t *testing.T) {
+	dir := t.TempDir()
+	schema := filepath.Join(dir, "schema.sql")
+	if err := os.WriteFile(schema, []byte("\\include extra.sql\nCREATE TABLE foo (id int);\n"), 0600); err != nil {
+		t.Fatal(err)
+	}
+
+	c, err := NewCompiler(config.SQL{
+		Engine: config.EnginePostgreSQL,
+		Schema: []string{schema},
+		Database: &config.Database{
+			Managed: true,
+		},
+	}, config.CombinedSettings{}, opts.Parser{})
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	err = c.ParseCatalog([]string{schema})
+	if err == nil {
+		t.Fatal("expected managed analyzer schema preprocessing to reject semantic psql command")
+	}
+	merr, ok := err.(*multierr.Error)
+	if !ok || len(merr.Errs()) != 1 {
+		t.Fatalf("expected one schema error, got %T: %v", err, err)
+	}
+	if !strings.Contains(merr.Errs()[0].Err.Error(), `psql meta-command \include is not supported`) {
+		t.Fatalf("unexpected error: %v", err)
+	}
+}

internal/compiler/engine.go

@@ -39,7 +39,9 @@ type Compiler struct {
 	coreAnalysis bool
 	coreDialect  core.Option
 
-	schema []string
+	schema      []string
+	applySchema []string
+	warns       []string
 }
 
 // Option configures a Compiler.
@@ -159,6 +161,17 @@ func (c *Compiler) ParseCatalog(schema []string) error {
 	return c.parseCatalog(schema)
 }
 
+func (c *Compiler) usesManagedAnalyzer() bool {
+	return c.analyzer != nil && c.conf.Database != nil && c.conf.Database.Managed
+}
+
+func (c *Compiler) analyzerMigrations() []string {
+	if c.usesManagedAnalyzer() {
+		return c.applySchema
+	}
+	return c.schema
+}
+
 func (c *Compiler) ParseQueries(queries []string, o opts.Parser) error {
 	r, err := c.parseQueries(o)
 	if err != nil {
@@ -172,6 +185,12 @@ func (c *Compiler) Result() *Result {
 	return c.result
 }
 
+// Warnings returns a copy of any non-fatal schema preprocessing warnings
+// collected while parsing the catalog.
+func (c *Compiler) Warnings() []string {
+	return append([]string(nil), c.warns...)
+}
+
 func (c *Compiler) Close(ctx context.Context) {
 	if c.analyzer != nil {
 		c.analyzer.Close(ctx)

internal/compiler/parse.go

@@ -93,7 +93,7 @@ func (c *Compiler) parseQuery(stmt ast.Node, pp *preprocess.Result, o opts.Parse
 			inference.Query = rawSQL
 		}
 
-		result, err := c.analyzer.Analyze(ctx, raw, inference.Query, c.schema, inference.Named)
+		result, err := c.analyzer.Analyze(ctx, raw, inference.Query, c.analyzerMigrations(), inference.Named)
 		if err != nil {
 			return nil, err
 		}