tools: Generate functions for all of contrib

Author: kyleconroy

internal/tools/sqlc-pg-gen/main.go

@@ -59,14 +59,16 @@ ORDER BY 1;
 `
 
 const catalogTmpl = `
-package postgresql
+// Code generated by sqlc-pg-gen. DO NOT EDIT.
+
+package {{.Pkg}}
 
 import (
 	"github.com/kyleconroy/sqlc/internal/sql/ast"
 	"github.com/kyleconroy/sqlc/internal/sql/catalog"
 )
 
-func gen{{.Name}}() *catalog.Schema {
+func {{.Name}}() *catalog.Schema {
 	s := &catalog.Schema{Name: "pg_catalog"}
 	s.Funcs = []*catalog.Function{
 	    {{- range .Funcs}}
@@ -92,7 +94,29 @@ func gen{{.Name}}() *catalog.Schema {
 }
 `
 
+const loaderFuncTmpl = `
+// Code generated by sqlc-pg-gen. DO NOT EDIT.
+
+package postgresql
+
+import (
+	"github.com/kyleconroy/sqlc/internal/engine/postgresql/contrib"
+	"github.com/kyleconroy/sqlc/internal/sql/catalog"
+)
+
+func loadExtension(name string) *catalog.Schema {
+	switch name {
+	{{- range .}}
+	case "{{.Name}}":
+		return contrib.{{.Func}}()
+	{{- end}}
+	}
+	return nil
+}
+`
+
 type tmplCtx struct {
+	Pkg   string
 	Name  string
 	Funcs []catalog.Function
 }
@@ -220,48 +244,134 @@ func run(ctx context.Context) error {
 		return err
 	}
 	out := bytes.NewBuffer([]byte{})
-	if err := tmpl.Execute(out, tmplCtx{Name: "PGCatalog", Funcs: funcs}); err != nil {
+	if err := tmpl.Execute(out, tmplCtx{Pkg: "postgresql", Name: "genPGCatalog", Funcs: funcs}); err != nil {
 		return err
 	}
 	code, err := format.Source(out.Bytes())
 	if err != nil {
 		return err
 	}
-	err = ioutil.WriteFile(filepath.Join("internal", "engine", "postgresql", "pg_catalog.gen.go"), code, 0644)
+	err = ioutil.WriteFile(filepath.Join("internal", "engine", "postgresql", "pg_catalog.go"), code, 0644)
 	if err != nil {
 		return err
 	}
 
-	// https://www.postgresql.org/docs/current/contrib.html
-	extensions := map[string]string{
-		"citext":    "CIText",
-		"pg_trgm":   "PGTrigram",
-		"pgcrypto":  "PGCrypto",
-		"uuid-ossp": "UUIDOSSP",
-	}
+	loaded := []extensionPair{}
+
+	for _, extension := range extensions {
+		name := strings.Replace(extension, "-", "_", -1)
+
+		var funcName string
+		for _, part := range strings.Split(name, "_") {
+			funcName += strings.Title(part)
+		}
 
-	for extension, name := range extensions {
 		_, err := conn.Exec(ctx, fmt.Sprintf("CREATE EXTENSION IF NOT EXISTS \"%s\"", extension))
 		if err != nil {
-			return err
+			log.Printf("error creating %s: %s", extension, err)
+			continue
 		}
+
 		rows, err := conn.Query(ctx, extensionFuncs, extension)
+		if err != nil {
+			return err
+		}
 		funcs, err := scanFuncs(rows)
 		if err != nil {
 			return err
 		}
+		if len(funcs) == 0 {
+			log.Printf("no functions in %s, skipping", extension)
+			continue
+		}
 		out := bytes.NewBuffer([]byte{})
-		if err := tmpl.Execute(out, tmplCtx{Name: name, Funcs: funcs}); err != nil {
+		if err := tmpl.Execute(out, tmplCtx{Pkg: "contrib", Name: funcName, Funcs: funcs}); err != nil {
 			return err
 		}
 		code, err := format.Source(out.Bytes())
 		if err != nil {
 			return err
 		}
-		err = ioutil.WriteFile(filepath.Join("internal", "engine", "postgresql", "extension_"+strings.Replace(extension, "-", "_", -1)+".gen.go"), code, 0644)
+		err = ioutil.WriteFile(filepath.Join("internal", "engine", "postgresql", "contrib", name+".go"), code, 0644)
 		if err != nil {
 			return err
 		}
+
+		loaded = append(loaded, extensionPair{Name: extension, Func: funcName})
 	}
+
+	{
+		tmpl, err := template.New("").Parse(loaderFuncTmpl)
+		if err != nil {
+			return err
+		}
+		out := bytes.NewBuffer([]byte{})
+		if err := tmpl.Execute(out, loaded); err != nil {
+			return err
+		}
+		code, err := format.Source(out.Bytes())
+		if err != nil {
+			return err
+		}
+		err = ioutil.WriteFile(filepath.Join("internal", "engine", "postgresql", "extension.go"), code, 0644)
+		if err != nil {
+			return err
+		}
+	}
+
 	return nil
 }
+
+type extensionPair struct {
+	Name string
+	Func string
+}
+
+// https://www.postgresql.org/docs/current/contrib.html
+var extensions = []string{
+	"adminpack",
+	"amcheck",
+	"auth_delay",
+	"auto_explain",
+	"bloom",
+	"btree_gin",
+	"btree_gist",
+	"citext",
+	"cube",
+	"dblink",
+	"dict_int",
+	"dict_xsyn",
+	"earthdistance",
+	"file_fdw",
+	"fuzzystrmatch",
+	"hstore",
+	"intagg",
+	"intarray",
+	"isn",
+	"lo",
+	"ltree",
+	"pageinspect",
+	"passwordcheck",
+	"pg_buffercache",
+	"pgcrypto",
+	"pg_freespacemap",
+	"pg_prewarm",
+	"pgrowlocks",
+	"pg_stat_statements",
+	"pgstattuple",
+	"pg_trgm",
+	"pg_visibility",
+	"postgres_fdw",
+	"seg",
+	"sepgsql",
+	"spi",
+	"sslinfo",
+	"tablefunc",
+	"tcn",
+	"test_decoding",
+	"tsm_system_rows",
+	"tsm_system_time",
+	"unaccent",
+	"uuid-ossp",
+	"xml2",
+}