66#include " duckdb/parser/statement/extension_statement.hpp"
77#include " duckdb/function/table_function.hpp"
88#include " duckdb/main/connection.hpp"
9+ #include " duckdb/logging/logger.hpp"
910
1011// Include FFI types header
1112#include " yardstick_ffi.h"
@@ -222,6 +223,8 @@ static std::string RewritePercentileWithinGroup(const std::string &sql) {
222223// TABLE FUNCTION: yardstick(sql) - Execute SQL with AGGREGATE() expansion
223224// =============================================================================
224225
226+ static void HandleAggregateWarnings (ClientContext &context, const string &warnings);
227+
225228struct YardstickQueryData : public TableFunctionData {
226229 string original_sql;
227230 string rewritten_sql;
@@ -235,6 +238,9 @@ static unique_ptr<FunctionData> YardstickQueryBind(ClientContext &context,
235238 vector<string> &names) {
236239 auto data = make_uniq<YardstickQueryData>();
237240 data->original_sql = input.inputs [0 ].GetValue <string>();
241+ if (input.inputs .size () > 1 ) {
242+ HandleAggregateWarnings (context, input.inputs [1 ].GetValue <string>());
243+ }
238244
239245 // Rewrite the SQL using our Rust code
240246 if (yardstick_has_aggregate (data->original_sql .c_str ())) {
@@ -330,6 +336,49 @@ static bool StartsWithSemantic(const std::string &query, std::string &stripped_q
330336 return true ;
331337}
332338
339+ static std::string EscapeSqlStringLiteral (const std::string &value) {
340+ string escaped;
341+ for (char c : value) {
342+ if (c == ' \' ' ) {
343+ escaped += " ''" ;
344+ } else {
345+ escaped += c;
346+ }
347+ }
348+ return escaped;
349+ }
350+
351+ static std::string AggregateWarnings (YardstickAggregateResult &result) {
352+ return result.warnings ? string (result.warnings ) : string ();
353+ }
354+
355+ static bool WarningsAsErrors (ClientContext &context) {
356+ Value setting;
357+ if (!context.TryGetCurrentSetting (" warnings_as_errors" , setting) || setting.IsNull ()) {
358+ return false ;
359+ }
360+ return setting.GetValue <bool >();
361+ }
362+
363+ static void HandleAggregateWarnings (ClientContext &context, const string &warnings) {
364+ if (warnings.empty ()) {
365+ return ;
366+ }
367+ if (WarningsAsErrors (context)) {
368+ throw InvalidInputException (" %s" , warnings);
369+ }
370+ DUCKDB_LOG_WARNING (context, " %s" , warnings.c_str ());
371+ }
372+
373+ static std::string YardstickWrapperSql (const std::string &expanded_sql, const std::string &warnings) {
374+ string wrapper_sql = " SELECT * FROM yardstick('" + EscapeSqlStringLiteral (expanded_sql) + " '" ;
375+ if (!warnings.empty ()) {
376+ wrapper_sql += " , '" + EscapeSqlStringLiteral (warnings) + " '" ;
377+ }
378+ wrapper_sql += " )" ;
379+ return wrapper_sql;
380+ }
381+
333382ParserExtensionParseResult yardstick_parse (ParserExtensionInfo *,
334383 const std::string &query) {
335384
@@ -357,20 +406,10 @@ ParserExtensionParseResult yardstick_parse(ParserExtensionInfo *,
357406
358407 if (result.had_aggregate ) {
359408 string expanded_sql (result.expanded_sql );
409+ string warnings = AggregateWarnings (result);
360410 yardstick_free_aggregate_result (result);
361411
362- // Escape single quotes for embedding in string literal
363- string escaped_sql;
364- for (char c : expanded_sql) {
365- if (c == ' \' ' ) {
366- escaped_sql += " ''" ;
367- } else {
368- escaped_sql += c;
369- }
370- }
371-
372- // Wrap in table function call
373- string wrapper_sql = " SELECT * FROM yardstick('" + escaped_sql + " ')" ;
412+ string wrapper_sql = YardstickWrapperSql (expanded_sql, warnings);
374413
375414 Parser parser;
376415 parser.ParseQuery (wrapper_sql);
@@ -458,6 +497,7 @@ ParserOverrideResult yardstick_parser_override(ParserExtensionInfo *,
458497
459498 if (result.had_aggregate ) {
460499 string expanded_sql (result.expanded_sql );
500+ string warnings = AggregateWarnings (result);
461501 yardstick_free_aggregate_result (result);
462502
463503 // Validate the expanded SQL parses. If expansion produced garbage
@@ -478,16 +518,7 @@ ParserOverrideResult yardstick_parser_override(ParserExtensionInfo *,
478518 validation_parser.statements [0 ]->type == StatementType::SELECT_STATEMENT ;
479519
480520 if (is_select) {
481- string escaped_sql;
482- for (char c : expanded_sql) {
483- if (c == ' \' ' ) {
484- escaped_sql += " ''" ;
485- } else {
486- escaped_sql += c;
487- }
488- }
489-
490- string wrapper_sql = " SELECT * FROM yardstick('" + escaped_sql + " ')" ;
521+ string wrapper_sql = YardstickWrapperSql (expanded_sql, warnings);
491522 Parser parser;
492523 parser.ParseQuery (wrapper_sql);
493524 return ParserOverrideResult (std::move (parser.statements ));
@@ -583,20 +614,11 @@ BoundStatement yardstick_bind(ClientContext &context, Binder &binder,
583614
584615 if (result.had_aggregate ) {
585616 string expanded_sql (result.expanded_sql );
617+ string warnings = AggregateWarnings (result);
586618 yardstick_free_aggregate_result (result);
587619
588- // Escape single quotes for embedding in string literal
589- string escaped_sql;
590- for (char c : expanded_sql) {
591- if (c == ' \' ' ) {
592- escaped_sql += " ''" ;
593- } else {
594- escaped_sql += c;
595- }
596- }
597-
598620 // Rebind through table function so rewritten SQL executes with normal planning
599- string wrapper_sql = " SELECT * FROM yardstick(' " + escaped_sql + " ') " ;
621+ string wrapper_sql = YardstickWrapperSql (expanded_sql, warnings) ;
600622 Parser parser;
601623 parser.ParseQuery (wrapper_sql);
602624 auto statements = std::move (parser.statements );
@@ -667,6 +689,9 @@ static void LoadInternal(ExtensionLoader &loader) {
667689 TableFunction query_func (" yardstick" , {LogicalType::VARCHAR },
668690 YardstickQueryFunction, YardstickQueryBind);
669691 loader.RegisterFunction (query_func);
692+ TableFunction query_func_with_warnings (" yardstick" , {LogicalType::VARCHAR , LogicalType::VARCHAR },
693+ YardstickQueryFunction, YardstickQueryBind);
694+ loader.RegisterFunction (query_func_with_warnings);
670695}
671696
672697void YardstickExtension::Load (ExtensionLoader &loader) {
0 commit comments