Support BCP on TSQL #temp table

contrib/babelfishpg_tsql/runtime/functions.c

@@ -25,6 +25,7 @@
 #include "catalog/pg_namespace.h"
 #include "catalog/pg_type.h"
 #include "catalog/pg_attrdef.h"
+#include "catalog/pg_attribute.h"
 #include "catalog/pg_depend.h"
 #include "commands/dbcommands.h"
 #include "commands/extension.h"
@@ -135,6 +136,7 @@ PG_FUNCTION_INFO_V1(servername);
 PG_FUNCTION_INFO_V1(servicename);
 PG_FUNCTION_INFO_V1(xact_state);
 PG_FUNCTION_INFO_V1(get_enr_list);
+PG_FUNCTION_INFO_V1(get_enr_attributes);
 PG_FUNCTION_INFO_V1(tsql_random);
 PG_FUNCTION_INFO_V1(timezone_mapping);
 PG_FUNCTION_INFO_V1(pltsql_timezone_mapping_pg_to_windows);
@@ -6094,3 +6096,136 @@ openxml_simple(PG_FUNCTION_ARGS)
 #endif							/* USE_LIBXML */
 }
 
+/*
+ * get_enr_attributes - Get pg_attribute rows for temp tables (ENR and non-ENR)
+ *
+ * This function returns all pg_attribute columns for a temp table.
+ * It works for both ENR temp tables (reads from ENR cache) and non-ENR temp tables
+ * (reads from actual pg_attribute catalog).
+ */
+Datum
+get_enr_attributes(PG_FUNCTION_ARGS)
+{
+	char	   *table_name_input = PG_ARGISNULL(0) ? NULL : text_to_cstring(PG_GETARG_TEXT_PP(0));
+	ReturnSetInfo *rsinfo = (ReturnSetInfo *) fcinfo->resultinfo;
+	TupleDesc	tupdesc;
+	Tuplestorestate *tupstore;
+	MemoryContext per_query_ctx;
+	MemoryContext oldcontext;
+	EphemeralNamedRelation enr = NULL;
+	char	   *table_name;
+	char	   *temp_name_ptr;
+	Relation	pg_attribute_rel;
+	Oid			relid = InvalidOid;
+	bool		is_enr = false;
+	int			i;
+
+	/* check to see if caller supports us returning a tuplestore */
+	if (rsinfo == NULL || !IsA(rsinfo, ReturnSetInfo))
+		ereport(ERROR,
+				(errcode(ERRCODE_FEATURE_NOT_SUPPORTED),
+				 errmsg("set-valued function called in context that cannot accept a set")));
+	if (!(rsinfo->allowedModes & SFRM_Materialize))
+		ereport(ERROR,
+				(errcode(ERRCODE_FEATURE_NOT_SUPPORTED),
+				 errmsg("materialize mode required, but it is not allowed in this context")));
+
+	if (!table_name_input)
+		PG_RETURN_NULL();
+
+	/* Handle formats like ".[#TempTable]" or "#TempTable" to find the # */
+	temp_name_ptr = strchr(table_name_input, '#');
+	if (!temp_name_ptr)
+		PG_RETURN_NULL();
+
+	/* Lowercase the table name and strip trailing "]" */
+	table_name = pstrdup(temp_name_ptr);
+	for (i = 0; table_name[i]; i++)
+	{
+		if (table_name[i] == ']')
+		{
+			table_name[i] = '\0';
+			break;
+		}
+		table_name[i] = tolower((unsigned char) table_name[i]);
+	}
+
+	/* Try ENR first */
+	if (currentQueryEnv != NULL)
+	{
+		enr = get_ENR(currentQueryEnv, table_name, true);
+		if (enr != NULL && enr->md.enrtype == ENR_TSQL_TEMP)
+		{
+			relid = enr->md.reliddesc;
+			is_enr = true;
+		}
+	}
+
+	/* Fallback: search pg_temp schema for non-ENR temp tables */
+	if (!is_enr)
+	{
+		Oid temp_ns = LookupNamespaceNoError("pg_temp");
+		if (OidIsValid(temp_ns))
+			relid = get_relname_relid(table_name, temp_ns);
+	}
+
+	if (!OidIsValid(relid))
+		PG_RETURN_NULL();
+
+	/* Setup return - use pg_attribute's tuple descriptor */
+	per_query_ctx = rsinfo->econtext->ecxt_per_query_memory;
+	oldcontext = MemoryContextSwitchTo(per_query_ctx);
+
+	pg_attribute_rel = table_open(AttributeRelationId, AccessShareLock);
+	tupdesc = CreateTupleDescCopy(RelationGetDescr(pg_attribute_rel));
+
+	tupstore = tuplestore_begin_heap(rsinfo->allowedModes & SFRM_Materialize_Random,
+									 false, work_mem);
+
+	rsinfo->returnMode = SFRM_Materialize;
+	rsinfo->setResult = tupstore;
+	rsinfo->setDesc = BlessTupleDesc(tupdesc);
+
+	MemoryContextSwitchTo(oldcontext);
+
+	if (is_enr)
+	{
+		/* ENR path: return tuples directly from ENR cache */
+		ListCell   *lc;
+
+		foreach(lc, enr->md.cattups[ENR_CATTUP_ATTRIBUTE])
+		{
+			HeapTuple	tup = (HeapTuple) lfirst(lc);
+
+			tuplestore_puttuple(tupstore, tup);
+		}
+	}
+	else
+	{
+		/* Non-ENR path: scan pg_attribute catalog */
+		ScanKeyData skey[1];
+		SysScanDesc scan;
+		HeapTuple	tup;
+
+		ScanKeyInit(&skey[0],
+					Anum_pg_attribute_attrelid,
+					BTEqualStrategyNumber, F_OIDEQ,
+					ObjectIdGetDatum(relid));
+
+		scan = systable_beginscan(pg_attribute_rel, AttributeRelidNumIndexId,
+								  true, NULL, 1, skey);
+
+		while (HeapTupleIsValid(tup = systable_getnext(scan)))
+		{
+			tuplestore_puttuple(tupstore, tup);
+		}
+
+		systable_endscan(scan);
+	}
+
+	table_close(pg_attribute_rel, AccessShareLock);
+
+	tuplestore_donestoring(tupstore);
+
+	PG_RETURN_NULL();
+}

contrib/babelfishpg_tsql/sql/babelfishpg_tsql.sql

@@ -802,24 +802,62 @@ CREATE OR REPLACE VIEW sys.spt_tablecollations_view AS
         c.is_sparse = 0;
 GRANT SELECT ON sys.spt_tablecollations_view TO PUBLIC;
 
+-- Function to get pg_attribute rows for temp tables (ENR and non-ENR)
+CREATE OR REPLACE FUNCTION sys.babelfish_get_enr_attributes(IN table_name sys.varchar(4000))
+RETURNS SETOF pg_catalog.pg_attribute
+AS 'babelfishpg_tsql', 'get_enr_attributes'
+LANGUAGE C STABLE PARALLEL UNSAFE;
+GRANT EXECUTE ON FUNCTION sys.babelfish_get_enr_attributes(IN sys.varchar(4000)) TO PUBLIC;
+
+-- Wrapper function for sp_tablecollations_100 that uses the babelfish_get_enr_attributes function
+CREATE OR REPLACE FUNCTION sys.sp_tablecollations_100_enr(IN table_name sys.varchar(4000))
+RETURNS TABLE(colid INT, name sys.varchar, collation_name sys.nvarchar(128))
+AS $$
+    SELECT 
+        CAST(a.attnum AS INT) AS colid,
+        CAST(a.attname AS sys.varchar) AS name,
+        CAST(c.collname AS sys.nvarchar(128)) AS collation_name
+    FROM sys.babelfish_get_enr_attributes(table_name) a
+    LEFT JOIN pg_catalog.pg_collation c ON a.attcollation = c.oid
+    WHERE a.attnum > 0 AND NOT a.attisdropped
+    ORDER BY a.attnum;
+$$
+LANGUAGE SQL STABLE PARALLEL UNSAFE;
+GRANT EXECUTE ON FUNCTION sys.sp_tablecollations_100_enr(IN sys.varchar(4000)) TO PUBLIC;
+
 -- We are limited by what postgres procedures can return here, but IEW may not
 -- need it for initial compatibility
+-- Modified to handle temp tables (starting with #) using ENR function
 CREATE OR REPLACE PROCEDURE sys.sp_tablecollations_100
 (
     IN "@object" nvarchar(4000)
 )
 AS $$
 BEGIN
-    select
-        s_tcv.colid         AS colid,
-        s_tcv.name          AS name,
-        s_tcv.tds_collation_100 AS tds_collation,
-        s_tcv.collation_100 AS collation
-    from
-        sys.spt_tablecollations_view s_tcv
-    where
-        s_tcv.object_id = (SELECT sys.object_id(@object))
-    order by colid;
+    -- Check if this is a temp table (starts with # or [# or contains .# or .[#)
+    IF LEFT(@object, 1) = '#' OR LEFT(@object, 2) = '[#' OR @object LIKE '%.#%' OR @object LIKE '%.[[]#%'
+    BEGIN
+        -- Use ENR function for temp tables
+        SELECT
+            t.colid AS colid,
+            t.name AS name,
+            CAST(CollationProperty(t.collation_name, 'tdscollation') AS sys.binary(5)) AS tds_collation,
+            t.collation_name AS collation
+        FROM sys.sp_tablecollations_100_enr(@object) t
+        ORDER BY t.colid;
+    END
+    ELSE
+    BEGIN
+        -- Existing logic for regular tables
+        SELECT
+            s_tcv.colid AS colid,
+            s_tcv.name AS name,
+            s_tcv.tds_collation_100 AS tds_collation,
+            s_tcv.collation_100 AS collation
+        FROM sys.spt_tablecollations_view s_tcv
+        WHERE s_tcv.object_id = (SELECT sys.object_id(@object))
+        ORDER BY colid;
+    END
 END;
 $$
 LANGUAGE 'pltsql';

contrib/babelfishpg_tsql/sql/upgrades/babelfishpg_tsql--5.5.0--5.6.0.sql

@@ -360,6 +360,65 @@ END;
 $$;
 GRANT EXECUTE ON FUNCTION sys.textsize() TO PUBLIC;
 
+-- Function to get pg_attribute rows for temp tables (ENR and non-ENR)
+CREATE OR REPLACE FUNCTION sys.babelfish_get_enr_attributes(IN table_name sys.varchar(4000))
+RETURNS SETOF pg_catalog.pg_attribute
+AS 'babelfishpg_tsql', 'get_enr_attributes'
+LANGUAGE C STABLE PARALLEL UNSAFE;
+GRANT EXECUTE ON FUNCTION sys.babelfish_get_enr_attributes(IN sys.varchar(4000)) TO PUBLIC;
+
+-- Wrapper function for sp_tablecollations_100 that uses the babelfish_get_enr_attributes function
+CREATE OR REPLACE FUNCTION sys.sp_tablecollations_100_enr(IN table_name sys.varchar(4000))
+RETURNS TABLE(colid INT, name sys.varchar, collation_name sys.nvarchar(128))
+AS $$
+    SELECT 
+        CAST(a.attnum AS INT) AS colid,
+        CAST(a.attname AS sys.varchar) AS name,
+        CAST(c.collname AS sys.nvarchar(128)) AS collation_name
+    FROM sys.babelfish_get_enr_attributes(table_name) a
+    LEFT JOIN pg_catalog.pg_collation c ON a.attcollation = c.oid
+    WHERE a.attnum > 0 AND NOT a.attisdropped
+    ORDER BY a.attnum;
+$$
+LANGUAGE SQL STABLE PARALLEL UNSAFE;
+GRANT EXECUTE ON FUNCTION sys.sp_tablecollations_100_enr(IN sys.varchar(4000)) TO PUBLIC;
+
+-- Modify sys.sp_tablecollations_100 to handle temp tables (starting with #)
+-- This enables SqlBulkCopy to work with temp tables by providing column collation metadata.
+CREATE OR REPLACE PROCEDURE sys.sp_tablecollations_100
+(
+    IN "@object" nvarchar(4000)
+)
+AS $$
+BEGIN
+    -- Check if this is a temp table (starts with # or [# or contains .# or .[#)
+    IF LEFT(@object, 1) = '#' OR LEFT(@object, 2) = '[#' OR @object LIKE '%.#%' OR @object LIKE '%.[[]#%'
+    BEGIN
+        -- Use ENR function for temp tables
+        SELECT
+            t.colid AS colid,
+            t.name AS name,
+            CAST(CollationProperty(t.collation_name, 'tdscollation') AS sys.binary(5)) AS tds_collation,
+            t.collation_name AS collation
+        FROM sys.sp_tablecollations_100_enr(@object) t
+        ORDER BY t.colid;
+    END
+    ELSE
+    BEGIN
+        -- Existing logic for regular tables
+        SELECT
+            s_tcv.colid AS colid,
+            s_tcv.name AS name,
+            s_tcv.tds_collation_100 AS tds_collation,
+            s_tcv.collation_100 AS collation
+        FROM sys.spt_tablecollations_view s_tcv
+        WHERE s_tcv.object_id = (SELECT sys.object_id(@object))
+        ORDER BY colid;
+    END
+END;
+$$
+LANGUAGE 'pltsql';
+
 -- Drops the temporary procedure used by the upgrade script.
 -- Please have this be one of the last statements executed in this upgrade script.
 DROP PROCEDURE sys.babelfish_drop_deprecated_object(varchar, varchar, varchar, varchar);

contrib/babelfishpg_tsql/src/procedures.c

@@ -36,6 +36,7 @@
 #include "utils/rel.h"
 #include "utils/syscache.h"
 #include "utils/fmgroids.h"
+#include "utils/lsyscache.h"
 #include "utils/formatting.h"
 #include "pltsql_instr.h"
 #include "pltsql.h"
@@ -55,6 +56,7 @@
 #include "catalog.h"
 #include "catalog/toasting.h"
 #include "extendedproperty.h"
+#include "collation.h"
 #include "multidb.h"
 #include "session.h"
 #include "rolecmds.h"
@@ -5115,4 +5117,4 @@ tsql_openxml_get_colpattern(PG_FUNCTION_ARGS)
 	}
 
 	PG_RETURN_TEXT_P(cstring_to_text(xpath_expr));
-}
+}

test/JDBC/expected/BABEL-5264-vu-cleanup.out

@@ -0,0 +1,12 @@
+-- BABEL-5264: Cleanup for sp_tablecollations_100 tests
+-- Cleanup UDTs created for non-ENR tests
+DROP TYPE IF EXISTS dbo.TestVarcharType1
+GO
+DROP TYPE IF EXISTS dbo.TestVarcharType2
+GO
+DROP TYPE IF EXISTS dbo.TestVarcharType3
+GO
+DROP TYPE IF EXISTS dbo.TestVarcharType4
+GO
+DROP TYPE IF EXISTS dbo.TestIntType
+GO

test/JDBC/expected/BABEL-5264-vu-prepare.out

@@ -0,0 +1,12 @@
+-- BABEL-5264: Test setup for sp_tablecollations_100 functions
+-- Create UDTs for non-ENR temp table tests
+CREATE TYPE dbo.TestVarcharType1 FROM varchar(100)
+GO
+CREATE TYPE dbo.TestVarcharType2 FROM varchar(50)
+GO
+CREATE TYPE dbo.TestVarcharType3 FROM varchar(50)
+GO
+CREATE TYPE dbo.TestVarcharType4 FROM varchar(50)
+GO
+CREATE TYPE dbo.TestIntType FROM int
+GO