Fix datatype handling for plpgsql function parameters (#256)

* Fix datatype handling for plpgsql function parameters Currently all plpgql function parameters are returned as UNKNOWN. This patch changes function parameters to return actual type names used in the function definition. This patch also fixes a problem with cursor variables parsing and adds a test for it. * Resolve `quote_qualified_identifier` symbol --------- Co-authored-by: msepga <[email protected]>
pganalyze · Sep 24, 2024 · 680f5ee · 680f5ee
1 parent 972c748
commit 680f5ee
Show file tree

Hide file tree

Showing 6 changed files with 139 additions and 25 deletions.
diff --git a/scripts/extract_source.rb b/scripts/extract_source.rb
@@ -525,7 +525,51 @@ def write_out
 # Mocks REQUIRED for PL/pgSQL parsing
 runner.mock('format_type_be', 'char * format_type_be(Oid type_oid) { return pstrdup("-"); }')
 runner.mock('build_row_from_class', 'static PLpgSQL_row *build_row_from_class(Oid classOid) { return NULL; }')
-runner.mock('plpgsql_build_datatype', 'PLpgSQL_type * plpgsql_build_datatype(Oid typeOid, int32 typmod, Oid collation, TypeName *origtypname) { PLpgSQL_type *typ; typ = (PLpgSQL_type *) palloc0(sizeof(PLpgSQL_type)); typ->typname = pstrdup("UNKNOWN"); typ->ttype = PLPGSQL_TTYPE_SCALAR; return typ; }')
+runner.mock('plpgsql_build_datatype', %(
+PLpgSQL_type * plpgsql_build_datatype(Oid typeOid, int32 typmod, Oid collation, TypeName *origtypname)
+{
+	PLpgSQL_type *typ;
+	char *ident = NULL, *ns = NULL;
+	typ = (PLpgSQL_type *) palloc0(sizeof(PLpgSQL_type));
+
+	typ->ttype = PLPGSQL_TTYPE_SCALAR;
+	typ->atttypmod = typmod;
+	typ->collation = collation;
+
+	if (origtypname) {
+		typ->typoid = origtypname->typeOid;
+
+		if (list_length(origtypname->names) == 1) {
+			ident = linitial_node(String, origtypname->names)->sval;
+		} else if (list_length(origtypname->names) == 2) {
+			ns = linitial_node(String, origtypname->names)->sval;
+			ident = lsecond_node(String, origtypname->names)->sval;
+		}
+	} else {
+		typ->typoid = typeOid;
+		ns = "pg_catalog";
+		switch(typeOid)
+		{
+			case BOOLOID:
+				ident = "boolean";
+				break;
+			case INT4OID:
+				ident = "integer";
+				break;
+			case TEXTOID:
+				ident = "text";
+				break;
+			case REFCURSOROID:
+				ident = "refcursor";
+				break;
+		}
+	}
+	if (ident) {
+		typ->typname = quote_qualified_identifier(ns, ident);
+	}
+	return typ;
+}
+))
 runner.mock('parse_datatype', 'static PLpgSQL_type * parse_datatype(const char *string, int location) { PLpgSQL_type *typ; typ = (PLpgSQL_type *) palloc0(sizeof(PLpgSQL_type)); typ->typname = pstrdup(string); typ->ttype = strcmp(string, "RECORD") == 0 ? PLPGSQL_TTYPE_REC : PLPGSQL_TTYPE_SCALAR; return typ; }')
 runner.mock('get_collation_oid', 'Oid get_collation_oid(List *name, bool missing_ok) { return -1; }')
 runner.mock('plpgsql_parse_wordtype', 'PLpgSQL_type * plpgsql_parse_wordtype(char *ident) { return NULL; }')
@@ -582,6 +626,7 @@ def write_out
 # PL/pgSQL Parsing
 runner.deep_resolve('plpgsql_compile_inline')
 runner.deep_resolve('plpgsql_free_function_memory')
+runner.deep_resolve('quote_qualified_identifier')
 
 # Basic Postgres needed to call parser
 runner.deep_resolve('SetDatabaseEncoding')

diff --git a/src/pg_query_parse_plpgsql.c b/src/pg_query_parse_plpgsql.c
@@ -226,7 +226,7 @@ static PLpgSQL_function *compile_create_function_stmt(CreateFunctionStmt* stmt)
 			PLpgSQL_variable *argvariable;
 			PLpgSQL_nsitem_type argitemtype;
 			snprintf(buf, sizeof(buf), "$%d", foreach_current_index(lc) + 1);
-			argdtype = plpgsql_build_datatype(UNKNOWNOID, -1, InvalidOid, NULL);
+			argdtype = plpgsql_build_datatype(UNKNOWNOID, -1, InvalidOid, param->argType);
 			argvariable = plpgsql_build_variable(param->name ? param->name : buf, 0, argdtype, false);
 			argitemtype = argvariable->dtype == PLPGSQL_DTYPE_VAR ? PLPGSQL_NSTYPE_VAR : PLPGSQL_NSTYPE_REC;
 			plpgsql_ns_additem(argitemtype, argvariable->dno, buf);

diff --git a/src/postgres/src_backend_utils_adt_ruleutils.c b/src/postgres/src_backend_utils_adt_ruleutils.c
@@ -2,6 +2,7 @@
  * Symbols referenced in this file:
  * - quote_identifier
  * - quote_all_identifiers
+ * - quote_qualified_identifier
  *--------------------------------------------------------------------
  */
 
@@ -1718,7 +1719,18 @@ quote_identifier(const char *ident)
  * Return a name of the form qualifier.ident, or just ident if qualifier
  * is NULL, quoting each component if necessary.  The result is palloc'd.
  */
+char *
+quote_qualified_identifier(const char *qualifier,
+						   const char *ident)
+{
+	StringInfoData buf;
 
+	initStringInfo(&buf);
+	if (qualifier)
+		appendStringInfo(&buf, "%s.", quote_identifier(qualifier));
+	appendStringInfoString(&buf, quote_identifier(ident));
+	return buf.data;
+}
 
 /*
  * get_relation_name

diff --git a/src/postgres/src_pl_plpgsql_src_pl_comp.c b/src/postgres/src_pl_plpgsql_src_pl_comp.c
@@ -879,7 +879,51 @@ plpgsql_build_recfield(PLpgSQL_rec *rec, const char *fldname)
  * It can be NULL if the type could not be a composite type, or if it was
  * identified by OID to begin with (e.g., it's a function argument type).
  */
-PLpgSQL_type * plpgsql_build_datatype(Oid typeOid, int32 typmod, Oid collation, TypeName *origtypname) { PLpgSQL_type *typ; typ = (PLpgSQL_type *) palloc0(sizeof(PLpgSQL_type)); typ->typname = pstrdup("UNKNOWN"); typ->ttype = PLPGSQL_TTYPE_SCALAR; return typ; }
+
+PLpgSQL_type * plpgsql_build_datatype(Oid typeOid, int32 typmod, Oid collation, TypeName *origtypname)
+{
+	PLpgSQL_type *typ;
+	char *ident = NULL, *ns = NULL;
+	typ = (PLpgSQL_type *) palloc0(sizeof(PLpgSQL_type));
+
+	typ->ttype = PLPGSQL_TTYPE_SCALAR;
+	typ->atttypmod = typmod;
+	typ->collation = collation;
+
+	if (origtypname) {
+		typ->typoid = origtypname->typeOid;
+
+		if (list_length(origtypname->names) == 1) {
+			ident = linitial_node(String, origtypname->names)->sval;
+		} else if (list_length(origtypname->names) == 2) {
+			ns = linitial_node(String, origtypname->names)->sval;
+			ident = lsecond_node(String, origtypname->names)->sval;
+		}
+	} else {
+		typ->typoid = typeOid;
+		ns = "pg_catalog";
+		switch(typeOid)
+		{
+			case BOOLOID:
+				ident = "boolean";
+				break;
+			case INT4OID:
+				ident = "integer";
+				break;
+			case TEXTOID:
+				ident = "text";
+				break;
+			case REFCURSOROID:
+				ident = "refcursor";
+				break;
+		}
+	}
+	if (ident) {
+		typ->typname = quote_qualified_identifier(ns, ident);
+	}
+	return typ;
+}
+
 
 
 /*