Add support for the method .query() for XML data type (#4845)

Currently, Babelfish supports .exist() and .value() XML methods but .query() was unsupported and threw an error.
With this change, .query(xpath) is now supported and returns XML results (for XPath 1.0 compatible queries).

The implementation follows the same 3-stage architecture like .exist() and .value():

ANTLR rewriting: @x.query('/path') is rewritten to sys.bbf_xmlquery('/path', @x)
Grammar: BBF_XMLQUERY is mapped in gram-tsql-rule.y
PL/pgSQL function sys.bbf_xmlquery(): validates the input type, checks QUOTED_IDENTIFIER, calls PostgreSQL's
xpath(), and returns the result as XML

Task: BABEL-5224
Signed-off-by: Japleen Kaur <amjj@amazon.com>

Author: JapleenKaur3

contrib/babelfishpg_tsql/runtime/functions.c

@@ -6094,3 +6094,115 @@ openxml_simple(PG_FUNCTION_ARGS)
 #endif							/* USE_LIBXML */
 }
 
+
+PG_FUNCTION_INFO_V1(bbf_xmlquery);
+
+/*
+ * bbf_xmlquery - C implementation of XML .query() method
+ *
+ * Signature:
+ *   sys.bbf_xmlquery(xpath_pattern TEXT, xml_element ANYELEMENT)
+ *
+ * Returns XML result of evaluating the XPath expression against the input.
+ * Returns empty XML if no nodes match.
+ *
+ * Validates:
+ *   - Input must be XML type (or UDT based on XML)
+ *   - QUOTED_IDENTIFIER must be ON
+ */
+Datum
+bbf_xmlquery(PG_FUNCTION_ARGS)
+{
+	text	   *xpath_expr;
+	Datum		xml_datum;
+	Oid			arg_type;
+	Oid			immediate_base_type;
+	ArrayType  *namespaces;
+	Datum		xpath_result;
+	ArrayType  *result_arr;
+	Datum	   *elems;
+	bool	   *nulls;
+	int			nitems;
+	StringInfoData buf;
+	int			i;
+
+	xpath_expr = PG_GETARG_TEXT_PP(0);
+	xml_datum = PG_GETARG_DATUM(1);
+
+	/* Lookup the datatype of the supplied argument */
+	arg_type = get_fn_expr_argtype(fcinfo->flinfo, 1);
+
+	/* UDT handling: resolve to immediate base type if it's a UDT */
+	immediate_base_type = get_immediate_base_type_of_UDT_internal(arg_type);
+	if (OidIsValid(immediate_base_type))
+		arg_type = immediate_base_type;
+
+	if (arg_type != XMLOID)
+	{
+		const char *typname = NULL;
+
+		/* Get T-SQL type name for error message */
+		if (common_utility_plugin_ptr)
+			typname = (*common_utility_plugin_ptr->resolve_pg_type_to_tsql)(arg_type);
+		if (typname == NULL)
+			typname = format_type_be(arg_type);
+
+		ereport(ERROR,
+				(errcode(ERRCODE_DATATYPE_MISMATCH),
+				 errmsg("Cannot call methods on %s.", typname)));
+	}
+
+	/* Check QUOTED_IDENTIFIER setting (required for XML methods in T-SQL) */
+	if (!pltsql_quoted_identifier)
+		ereport(ERROR,
+				(errcode(ERRCODE_INVALID_PARAMETER_VALUE),
+				 errmsg("SELECT failed because the following SET options have "
+						"incorrect settings: 'QUOTED_IDENTIFIER'. Verify that "
+						"SET options are correct for XML data type methods.")));
+
+	/*
+	 * Call the built-in xpath(text, xml, text[][]) directly with an empty
+	 * namespace array. Returns xml[] (array of XML fragments).
+	 *
+	 * TODO: when WITH XMLNAMESPACES is supported, populate this array with
+	 * the declared (prefix, uri) pairs from the active namespace context.
+	 */
+	namespaces = construct_empty_array(TEXTOID);
+	xpath_result = DirectFunctionCall3(xpath,
+									   PointerGetDatum(xpath_expr),
+									   xml_datum,
+									   PointerGetDatum(namespaces));
+
+	result_arr = DatumGetArrayTypeP(xpath_result);
+
+	/* Deconstruct the result array */
+	deconstruct_array(result_arr, XMLOID, -1, false, TYPALIGN_INT,
+					  &elems, &nulls, &nitems);
+
+	/* Empty result → return empty string as XML (matches T-SQL behavior) */
+	if (nitems == 0)
+		PG_RETURN_XML_P((xmltype *) cstring_to_text(""));
+
+	/* Single result → return directly (common fast path) */
+	if (nitems == 1 && !nulls[0])
+		PG_RETURN_DATUM(elems[0]);
+
+	/*
+	 * Multiple results - concatenate all XML fragments.
+	 * Equivalent to: SELECT xmlagg(x) FROM unnest(result_set) AS x
+	 */
+	initStringInfo(&buf);
+	for (i = 0; i < nitems; i++)
+	{
+		if (!nulls[i])
+		{
+			text *fragment = DatumGetTextPP(elems[i]);
+
+			appendBinaryStringInfo(&buf,
+								   VARDATA_ANY(fragment),
+								   VARSIZE_ANY_EXHDR(fragment));
+		}
+	}
+
+	PG_RETURN_XML_P((xmltype *) cstring_to_text_with_len(buf.data, buf.len));
+}

contrib/babelfishpg_tsql/sql/sys_functions.sql

@@ -159,11 +159,17 @@ BEGIN
         result := pg_catalog.replace(result, '"', '"');
         result := pg_catalog.replace(result, '&', '&');
         return result;
-    END IF; 
+    END IF;
 END
 $BODY$
 LANGUAGE plpgsql STABLE STRICT PARALLEL SAFE;
 
+-- helper function for XML QUERY(xpath)
+CREATE OR REPLACE FUNCTION sys.bbf_xmlquery(xpath_pattern TEXT, xml_element ANYELEMENT)
+RETURNS XML
+AS 'babelfishpg_tsql', 'bbf_xmlquery'
+LANGUAGE C STABLE STRICT PARALLEL SAFE;
+
 -- SELECT FOR JSON
 CREATE OR REPLACE FUNCTION sys.tsql_query_to_json_sfunc(
     state INTERNAL,

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

@@ -66,6 +66,12 @@ $$
     end;
 $$;
 
+-- helper function for XML QUERY(xpath)
+CREATE OR REPLACE FUNCTION sys.bbf_xmlquery(xpath_pattern TEXT, xml_element ANYELEMENT)
+RETURNS XML
+AS 'babelfishpg_tsql', 'bbf_xmlquery'
+LANGUAGE C STABLE STRICT PARALLEL SAFE;
+
 -- BABELFISH_FUNCTION_EXT (antlr_parse_cache)
 SET allow_system_table_mods = on;
 ALTER TABLE sys.babelfish_function_ext ADD COLUMN IF NOT EXISTS antlr_parse_cache_tree TEXT DEFAULT NULL;

contrib/babelfishpg_tsql/src/pltsql_instr.h

@@ -45,7 +45,6 @@ typedef enum PgTsqlInstrMetricType
 	INSTR_UNSUPPORTED_TSQL_FREETEXT,
 	INSTR_UNSUPPORTED_TSQL_NEXT_VALUE_FOR,
 	INSTR_UNSUPPORTED_TSQL_XML_NODES,
-	INSTR_UNSUPPORTED_TSQL_XML_QUERY,
 	INSTR_UNSUPPORTED_TSQL_XML_MODIFY,
 	INSTR_UNSUPPORTED_TSQL_ALTER_FUNCTION,
 	INSTR_UNSUPPORTED_TSQL_ALTER_FUNCTION_ENCRYPTION_OPTION,

contrib/babelfishpg_tsql/src/tsqlIface.cpp

@@ -1050,7 +1050,7 @@ class tsqlCommonMutator : public TSqlParserBaseListener
 
 	void exitXml_func_arg(TSqlParser::Xml_func_argContext *ctx) override
 	{
-		if (ctx->EXIST() || ctx->VALUE())
+		if (ctx->EXIST() || ctx->VALUE() || ctx->QUERY())
 		{
 			size_t startPosition = ctx->start->getStartIndex();
 			rewritten_query_fragment.emplace(std::make_pair(startPosition, std::make_pair("", "bbf_xml")));
@@ -9477,6 +9477,10 @@ validateXMLFunctionArgs(TSqlParser::Xml_func_argContext *xml_func, TSqlParser::E
 	if (xml_func->VALUE() && (expr_list == NULL || expr_list->expression().size() != 2))
 		throw PGErrorWrapperException(ERROR, ERRCODE_UNDEFINED_FUNCTION, "The value function requires 2 argument(s).", getLineAndPos(xml_func));
 
+	/* XML QUERY function requires only 1 argument */
+	if (xml_func->QUERY() && (expr_list == NULL || expr_list->expression().size() != 1))
+		throw PGErrorWrapperException(ERROR, ERRCODE_UNDEFINED_FUNCTION, "The query function requires 1 argument(s).", getLineAndPos(xml_func));
+
 	/* Only String Literal is allowed as agument for XML Functions */
 	if (expr_list)
 	{

contrib/babelfishpg_tsql/src/tsqlUnsupportedFeatureHandler.cpp

@@ -1518,10 +1518,8 @@ antlrcpp::Any TsqlUnsupportedFeatureHandlerImpl::visitId(TSqlParser::IdContext *
 
 antlrcpp::Any TsqlUnsupportedFeatureHandlerImpl::visitXml_func_arg(TSqlParser::Xml_func_argContext *ctx)
 {
-	if (ctx->QUERY())
-		handle(INSTR_UNSUPPORTED_TSQL_XML_QUERY, "XML QUERY", getLineAndPos(ctx));
-	else if (ctx->MODIFY())
-		handle(INSTR_UNSUPPORTED_TSQL_XML_QUERY, "XML MODIFY", getLineAndPos(ctx));
+	if (ctx->MODIFY())
+		handle(INSTR_UNSUPPORTED_TSQL_XML_MODIFY, "XML MODIFY", getLineAndPos(ctx));
 	return visitChildren(ctx);
 }
 

test/JDBC/expected/forxml-path-elements-before-17_10-or-18_4-vu-verify.out

@@ -1116,9 +1116,10 @@ varchar
 -- Using .query() method on XSINIL result
 SELECT (SELECT ID, Name, Department FROM forxml_path_elements_t1 FOR XML PATH('Employee'), ROOT('Data'), ELEMENTS XSINIL, TYPE).query('/Data/Employee[1]') AS query_result;
 GO
-~~ERROR (Code: 33557097)~~
-
-~~ERROR (Message: 'XML QUERY' is not currently supported in Babelfish)~~
+~~START~~
+xml
+<Employee xmlns:xsi="http://www.w3.org/2001/XMLSchema-instance"><ID>1</ID><Name>John</Name><Department>Sales</Department></Employee>
+~~END~~
 
 
 -- Using .exist() method on XSINIL result

test/JDBC/expected/forxml-path-elements-vu-verify.out

@@ -1202,9 +1202,10 @@ varchar
 -- Using .query() method on XSINIL result
 SELECT (SELECT ID, Name, Department FROM forxml_path_elements_t1 FOR XML PATH('Employee'), ROOT('Data'), ELEMENTS XSINIL, TYPE).query('/Data/Employee[1]') AS query_result;
 GO
-~~ERROR (Code: 33557097)~~
-
-~~ERROR (Message: 'XML QUERY' is not currently supported in Babelfish)~~
+~~START~~
+xml
+<Employee xmlns:xsi="http://www.w3.org/2001/XMLSchema-instance"><ID>1</ID><Name>John</Name><Department>Sales</Department></Employee>
+~~END~~
 
 
 -- Using .exist() method on XSINIL result

test/JDBC/expected/xml_exist-before-16_5-vu-verify.out

@@ -823,11 +823,12 @@ GO
 
 -- Exist function called on XML Query
 DECLARE @xml XML = '<artists> <artist name="John Doe"/> <artist name="Edward Poe"/> <artist name="Mark The Great"/> </artists>'
-SELECT @xml.query('/artists/artist').exist('/artist/@name')
+SELECT @xml.query('/artists').exist('/artists/artist/@name')
 GO
-~~ERROR (Code: 33557097)~~
-
-~~ERROR (Message: 'XML QUERY' is not currently supported in Babelfish)~~
+~~START~~
+bit
+1
+~~END~~
 
 
 -- Dependent objects
@@ -1086,9 +1087,10 @@ John Doe
 DECLARE @xml XML = '<artists> <artist name="John Doe"/> <artist name="Edward Poe"/> <artist name="Mark The Great"/> </artists>'
 SELECT @xml.query('/artists/artist')
 GO
-~~ERROR (Code: 33557097)~~
-
-~~ERROR (Message: 'XML QUERY' is not currently supported in Babelfish)~~
+~~START~~
+xml
+<artist name="John Doe"/><artist name="Edward Poe"/><artist name="Mark The Great"/>
+~~END~~
 
 
 DECLARE @xml XML = '<Root><Child1>Value1</Child1><Child2>Value2</Child2></Root>';

test/JDBC/expected/xml_exist-vu-verify.out

@@ -823,11 +823,12 @@ GO
 
 -- Exist function called on XML Query
 DECLARE @xml XML = '<artists> <artist name="John Doe"/> <artist name="Edward Poe"/> <artist name="Mark The Great"/> </artists>'
-SELECT @xml.query('/artists/artist').exist('/artist/@name')
+SELECT @xml.query('/artists').exist('/artists/artist/@name')
 GO
-~~ERROR (Code: 33557097)~~
-
-~~ERROR (Message: 'XML QUERY' is not currently supported in Babelfish)~~
+~~START~~
+bit
+1
+~~END~~
 
 
 -- Dependent objects
@@ -1274,9 +1275,10 @@ John Doe
 DECLARE @xml XML = '<artists> <artist name="John Doe"/> <artist name="Edward Poe"/> <artist name="Mark The Great"/> </artists>'
 SELECT @xml.query('/artists/artist')
 GO
-~~ERROR (Code: 33557097)~~
-
-~~ERROR (Message: 'XML QUERY' is not currently supported in Babelfish)~~
+~~START~~
+xml
+<artist name="John Doe"/><artist name="Edward Poe"/><artist name="Mark The Great"/>
+~~END~~
 
 
 DECLARE @xml XML = '<Root><Child1>Value1</Child1><Child2>Value2</Child2></Root>';