Merge pull request #443 from TotallyGamerJet/441_gen

generator/*: allow passing DB directly
go-jet · Jan 23, 2025 · 00b8155 · 00b8155
2 parents 269bb4a + 8975f87
commit 00b8155
Show file tree

Hide file tree

Showing 3 changed files with 15 additions and 6 deletions.
diff --git a/generator/mysql/mysql_generator.go b/generator/mysql/mysql_generator.go
@@ -38,7 +38,7 @@ func Generate(destDir string, dbConn DBConnection, generatorTemplate ...template
 	}
 	defer db.Close()
 
-	err = generate(db, dbConn.DBName, destDir, generatorTemplate...)
+	err = GenerateDB(db, dbConn.DBName, destDir, generatorTemplate...)
 	if err != nil {
 		return err
 	}
@@ -70,7 +70,7 @@ func GenerateDSN(dsn, destDir string, templates ...template.Template) error {
 	}
 	defer db.Close()
 
-	err = generate(db, cfg.DBName, destDir, templates...)
+	err = GenerateDB(db, cfg.DBName, destDir, templates...)
 	if err != nil {
 		return fmt.Errorf("failed to generate: %w", err)
 	}
@@ -96,7 +96,8 @@ func openConnection(connectionString string) (*sql.DB, error) {
 	return db, nil
 }
 
-func generate(db *sql.DB, dbName, destDir string, templates ...template.Template) error {
+// GenerateDB generates jet files using the provided *sql.DB
+func GenerateDB(db *sql.DB, dbName, destDir string, templates ...template.Template) error {
 	fmt.Println("Retrieving database information...")
 	// No schemas in MySQL
 	schemaMetaData, err := metadata.GetSchema(db, &mySqlQuerySet{}, dbName)

diff --git a/generator/postgres/postgres_generator.go b/generator/postgres/postgres_generator.go
@@ -56,6 +56,11 @@ func GenerateDSN(dsn, schema, destDir string, templates ...template.Template) er
 	defer db.Close()
 
 	fmt.Println("Retrieving schema information...")
+	return GenerateDB(db, schema, filepath.Join(destDir, cfg.Database), templates...)
+}
+
+// GenerateDB generates jet files using the provided *sql.DB
+func GenerateDB(db *sql.DB, schema, destDir string, templates ...template.Template) error {
 	generatorTemplate := template.Default(postgres.Dialect)
 	if len(templates) > 0 {
 		generatorTemplate = templates[0]
@@ -66,9 +71,7 @@ func GenerateDSN(dsn, schema, destDir string, templates ...template.Template) er
 		return fmt.Errorf("failed to get '%s' schema metadata: %w", schema, err)
 	}
 
-	dirPath := filepath.Join(destDir, cfg.Database)
-
-	err = template.ProcessSchema(dirPath, schemaMetadata, generatorTemplate)
+	err = template.ProcessSchema(destDir, schemaMetadata, generatorTemplate)
 	if err != nil {
 		return fmt.Errorf("failed to generate schema %s: %d", schemaMetadata.Name, err)
 	}

diff --git a/generator/sqlite/sqlite_generator.go b/generator/sqlite/sqlite_generator.go
@@ -3,6 +3,7 @@ package sqlite
 import (
 	"database/sql"
 	"fmt"
+
 	"github.com/go-jet/jet/v2/generator/metadata"
 	"github.com/go-jet/jet/v2/generator/template"
 	"github.com/go-jet/jet/v2/sqlite"
@@ -17,7 +18,11 @@ func GenerateDSN(dsn, destDir string, templates ...template.Template) error {
 	defer db.Close()
 
 	fmt.Println("Retrieving schema information...")
+	return GenerateDB(db, destDir, templates...)
+}
 
+// GenerateDB generates jet files using the provided *sql.DB
+func GenerateDB(db *sql.DB, destDir string, templates ...template.Template) error {
 	generatorTemplate := template.Default(sqlite.Dialect)
 	if len(templates) > 0 {
 		generatorTemplate = templates[0]