From 2f29f19becee27287dbb5eea74b75bf9f63ab7cc Mon Sep 17 00:00:00 2001
From: Kyle Conroy <kyle@conroy.org>
Date: Mon, 16 Dec 2019 13:01:15 -0800
Subject: [PATCH] internal/dinosql: Support the pg_temp schema

The schema is removed from the catalog so that queries can not
reference it, since it will be empty in a new session.
---
 internal/catalog/build_test.go | 29 ++++++++++++++++++++++++++---
 internal/dinosql/parser.go     |  5 +++++
 internal/pg/catalog.go         |  9 +++++++++
 3 files changed, 40 insertions(+), 3 deletions(-)

diff --git a/internal/catalog/build_test.go b/internal/catalog/build_test.go
index 76416eca4b..07142a0382 100644
--- a/internal/catalog/build_test.go
+++ b/internal/catalog/build_test.go
@@ -481,6 +481,27 @@ func TestUpdate(t *testing.T) {
 				},
 			},
 		},
+		{
+			`
+			CREATE TABLE pg_temp.migrate (val INT);
+			INSERT INTO pg_temp.migrate (val) SELECT val FROM old;
+			INSERT INTO new (val) SELECT val FROM pg_temp.migrate;
+			`,
+			pg.Catalog{
+				Schemas: map[string]pg.Schema{
+					"pg_temp": {
+						Tables: map[string]pg.Table{
+							"migrate": pg.Table{
+								Name: "migrate",
+								Columns: []pg.Column{
+									{Name: "val", DataType: "pg_catalog.int4", NotNull: false, Table: pg.FQN{Schema: "pg_temp", Rel: "migrate"}},
+								},
+							},
+						},
+					},
+				},
+			},
+		},
 	} {
 		test := tc
 		t.Run(strconv.Itoa(i), func(t *testing.T) {
@@ -490,10 +511,12 @@ func TestUpdate(t *testing.T) {
 				t.Fatal(err)
 			}
 
-			delete(c.Schemas, "pg_catalog")
-			delete(test.c.Schemas, "pg_catalog")
+			expected := pg.NewCatalog()
+			for name, schema := range test.c.Schemas {
+				expected.Schemas[name] = schema
+			}
 
-			if diff := cmp.Diff(test.c, c, cmpopts.EquateEmpty()); diff != "" {
+			if diff := cmp.Diff(expected, c, cmpopts.EquateEmpty()); diff != "" {
 				t.Log(test.stmt)
 				t.Errorf("catalog mismatch:\n%s", diff)
 			}
diff --git a/internal/dinosql/parser.go b/internal/dinosql/parser.go
index 486ec8fa3b..86df0f4b9c 100644
--- a/internal/dinosql/parser.go
+++ b/internal/dinosql/parser.go
@@ -108,6 +108,11 @@ func ParseCatalog(schema string) (core.Catalog, error) {
 			}
 		}
 	}
+
+	// The pg_temp schema is scoped to the current session. Remove it from the
+	// catalog so that other queries can not read from it.
+	delete(c.Schemas, "pg_temp")
+
 	if len(merr.Errs) > 0 {
 		return c, merr
 	}
diff --git a/internal/pg/catalog.go b/internal/pg/catalog.go
index 7e0ce9cbf5..7b6f94f25a 100644
--- a/internal/pg/catalog.go
+++ b/internal/pg/catalog.go
@@ -5,6 +5,15 @@ func NewCatalog() Catalog {
 		Schemas: map[string]Schema{
 			"public":     NewSchema(),
 			"pg_catalog": pgCatalog(),
+			// Likewise, the current session's temporary-table schema, pg_temp_nnn, is
+			// always searched if it exists. It can be explicitly listed in the path by
+			// using the alias pg_temp. If it is not listed in the path then it is
+			// searched first (even before pg_catalog). However, the temporary schema is
+			// only searched for relation (table, view, sequence, etc) and data type
+			// names. It is never searched for function or operator names.
+			//
+			// https://www.postgresql.org/docs/current/runtime-config-client.html
+			"pg_temp": NewSchema(),
 		},
 	}
 }