From b7bcb3527e69075703a0027b40d475b9a994f04a Mon Sep 17 00:00:00 2001
From: Joonas Haapsaari <joonas.haapsaari@gmail.com>
Date: Tue, 18 Aug 2020 15:47:10 +0300
Subject: [PATCH] Schema rename support

- Support for renaming table schemas
 * Table support for renaming schema
 * Empty schema name is left out (using default schema for the
   database connection)
 * Generated code support for obtaining a version of the table with
   renamed schema, similarly as the `AS` function works
 * Unit tests for setting and clearing the schema name
---
 generator/internal/template/templates.go | 18 ++++++++++++
 internal/jet/table.go                    | 25 +++++++++++++++--
 internal/jet/table_test.go               | 35 ++++++++++++++++++++++++
 3 files changed, 76 insertions(+), 2 deletions(-)

diff --git a/generator/internal/template/templates.go b/generator/internal/template/templates.go
index 56dba2c..5288ef4 100644
--- a/generator/internal/template/templates.go
+++ b/generator/internal/template/templates.go
@@ -41,6 +41,15 @@ type {{.GoStructName}} struct {
 func (a *{{.GoStructName}}) AS(alias string) {{.GoStructName}} {
 	aliasTable := new{{.GoStructName}}()
 	aliasTable.Table.AS(alias)
+	aliasTable.Table.Schema(a.Table.SchemaName())
+	return aliasTable
+}
+
+// Schema creates new {{.GoStructName}} with assigned schema name
+func (a *{{.GoStructName}}) Schema(schemaName string) {{.GoStructName}} {
+	aliasTable := new{{.GoStructName}}()
+	aliasTable.Table.AS(a.Table.Alias())
+	aliasTable.Table.Schema(schemaName)
 	return aliasTable
 }
 
@@ -104,6 +113,15 @@ type {{.GoStructName}} struct {
 func (a *{{.GoStructName}}) AS(alias string) *{{.GoStructName}} {
 	aliasTable := new{{.GoStructName}}()
 	aliasTable.Table.AS(alias)
+	aliasTable.Table.Schema(a.Table.SchemaName())
+	return aliasTable
+}
+
+// Schema creates new {{.GoStructName}} with assigned schema name
+func (a *{{.GoStructName}}) Schema(schemaName string) *{{.GoStructName}} {
+	aliasTable := new{{.GoStructName}}()
+	aliasTable.Table.AS(a.Table.Alias())
+	aliasTable.Table.Schema(schemaName)
 	return aliasTable
 }
 
diff --git a/internal/jet/table.go b/internal/jet/table.go
index 45d46e9..68f5962 100644
--- a/internal/jet/table.go
+++ b/internal/jet/table.go
@@ -16,6 +16,8 @@ type Table interface {
 	SchemaName() string
 	TableName() string
 	AS(alias string)
+	Alias() string
+	Schema(schemaName string)
 }
 
 // NewTable creates new table with schema Name, table Name and list of columns
@@ -67,13 +69,25 @@ func (t *tableImpl) columns() []Column {
 	return ret
 }
 
+func (t *tableImpl) Alias() string {
+	return t.alias
+}
+
+func (t *tableImpl) Schema(schemaName string) {
+	t.schemaName = schemaName
+}
+
 func (t *tableImpl) serialize(statement StatementType, out *SQLBuilder, options ...SerializeOption) {
 	if t == nil {
 		panic("jet: tableImpl is nil")
 	}
 
-	out.WriteIdentifier(t.schemaName)
-	out.WriteString(".")
+	// Use default schema if the schema name is not set
+	if len(t.schemaName) > 0 {
+		out.WriteIdentifier(t.schemaName)
+		out.WriteString(".")
+	}
+
 	out.WriteIdentifier(t.name)
 
 	if len(t.alias) > 0 {
@@ -145,6 +159,13 @@ func (t *joinTableImpl) columns() []Column {
 	return ret
 }
 
+func (t *joinTableImpl) Alias() string {
+	return ""
+}
+
+func (t *joinTableImpl) Schema(schemaName string) {
+}
+
 func (t *joinTableImpl) serialize(statement StatementType, out *SQLBuilder, options ...SerializeOption) {
 	if t == nil {
 		panic("jet: Join table is nil. ")
diff --git a/internal/jet/table_test.go b/internal/jet/table_test.go
index 66646b2..47923cd 100644
--- a/internal/jet/table_test.go
+++ b/internal/jet/table_test.go
@@ -31,3 +31,38 @@ INNER JOIN schema.table2 ON ("intCol1" = "intCol2")`)
 	require.Equal(t, joinTable.columns()[0].Name(), "intCol1")
 	require.Equal(t, joinTable.columns()[1].Name(), "intCol2")
 }
+
+func TestSchemaNameSet(t *testing.T) {
+	newTable := NewTable("schema", "table")
+	newTable.Schema("foo")
+	newTable.AS("bar")
+	assertClauseSerialize(t, newTable, `foo.table AS bar`)
+}
+
+func TestSchemaNameClear(t *testing.T) {
+	newTable := NewTable("schema", "table")
+	newTable.Schema("")
+	newTable.AS("bar")
+	assertClauseSerialize(t, newTable, `table AS bar`)
+}
+
+func TestNewJoinTableSchemaNameSet(t *testing.T) {
+	newTable1 := NewTable("schema", "table", IntegerColumn("intCol1"))
+	newTable2 := NewTable("schema", "table2", IntegerColumn("intCol2"))
+
+	newTable1.Schema("foo")
+	newTable2.Schema("foo")
+
+	joinTable := NewJoinTable(newTable1, newTable2, InnerJoin, IntegerColumn("intCol1").EQ(IntegerColumn("intCol2")))
+	joinTable.Schema("xxx")
+
+	assertClauseSerialize(t, joinTable, `foo.table
+INNER JOIN foo.table2 ON ("intCol1" = "intCol2")`)
+
+	require.Equal(t, joinTable.SchemaName(), "foo")
+	require.Equal(t, joinTable.TableName(), "")
+
+	require.Equal(t, len(joinTable.columns()), 2)
+	require.Equal(t, joinTable.columns()[0].Name(), "intCol1")
+	require.Equal(t, joinTable.columns()[1].Name(), "intCol2")
+}