Add support for pg_tablespace_location() function

BemiHQ · Dec 13, 2024 · edebf2b · edebf2b
1 parent ca71c7a
commit edebf2b
Show file tree

Hide file tree

Showing 3 changed files with 15 additions and 11 deletions.
diff --git a/src/query_handler_test.go b/src/query_handler_test.go
@@ -28,11 +28,14 @@ func TestHandleQuery(t *testing.T) {
 			"description": {"setting"},
 			"values":      {"nulls_last"},
 		},
-		// pg_get_partkeydef
 		"SELECT pg_catalog.pg_get_partkeydef(c.oid) FROM pg_catalog.pg_class c LIMIT 1": {
 			"description": {"pg_get_partkeydef"},
 			"values":      {""},
 		},
+		"SELECT pg_tablespace_location(t.oid) loc FROM pg_catalog.pg_tablespace": {
+			"description": {"loc"},
+			"values":      {""},
+		},
 		// PG system tables
 		"SELECT oid, typname AS typename FROM pg_type WHERE typname='geometry' OR typname='geography'": {
 			"description": {"oid", "typename"},

diff --git a/src/select_remapper.go b/src/select_remapper.go
@@ -6,16 +6,6 @@ import (
 	pgQuery "github.com/pganalyze/pg_query_go/v5"
 )
 
-var REMAPPED_CONSTANT_BY_PG_FUNCTION_NAME = map[string]string{
-	"version":                            "PostgreSQL " + PG_VERSION + ", compiled by Bemi",
-	"pg_get_userbyid":                    "bemidb",
-	"pg_get_function_identity_arguments": "",
-	"pg_total_relation_size":             "0",
-	"pg_table_size":                      "0",
-	"pg_indexes_size":                    "0",
-	"pg_get_partkeydef":                  "",
-}
-
 var KNOWN_SET_STATEMENTS = NewSet([]string{
 	"client_encoding",             // SET client_encoding TO 'UTF8'
 	"client_min_messages",         // SET client_min_messages TO 'warning'

diff --git a/src/select_remapper_select.go b/src/select_remapper_select.go
@@ -4,6 +4,17 @@ import (
 	pgQuery "github.com/pganalyze/pg_query_go/v5"
 )
 
+var REMAPPED_CONSTANT_BY_PG_FUNCTION_NAME = map[string]string{
+	"version":                            "PostgreSQL " + PG_VERSION + ", compiled by Bemi",
+	"pg_get_userbyid":                    "bemidb",
+	"pg_get_function_identity_arguments": "",
+	"pg_total_relation_size":             "0",
+	"pg_table_size":                      "0",
+	"pg_indexes_size":                    "0",
+	"pg_get_partkeydef":                  "",
+	"pg_tablespace_location":             "",
+}
+
 type SelectRemapperSelect struct {
 	parserSelect *QueryParserSelect
 	config       *Config