hacky support for Postgres schemas w/ runtime interpolation

sqlc-dev · May 8, 2024 · 5013ed7 · 5013ed7
1 parent 9ec7a68
commit 5013ed7
Show file tree

Hide file tree

Showing 7 changed files with 27 additions and 7 deletions.
diff --git a/internal/codegen/golang/gen.go b/internal/codegen/golang/gen.go
@@ -37,6 +37,7 @@ type tmplCtx struct {
 	EmitMethodsWithDBArgument bool
 	EmitEnumValidMethod       bool
 	EmitAllEnumValues         bool
+	UseSchemaPlaceholder      bool
 	UsesCopyFrom              bool
 	UsesBatch                 bool
 	OmitSqlcVersion           bool
@@ -177,6 +178,7 @@ func generate(req *plugin.GenerateRequest, options *opts.Options, enums []Enum,
 		EmitMethodsWithDBArgument: options.EmitMethodsWithDbArgument,
 		EmitEnumValidMethod:       options.EmitEnumValidMethod,
 		EmitAllEnumValues:         options.EmitAllEnumValues,
+		UseSchemaPlaceholder:      options.UseSchemaPlaceholder,
 		UsesCopyFrom:              usesCopyFrom(queries),
 		UsesBatch:                 usesBatch(queries),
 		SQLDriver:                 parseDriver(options.SqlPackage),

diff --git a/internal/codegen/golang/imports.go b/internal/codegen/golang/imports.go
@@ -124,6 +124,10 @@ func (i *importer) dbImports() fileImports {
 		{Path: "context"},
 	}
 
+	if i.Options.UseSchemaPlaceholder {
+		std = append(std, ImportSpec{Path: "strings"})
+	}
+
 	sqlpkg := parseDriver(i.Options.SqlPackage)
 	switch sqlpkg {
 	case opts.SQLDriverPGXV4:

diff --git a/internal/codegen/golang/opts/options.go b/internal/codegen/golang/opts/options.go
@@ -44,6 +44,7 @@ type Options struct {
 	OmitUnusedStructs           bool              `json:"omit_unused_structs,omitempty" yaml:"omit_unused_structs"`
 	BuildTags                   string            `json:"build_tags,omitempty" yaml:"build_tags"`
 	Initialisms                 *[]string         `json:"initialisms,omitempty" yaml:"initialisms"`
+	UseSchemaPlaceholder        bool              `json:"use_schema_placeholder,omitempty" yaml:"use_schema_placeholder"`
 
 	InitialismsMap map[string]struct{} `json:"-" yaml:"-"`
 }

diff --git a/internal/codegen/golang/query.go b/internal/codegen/golang/query.go
@@ -278,7 +278,7 @@ func (q Query) hasRetType() bool {
 func (q Query) TableIdentifierAsGoSlice() string {
 	escapedNames := make([]string, 0, 3)
 	for _, p := range []string{q.Table.Catalog, q.Table.Schema, q.Table.Name} {
-		if p != "" {
+		if p != "" && p != "sqlc_schema_placeholder" {
 			escapedNames = append(escapedNames, fmt.Sprintf("%q", p))
 		}
 	}

diff --git a/internal/codegen/golang/struct.go b/internal/codegen/golang/struct.go
@@ -17,6 +17,7 @@ type Struct struct {
 }
 
 func StructName(name string, options *opts.Options) string {
+	name = strings.TrimPrefix(name, "sqlc_schema_placeholder")
 	if rename := options.Rename[name]; rename != "" {
 		return rename
 	}

diff --git a/internal/codegen/golang/templates/pgx/dbCode.tmpl b/internal/codegen/golang/templates/pgx/dbCode.tmpl
@@ -12,7 +12,10 @@ type DBTX interface {
 {{- end }}
 }
 
-{{ if .EmitMethodsWithDBArgument}}
+{{ if and .EmitMethodsWithDBArgument .UseSchemaPlaceholder  }}
+func New(schema string) *Queries {
+	return &Queries{schema: schema}
+{{- else if .EmitMethodsWithDBArgument}}
 func New() *Queries {
 	return &Queries{}
 {{- else -}}
@@ -25,7 +28,16 @@ type Queries struct {
     {{if not .EmitMethodsWithDBArgument}}
 	db DBTX
     {{end}}
+    {{if .UseSchemaPlaceholder}}
+	schema string
+    {{end}}
+}
+
+{{if .UseSchemaPlaceholder}}
+func (q *Queries) interpolateSchemaName(query string) string {
+	return strings.ReplaceAll(query, "sqlc_schema_placeholder.", q.schema+".")
 }
+{{end -}}
 
 {{if not .EmitMethodsWithDBArgument}}
 func (q *Queries) WithTx(tx pgx.Tx) *Queries {

diff --git a/internal/codegen/golang/templates/pgx/queryCode.tmpl b/internal/codegen/golang/templates/pgx/queryCode.tmpl
@@ -28,7 +28,7 @@ type {{.Ret.Type}} struct { {{- range .Ret.Struct.Fields}}
 {{end -}}
 {{- if $.EmitMethodsWithDBArgument -}}
 func (q *Queries) {{.MethodName}}(ctx context.Context, db DBTX, {{.Arg.Pair}}) ({{.Ret.DefineType}}, error) {
-	row := db.QueryRow(ctx, {{.ConstantName}}, {{.Arg.Params}})
+	row := db.QueryRow(ctx, {{ if $.UseSchemaPlaceholder }}q.interpolateSchemaName({{.ConstantName}}){{else}}{{.ConstantName}}{{end}}, {{.Arg.Params}})
 {{- else -}}
 func (q *Queries) {{.MethodName}}(ctx context.Context, {{.Arg.Pair}}) ({{.Ret.DefineType}}, error) {
 	row := q.db.QueryRow(ctx, {{.ConstantName}}, {{.Arg.Params}})
@@ -46,7 +46,7 @@ func (q *Queries) {{.MethodName}}(ctx context.Context, {{.Arg.Pair}}) ({{.Ret.De
 {{end -}}
 {{- if $.EmitMethodsWithDBArgument -}}
 func (q *Queries) {{.MethodName}}(ctx context.Context, db DBTX, {{.Arg.Pair}}) ([]{{.Ret.DefineType}}, error) {
-	rows, err := db.Query(ctx, {{.ConstantName}}, {{.Arg.Params}})
+	rows, err := db.Query(ctx, {{ if $.UseSchemaPlaceholder }}q.interpolateSchemaName({{.ConstantName}}){{else}}{{.ConstantName}}{{end}}, {{.Arg.Params}})
 {{- else -}}
 func (q *Queries) {{.MethodName}}(ctx context.Context, {{.Arg.Pair}}) ([]{{.Ret.DefineType}}, error) {
 	rows, err := q.db.Query(ctx, {{.ConstantName}}, {{.Arg.Params}})
@@ -79,7 +79,7 @@ func (q *Queries) {{.MethodName}}(ctx context.Context, {{.Arg.Pair}}) ([]{{.Ret.
 {{end -}}
 {{- if $.EmitMethodsWithDBArgument -}}
 func (q *Queries) {{.MethodName}}(ctx context.Context, db DBTX, {{.Arg.Pair}}) error {
-	_, err := db.Exec(ctx, {{.ConstantName}}, {{.Arg.Params}})
+	_, err := db.Exec(ctx, {{ if $.UseSchemaPlaceholder }}q.interpolateSchemaName({{.ConstantName}}){{else}}{{.ConstantName}}{{end}}, {{.Arg.Params}})
 {{- else -}}
 func (q *Queries) {{.MethodName}}(ctx context.Context, {{.Arg.Pair}}) error {
 	_, err := q.db.Exec(ctx, {{.ConstantName}}, {{.Arg.Params}})
@@ -93,7 +93,7 @@ func (q *Queries) {{.MethodName}}(ctx context.Context, {{.Arg.Pair}}) error {
 {{end -}}
 {{if $.EmitMethodsWithDBArgument -}}
 func (q *Queries) {{.MethodName}}(ctx context.Context, db DBTX, {{.Arg.Pair}}) (int64, error) {
-	result, err := db.Exec(ctx, {{.ConstantName}}, {{.Arg.Params}})
+	result, err := db.Exec(ctx, {{ if $.UseSchemaPlaceholder }}q.interpolateSchemaName({{.ConstantName}}){{else}}{{.ConstantName}}{{end}}, {{.Arg.Params}})
 {{- else -}}
 func (q *Queries) {{.MethodName}}(ctx context.Context, {{.Arg.Pair}}) (int64, error) {
 	result, err := q.db.Exec(ctx, {{.ConstantName}}, {{.Arg.Params}})
@@ -110,7 +110,7 @@ func (q *Queries) {{.MethodName}}(ctx context.Context, {{.Arg.Pair}}) (int64, er
 {{end -}}
 {{- if $.EmitMethodsWithDBArgument -}}
 func (q *Queries) {{.MethodName}}(ctx context.Context, db DBTX, {{.Arg.Pair}}) (pgconn.CommandTag, error) {
-	return db.Exec(ctx, {{.ConstantName}}, {{.Arg.Params}})
+	return db.Exec(ctx, {{ if $.UseSchemaPlaceholder }}q.interpolateSchemaName({{.ConstantName}}){{else}}{{.ConstantName}}{{end}}, {{.Arg.Params}})
 {{- else -}}
 func (q *Queries) {{.MethodName}}(ctx context.Context, {{.Arg.Pair}}) (pgconn.CommandTag, error) {
 	return q.db.Exec(ctx, {{.ConstantName}}, {{.Arg.Params}})