Add scalar form of parse_tables

Author: zacMode

src/include/parse_tables.hpp

@@ -27,5 +27,6 @@ struct TableRefResult {
 void ExtractTablesFromSQL(const std::string &sql, std::vector<TableRefResult> &results);
 
 void RegisterParseTablesFunction(duckdb::DatabaseInstance &db);
+void RegisterParseTableScalarFunction(DatabaseInstance &db);
 
 } // namespace duckdb

src/parse_tables.cpp

@@ -7,6 +7,8 @@
 #include "duckdb/parser/tableref/joinref.hpp"
 #include "duckdb/parser/tableref/subqueryref.hpp"
 #include "duckdb/main/extension_util.hpp"
+#include "duckdb/function/scalar/nested_functions.hpp"
+
 
 namespace duckdb {
 
@@ -140,32 +142,28 @@ static void ExtractTablesFromQueryNode(
     }
 }
 
+void ExtractTablesFromSQL(const std::string &sql, std::vector<TableRefResult> &results) {
+    Parser parser;
+    parser.ParseQuery(sql);
+
+    for (auto &stmt : parser.statements) {
+        if (stmt->type == StatementType::SELECT_STATEMENT) {
+            auto &select_stmt = (SelectStatement &)*stmt;
+            if (select_stmt.node) {
+                ExtractTablesFromQueryNode(*select_stmt.node, results);
+            }
+        }
+    }
+}
+
 static void ParseTablesFunction(ClientContext &context,
                    TableFunctionInput &data,
                    DataChunk &output) {
     auto &state = (ParseTablesState &)*data.global_state;
     auto &bind_data = (ParseTablesBindData &)*data.bind_data;
 
     if (state.results.empty() && state.row == 0) {
-        try {
-            Parser parser;
-            parser.ParseQuery(bind_data.sql);
-
-            for (auto &stmt : parser.statements) {
-                if (stmt->type != StatementType::SELECT_STATEMENT) {
-                    throw InvalidInputException("parse_tables only supports SELECT statements");
-                }
-                
-                if (stmt->type == StatementType::SELECT_STATEMENT) {
-                    auto &select_stmt = (SelectStatement &)*stmt;
-                    if (select_stmt.node) {
-                        ExtractTablesFromQueryNode(*select_stmt.node, state.results);
-                    }
-                }
-            }
-        } catch (const std::exception &ex) {
-            throw InvalidInputException("Failed to parse SQL: %s", ex.what());
-        }
+        ExtractTablesFromSQL(bind_data.sql, state.results);
     }
 
     if (state.row >= state.results.size()) {
@@ -181,6 +179,41 @@ static void ParseTablesFunction(ClientContext &context,
     state.row++;
 }
 
+static void ParseTablesScalarFunction(DataChunk &args, ExpressionState &state, Vector &result) {
+    // Execute does the heavy lifting of iterating over the input data
+    // and calling the provided lambda function for each input value.
+    // The lambda function is responsible for parsing the SQL query and
+    // extracting the table names.
+    UnaryExecutor::Execute<string_t, list_entry_t>(args.data[0], result, args.size(), 
+    [&result](string_t query) -> list_entry_t {
+        // Parse the SQL query and extract table names
+        auto query_string = query.GetString();
+        std::vector<TableRefResult> parsed_tables;
+        ExtractTablesFromSQL(query_string, parsed_tables);
+    
+        auto current_size = ListVector::GetListSize(result);
+        auto number_of_tables = parsed_tables.size();
+        auto new_size = current_size + number_of_tables;
+
+        // grow list if needed
+        if (ListVector::GetListCapacity(result) < new_size) {
+            ListVector::Reserve(result, new_size);
+        }
+
+        // Write the string into the child vector
+        auto tables = FlatVector::GetData<string_t>(ListVector::GetEntry(result));
+        for (size_t i = 0; i < parsed_tables.size(); i++) {
+            auto &table = parsed_tables[i];
+            tables[current_size + i] = StringVector::AddStringOrBlob(ListVector::GetEntry(result), table.table);
+        }
+
+        // Update size
+        ListVector::SetListSize(result, new_size);
+
+        return list_entry_t(current_size, number_of_tables); 
+    });
+}
+
 // Extension scaffolding
 // ---------------------------------------------------
 
@@ -189,4 +222,9 @@ void RegisterParseTablesFunction(DatabaseInstance &db) {
     ExtensionUtil::RegisterFunction(db, tf);
 }
 
+void RegisterParseTableScalarFunction(DatabaseInstance &db) {
+    ScalarFunction sf( "parse_tables", {LogicalType::VARCHAR}, LogicalType::LIST(LogicalType::VARCHAR), ParseTablesScalarFunction);
+    ExtensionUtil::RegisterFunction(db, sf);
+}
+
 } // namespace duckdb

src/parser_tools_extension.cpp

@@ -22,6 +22,7 @@ namespace duckdb {
 
 static void LoadInternal(DatabaseInstance &instance) {
     RegisterParseTablesFunction(instance);
+	RegisterParseTableScalarFunction(instance);
 }
 
 void ParserToolsExtension::Load(DuckDB &db) {