Support View usage for PIVOT operator

Support create/select/drop view on stmt with pivot operator Task: BABEL-4673 Signed-off-by: Yanjie Xu <[email protected]>
amazon-aurora · Jul 29, 2024 · 1f59e37 · 1f59e37
1 parent 1596e1f
commit 1f59e37
Show file tree

Hide file tree

Showing 30 changed files with 6,656 additions and 173 deletions.
diff --git a/contrib/babelfishpg_tsql/runtime/functions.c b/contrib/babelfishpg_tsql/runtime/functions.c
@@ -75,6 +75,7 @@
 #define DATEPART_MIN_VALUE -53690               	/* minimun value for datepart general_integer_datatype */
 #define DATEPART_SMALLMONEY_MAX_VALUE 214748.3647	/* maximum value for datepart smallmoney */
 #define DATEPART_SMALLMONEY_MIN_VALUE -53690		/* minimum value for datepart smallmoney */
+#define PIVOT_METADATA_COUNT 3						/* total of 3 metadata is needed by bbf_pivot function */
 
 typedef enum
 {
@@ -198,9 +199,10 @@ void	   *get_host_id(void);
 
 Datum 		datepart_internal(char *field , Timestamp timestamp , float8 df_tz, bool general_integer_datatype);
 int 		SPI_execute_raw_parsetree(RawStmt *parsetree, const char *sourcetext, bool read_only, long tcount);
-static HTAB *load_categories_hash(RawStmt *cats_sql, const char *sourcetext, MemoryContext per_query_ctx);
+static HTAB *load_categories_hash(RawStmt *cats_sql, const char *sourcetext, const char	*view_uuid, MemoryContext per_query_ctx);
 static Tuplestorestate *get_bbf_pivot_tuplestore(RawStmt 	*sql,
 												const char 	*sourcetext,
+												const char	*view_uuid,
 												const char 	*funcName,
 												HTAB 		*bbf_pivot_hash,
 												TupleDesc 	tupdesc,
@@ -4407,6 +4409,7 @@ Datum
 bbf_pivot(PG_FUNCTION_ARGS)
 {	
 	ReturnSetInfo   *rsinfo = (ReturnSetInfo *) fcinfo->resultinfo;
+	char			*view_uuid = text_to_cstring(PG_GETARG_TEXT_PP(0));
 	TupleDesc		tupdesc;
 	MemoryContext 	per_query_ctx;
 	MemoryContext 	oldcontext;
@@ -4430,32 +4433,82 @@ bbf_pivot(PG_FUNCTION_ARGS)
 				(errcode(ERRCODE_FEATURE_NOT_SUPPORTED),
 				 errmsg("materialize mode required, but it is not allowed in this context")));
 
-	if (fcinfo->context == NULL || !IsA(fcinfo->context, List) || list_length((List *) fcinfo->context) != 3)
-		ereport(ERROR,
-			(errcode(ERRCODE_CHECK_VIOLATION),
-				errmsg("Babelfish PIVOT is not properly initialized.")));
+	bbf_pivot_src_sql = NULL;
+	bbf_pivot_cat_sql = NULL;
+	query_string = NULL;
+	funcName = NULL;
+
+	/*
+	 * if view_uuid is not defined (view_uuid is empty string), then we will get the pivot metadata from
+	 * FunctionCallInfo(fcinfo).
+	 */
+	if (!view_uuid || strlen(view_uuid) == 0)
+	{	
+		if (fcinfo->context == NULL || !IsA(fcinfo->context, List) || list_length((List *) fcinfo->context) != PIVOT_METADATA_COUNT)
+		{
+			ereport(ERROR,
+				(errcode(ERRCODE_CHECK_VIOLATION),
+					errmsg("Babelfish PIVOT is not properly initialized.")));
+		}
+
+		node = list_nth((List *)fcinfo->context, 0);
+		if (!IsA(node, List) 
+				|| !IsA(list_nth((List *)node, 0), String) 
+				|| strcmp(((String *)list_nth((List *)node, 0))->sval, "bbf_pivot_func") != 0)
+		{
+			ereport(ERROR,
+				(errcode(ERRCODE_CHECK_VIOLATION),
+					errmsg("Babelfish PIVOT is not properly initialized.")));
+		}
 
-	node = list_nth((List *)fcinfo->context, 0);
-	if (!IsA(node, List) 
-			|| !IsA(list_nth((List *)node, 0), String) 
-			|| strcmp(((String *)list_nth((List *)node, 0))->sval, "bbf_pivot_func") != 0)
+		pivot_parsetree = (List *) list_nth((List *) fcinfo->context, 1);
+		pivot_extrainfo = (List *) list_nth((List *) fcinfo->context, 2);
+
+		if (!IsA(pivot_parsetree, List) || !IsA(pivot_extrainfo, List))
+			ereport(ERROR,
+				(errcode(ERRCODE_CHECK_VIOLATION),
+					errmsg("Babelfish PIVOT is not properly initialized.")));
+
+		bbf_pivot_src_sql = (RawStmt *) list_nth(pivot_parsetree, 0);
+		bbf_pivot_cat_sql = (RawStmt *) list_nth(pivot_parsetree, 1);
+		query_string = ((String *) list_nth(pivot_extrainfo, 0))->sval;
+		funcName = ((String *) list_nth(pivot_extrainfo, 1))->sval;
+	}
+	else
 	{
-		ereport(ERROR,
-			(errcode(ERRCODE_CHECK_VIOLATION),
-				errmsg("Babelfish PIVOT is not properly initialized.")));
+	   /*
+		* If view_uuid is defined (not empty string), then we know current bbf_pivot function is within
+		* a view stmt.
+		* We will not fetch metadata from FunctionCallInfo (fcinfo) but rather get the 
+		* view_uuid from the function argument and get aggregate function argument from babel-
+		* fish_view_catalog.
+		*/
+		Relation	rel;
+		HeapTuple	tuple;
+		SysScanDesc	scan;
+		ScanKeyData	scanKey[1];
+
+		rel = table_open(get_bbf_pivot_view_oid(), RowExclusiveLock);
+
+		ScanKeyInit(&scanKey[0],
+			Anum_bbf_pivot_view_pivot_view_uuid,
+			BTEqualStrategyNumber, F_TEXTEQ,
+			CStringGetTextDatum(view_uuid));
+
+		scan = systable_beginscan(rel, get_bbf_pivot_view_idx_oid(),
+					false, NULL, 1, scanKey);
+
+		tuple = systable_getnext(scan);
+		if (HeapTupleIsValid(tuple))
+		{
+			bool isnull;
+			funcName = TextDatumGetCString(heap_getattr(tuple, Anum_bbf_pivot_view_agg_func_name, RelationGetDescr(rel), &isnull));
+		}
+
+		systable_endscan(scan);
+		table_close(rel, RowExclusiveLock);
 	}
-	pivot_parsetree = (List *) list_nth((List *) fcinfo->context, 1);
-	pivot_extrainfo = (List *) list_nth((List *) fcinfo->context, 2);
-
-	if (!IsA(pivot_parsetree, List) || !IsA(pivot_extrainfo, List))
-		ereport(ERROR,
-			(errcode(ERRCODE_CHECK_VIOLATION),
-				errmsg("Babelfish PIVOT is not properly initialized.")));
 
-	bbf_pivot_src_sql = (RawStmt *) list_nth(pivot_parsetree, 0);
-	bbf_pivot_cat_sql = (RawStmt *) list_nth(pivot_parsetree, 1);
-	query_string = ((String *) list_nth(pivot_extrainfo, 0))->sval;
-	funcName = ((String *) list_nth(pivot_extrainfo, 1))->sval;
 
 	per_query_ctx = rsinfo->econtext->ecxt_per_query_memory;
 	oldcontext = MemoryContextSwitchTo(per_query_ctx);
@@ -4477,14 +4530,15 @@ bbf_pivot(PG_FUNCTION_ARGS)
 						"bbf_pivot function are not compatible")));
 
 	/* load up the categories hash table */
-	bbf_pivot_hash = load_categories_hash(bbf_pivot_cat_sql, query_string, per_query_ctx);
+	bbf_pivot_hash = load_categories_hash(bbf_pivot_cat_sql, query_string, view_uuid, per_query_ctx);
 
 	/* let the caller know we're sending back a tuplestore */
 	rsinfo->returnMode = SFRM_Materialize;
 
 	/* now go build it */
 	rsinfo->setResult = get_bbf_pivot_tuplestore(bbf_pivot_src_sql,
 												query_string,
+												view_uuid,
 												funcName,
 												bbf_pivot_hash,
 												tupdesc,
@@ -4507,7 +4561,11 @@ bbf_pivot(PG_FUNCTION_ARGS)
  * load up the categories hash table
  */
 static HTAB *
-load_categories_hash(RawStmt *cats_sql, const char * sourcetext, MemoryContext per_query_ctx)
+load_categories_hash(RawStmt *cats_sql, 
+					 const char *sourcetext, 
+					 const char	*view_uuid,
+					 MemoryContext per_query_ctx
+					 )
 {
 	HTAB	   *bbf_pivot_hash;
 	HASHCTL		ctl;
@@ -4534,8 +4592,30 @@ load_categories_hash(RawStmt *cats_sql, const char * sourcetext, MemoryContext p
 		/* internal error */
 		elog(ERROR, "load_categories_hash: SPI_connect returned %d", ret);
 
-	/* Retrieve the category name rows */
-	ret = SPI_execute_raw_parsetree(cats_sql, sourcetext, true, 0);
+	/* 
+	 * Retrieve the category name rows    
+	 *
+	 * If view_uuid is defined, then we will use the provided uuid to form the catefory_sql and 
+	 * get result. category_sql name format is pvt_cv_[uuid]
+	 * 
+	 * if view_uuid is not defined, then we will use the rawparsetree of catefory_sql to get the
+	 * result
+	 */
+	if (!view_uuid || strlen(view_uuid) != 0)
+	{
+		StringInfoData 		buf;
+
+		initStringInfo(&buf);
+		appendStringInfoString(&buf, "select * from ");
+		appendStringInfoString(&buf, BBF_PIVOT_VIEW_CAT_PREFIX);
+		appendStringInfoString(&buf, view_uuid);
+
+		ret = SPI_execute(buf.data, true, 0);
+	}
+	else
+	{
+		ret = SPI_execute_raw_parsetree(cats_sql, sourcetext, true, 0);
+	}
 	tuple_processed = SPI_processed;
 
 	/* Check for qualifying tuples */
@@ -4602,6 +4682,7 @@ load_categories_hash(RawStmt *cats_sql, const char * sourcetext, MemoryContext p
 static Tuplestorestate *
 get_bbf_pivot_tuplestore(RawStmt 	*sql,
 						const char 	*sourcetext,
+						const char	*view_uuid,
 						const char	*funcName,
 						HTAB 		*bbf_pivot_hash,
 						TupleDesc 	tupdesc,
@@ -4623,8 +4704,30 @@ get_bbf_pivot_tuplestore(RawStmt 	*sql,
 		/* internal error */
 		elog(ERROR, "get_bbf_pivot_tuplestore: SPI_connect returned %d", ret);
 
-	/* Now retrieve the bbf_pivot source rows */
-	ret = SPI_execute_raw_parsetree(sql, sourcetext, true, 0);
+	/* 
+	 * Now retrieve the bbf_pivot source rows    
+	 *
+	 * If view_uuid is defined, then we will use the provided uuid to form the source_sql and 
+	 * get result. source_sql name format is pvt_sv_[uuid]
+	 * 
+	 * if view_uuid is not defined, then we will use the rawparsetree of source_sql to get the
+	 * result
+	 */
+	if (!view_uuid || strlen(view_uuid) != 0)
+	{
+		StringInfoData 		buf;
+
+		initStringInfo(&buf);
+		appendStringInfoString(&buf, "select * from ");
+		appendStringInfoString(&buf, BBF_PIVOT_VIEW_SRC_PREFIX);
+		appendStringInfoString(&buf, view_uuid);
+
+		ret = SPI_execute(buf.data, true, 0);
+	}
+	else
+	{
+		ret = SPI_execute_raw_parsetree(sql, sourcetext, true, 0);
+	}
 	tuple_processed = SPI_processed;
 
 	/* Check for qualifying tuples */

diff --git a/contrib/babelfishpg_tsql/sql/object_definition_tsql.sql b/contrib/babelfishpg_tsql/sql/object_definition_tsql.sql
@@ -11,4 +11,15 @@ CREATE TABLE sys.babelfish_view_def (
 );
 GRANT SELECT ON sys.babelfish_view_def TO PUBLIC;
 
-SELECT pg_catalog.pg_extension_config_dump('sys.babelfish_view_def', '');
+CREATE TABLE sys.babelfish_pivot_view
+(
+  dbid SMALLINT NOT NULL,
+  pivot_view_uuid sys.NVARCHAR(128) NOT NULL,
+  schema_name sys.NVARCHAR(128) NOT NULL COLLATE sys.database_default,
+  pivot_view_name sys.NVARCHAR(128) NOT NULL COLLATE sys.database_default,
+  agg_func_name sys.NVARCHAR(128) NOT NULL,
+  PRIMARY KEY(pivot_view_uuid)
+);
+
+SELECT pg_catalog.pg_extension_config_dump('sys.babelfish_view_def', '');
+SELECT pg_catalog.pg_extension_config_dump('sys.babelfish_pivot_view', '');
diff --git a/contrib/babelfishpg_tsql/sql/sys_functions.sql b/contrib/babelfishpg_tsql/sql/sys_functions.sql
@@ -5091,7 +5091,7 @@ END;
 $body$
 LANGUAGE plpgsql STABLE;
 
-CREATE OR REPLACE FUNCTION sys.bbf_pivot()
+CREATE OR REPLACE FUNCTION sys.bbf_pivot(IN arg TEXT)
 RETURNS setof record
 AS 'babelfishpg_tsql', 'bbf_pivot'
 LANGUAGE C STABLE;

diff --git a/contrib/babelfishpg_tsql/sql/upgrades/babelfishpg_tsql--3.6.0--3.7.0.sql b/contrib/babelfishpg_tsql/sql/upgrades/babelfishpg_tsql--3.6.0--3.7.0.sql
@@ -135,6 +135,23 @@ $$
 LANGUAGE 'pltsql';
 GRANT EXECUTE ON PROCEDURE sys.sp_tables TO PUBLIC;
 
+CREATE TABLE sys.babelfish_pivot_view
+(
+  dbid SMALLINT NOT NULL,
+  pivot_view_uuid sys.NVARCHAR(128) NOT NULL,
+  schema_name sys.NVARCHAR(128) NOT NULL COLLATE sys.database_default,
+  pivot_view_name sys.NVARCHAR(128) NOT NULL COLLATE sys.database_default,
+  agg_func_name sys.NVARCHAR(128) NOT NULL,
+  PRIMARY KEY(pivot_view_uuid)
+);
+
+CREATE OR REPLACE FUNCTION sys.bbf_pivot(IN arg TEXT)
+RETURNS setof record
+AS 'babelfishpg_tsql', 'bbf_pivot'
+LANGUAGE C STABLE;
+
+SELECT pg_catalog.pg_extension_config_dump('sys.babelfish_pivot_view', '');
+
 -- Update deprecated object_id function(s) since left function now restricts TEXT datatype
 DO $$
 BEGIN