Add support for udt column in information_schema.columns

BemiHQ · Feb 21, 2025 · 1f8953b · 1f8953b
1 parent fda4922
commit 1f8953b
Show file tree

Hide file tree

Showing 4 changed files with 57 additions and 9 deletions.
diff --git a/src/duckdb.go b/src/duckdb.go
@@ -50,8 +50,9 @@ func NewDuckdb(config *Config) *Duckdb {
 		CreatePgCatalogMacroQueries(config),
 		CreateInformationSchemaMacroQueries(config),
 
-		// Create pg-compatible tables
+		// Create pg-compatible tables and views
 		CreatePgCatalogTableQueries(config),
+		CreateInformationSchemaTableQueries(config),
 
 		// Use the public schema
 		[]string{"USE public"},

diff --git a/src/pg_constants.go b/src/pg_constants.go
@@ -16,7 +16,6 @@ const (
 	PG_TABLE_PG_NAMESPACE        = "pg_namespace"
 	PG_TABLE_PG_STAT_USER_TABLES = "pg_stat_user_tables"
 	PG_TABLE_TABLES              = "tables"
-	PG_TABLE_COLUMNS             = "columns"
 
 	PG_VAR_SEARCH_PATH = "search_path"
 )

diff --git a/src/query_handler_test.go b/src/query_handler_test.go
@@ -278,6 +278,11 @@ func TestHandleQuery(t *testing.T) {
 			"types":       {Uint32ToString(pgtype.TextOID)},
 			"values":      {"memory", "public", "test_table"},
 		},
+		"SELECT column_name, udt_schema, udt_name FROM information_schema.columns WHERE table_schema = 'public' ORDER BY ordinal_position LIMIT 1": {
+			"description": {"column_name", "udt_schema", "udt_name"},
+			"types":       {Uint32ToString(pgtype.TextOID), Uint32ToString(pgtype.TextOID), Uint32ToString(pgtype.TextOID)},
+			"values":      {"id", "pg_catalog", "int4"},
+		},
 
 		// DISCARD
 		"DISCARD ALL": {

diff --git a/src/query_remapper_table.go b/src/query_remapper_table.go
@@ -11,6 +11,7 @@ import (
 var MAX_REDUNDANT_PG_NAMESPACE_OID = 1265
 
 var PG_CATALOG_TABLE_NAMES = Set[string]{}
+var PG_INFORMATION_SCHEMA_TABLE_NAMES = Set[string]{}
 
 func CreatePgCatalogTableQueries(config *Config) []string {
 	result := []string{
@@ -46,6 +47,48 @@ func CreatePgCatalogTableQueries(config *Config) []string {
 	return result
 }
 
+func CreateInformationSchemaTableQueries(config *Config) []string {
+	result := []string{
+		// Dynamic views
+		// DuckDB does not support udt_schema, udt_name
+		`CREATE VIEW columns AS
+		SELECT
+			table_catalog, table_schema, table_name, column_name, ordinal_position, column_default, is_nullable, data_type, character_maximum_length, character_octet_length, numeric_precision, numeric_precision_radix, numeric_scale, datetime_precision, interval_type, interval_precision, character_set_catalog, character_set_schema, character_set_name, collation_catalog, collation_schema, collation_name, domain_catalog, domain_schema, domain_name,
+			'` + config.Database + `' AS udt_catalog,
+			'pg_catalog' AS udt_schema,
+			CASE data_type
+			WHEN 'BIGINT' THEN 'int8'
+			WHEN 'BIGINT[]' THEN '_int8'
+			WHEN 'BLOB' THEN 'bytea'
+			WHEN 'BLOB[]' THEN '_bytea'
+			WHEN 'BOOLEAN' THEN 'bool'
+			WHEN 'BOOLEAN[]' THEN '_bool'
+			WHEN 'DATE' THEN 'date'
+			WHEN 'DATE[]' THEN '_date'
+			WHEN 'FLOAT' THEN 'float8'
+			WHEN 'FLOAT[]' THEN '_float8'
+			WHEN 'INTEGER' THEN 'int4'
+			WHEN 'INTEGER[]' THEN '_int4'
+			WHEN 'VARCHAR' THEN 'text'
+			WHEN 'VARCHAR[]' THEN '_text'
+			WHEN 'TIME' THEN 'time'
+			WHEN 'TIME[]' THEN '_time'
+			WHEN 'TIMESTAMP' THEN 'timestamp'
+			WHEN 'TIMESTAMP[]' THEN '_timestamp'
+			WHEN 'UUID' THEN 'uuid'
+			WHEN 'UUID[]' THEN '_uuid'
+			ELSE
+				CASE
+				WHEN starts_with(data_type, 'DECIMAL') THEN 'numeric'
+				END
+			END AS udt_name,
+			scope_catalog, scope_schema, scope_name, maximum_cardinality, dtd_identifier, is_self_referencing, is_identity, identity_generation, identity_start, identity_increment, identity_maximum, identity_minimum, identity_cycle, is_generated, generation_expression, is_updatable
+		FROM information_schema.columns`,
+	}
+	PG_INFORMATION_SCHEMA_TABLE_NAMES = extractTableNames(result)
+	return result
+}
+
 type QueryRemapperTable struct {
 	parserTable         *ParserTable
 	parserWhere         *ParserWhere
@@ -109,16 +152,16 @@ func (remapper *QueryRemapperTable) RemapTable(node *pgQuery.Node) *pgQuery.Node
 		// information_schema.tables -> reload Iceberg tables
 		case PG_TABLE_TABLES:
 			remapper.reloadIceberSchemaTables()
-			return node
-
-		// information_schema.columns -> return hard-coded columns
-		// DuckDB does not support udt_schema, udt_name
-		case PG_TABLE_COLUMNS:
+		}
 
-		// information_schema.* other system tables -> return as is
-		default:
+		// information_schema.table -> main.table
+		if PG_INFORMATION_SCHEMA_TABLE_NAMES.Contains(qSchemaTable.Table) {
+			parser.RemapSchemaToMain(node)
 			return node
 		}
+
+		// information_schema.* other system tables -> return as is
+		return node
 	}
 
 	// public.table -> FROM iceberg_scan('path', skip_schema_inference = true) table