From 471499ea0063e3d341bc7a9f90da94d439da0ac6 Mon Sep 17 00:00:00 2001
From: Jay <me@bsavage.xyz>
Date: Thu, 1 Dec 2022 11:54:14 +0530
Subject: [PATCH] Added a global  method to table generation (#79)

---
 generator/template/file_templates.go |  8 ++++++
 generator/template/process.go        | 42 ++++++++++++++++++++++++++--
 2 files changed, 47 insertions(+), 3 deletions(-)

diff --git a/generator/template/file_templates.go b/generator/template/file_templates.go
index 35b4dcd..6ef375d 100644
--- a/generator/template/file_templates.go
+++ b/generator/template/file_templates.go
@@ -97,6 +97,14 @@ func new{{tableTemplate.TypeName}}Impl(schemaName, tableName, alias string) {{st
 }
 `
 
+var tableSqlBuilderSetSchemaTemplate = `
+func SetSchema(schema string) {
+{{- range .}}
+	{{ .InstanceName }} = {{ .InstanceName }}.FromSchema(schema)
+{{- end}}
+}
+`
+
 var tableModelFileTemplate = `package {{package}}
 
 {{ with modelImports }}
diff --git a/generator/template/process.go b/generator/template/process.go
index 09636dc..7643aaf 100644
--- a/generator/template/process.go
+++ b/generator/template/process.go
@@ -3,13 +3,15 @@ package template
 import (
 	"bytes"
 	"fmt"
+	"os"
+	"path"
+	"strings"
+	"text/template"
+
 	"github.com/go-jet/jet/v2/generator/metadata"
 	"github.com/go-jet/jet/v2/internal/jet"
 	"github.com/go-jet/jet/v2/internal/utils"
 	"github.com/go-jet/jet/v2/internal/utils/throw"
-	"path"
-	"strings"
-	"text/template"
 )
 
 // ProcessSchema will process schema metadata and constructs go files using generator Template
@@ -61,6 +63,40 @@ func processSQLBuilder(dirPath string, dialect jet.Dialect, schemaMetaData metad
 	processTableSQLBuilder("table", sqlBuilderPath, dialect, schemaMetaData, schemaMetaData.TablesMetaData, sqlBuilderTemplate)
 	processTableSQLBuilder("view", sqlBuilderPath, dialect, schemaMetaData, schemaMetaData.ViewsMetaData, sqlBuilderTemplate)
 	processEnumSQLBuilder(sqlBuilderPath, dialect, schemaMetaData.EnumsMetaData, sqlBuilderTemplate)
+	processTableSQLBuilderSetSchema(sqlBuilderPath, schemaMetaData.TablesMetaData, sqlBuilderTemplate)
+}
+
+func processTableSQLBuilderSetSchema(dirPath string, tablesMetadata []metadata.Table, builderTemplate SQLBuilder) {
+	if len(tablesMetadata) == 0 {
+		return
+	}
+
+	fmt.Println("Generating global `SetSchema` method...")
+
+	err := utils.EnsureDirPath(dirPath)
+	throw.OnError(err)
+
+	var builders []TableSQLBuilder
+	for _, tm := range tablesMetadata {
+		builders = append(builders, builderTemplate.Table(tm))
+	}
+
+	funcPath := path.Join(dirPath, builders[0].Path)
+
+	origText, err := os.ReadFile(path.Join(funcPath, builders[0].FileName+".go"))
+	throw.OnError(err)
+
+	text, err := generateTemplate(
+		tableSqlBuilderSetSchemaTemplate,
+		builders,
+		nil,
+	)
+	throw.OnError(err)
+
+	text = append(origText, text...)
+
+	err = utils.SaveGoFile(funcPath, builders[0].FileName, text)
+	throw.OnError(err)
 }
 
 func processEnumSQLBuilder(dirPath string, dialect jet.Dialect, enumsMetaData []metadata.Enum, sqlBuilder SQLBuilder) {