Skip to content

Commit

Permalink
change query_table() to get a varchar or a list of varchars as argument
Browse files Browse the repository at this point in the history
  • Loading branch information
@chrisiou
chrisiou committed Feb 28, 2024
1 parent 2b132b3 commit f5c2617
Show file tree
Hide file tree
Showing 2 changed files with 96 additions and 64 deletions.
78 changes: 48 additions & 30 deletions src/function/table/query_function.cpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -2,56 +2,74 @@
#include "duckdb/common/helper.hpp"
#include "duckdb/main/client_context.hpp"
#include "duckdb/function/table/range.hpp"
#include "duckdb/function/function_set.hpp"
#include "duckdb/parser/tableref/subqueryref.hpp"

namespace duckdb {

static unique_ptr<TableRef> QueryBindReplace(ClientContext &context, TableFunctionBindInput &input) {
string query = input.inputs[0].ToString();
Parser parser(context.GetParserOptions());
static unique_ptr<SubqueryRef> parse_subquery(const string &query, const ParserOptions &options,
const string &error_msg) {
Parser parser(options);
parser.ParseQuery(query);
if (parser.statements.size() != 1 || parser.statements[0]->type != StatementType::SELECT_STATEMENT) {
throw ParserException("Expected a single SELECT statement");
throw ParserException(error_msg);
}

auto select_stmt = unique_ptr_cast<SQLStatement, SelectStatement>(std::move(parser.statements[0]));
auto subquery_ref = duckdb::make_uniq<SubqueryRef>(std::move(select_stmt));
return std::move(subquery_ref);
return duckdb::make_uniq<SubqueryRef>(std::move(select_stmt));
}

static unique_ptr<TableRef> TableBindReplace(ClientContext &context, TableFunctionBindInput &input) {
D_ASSERT(input.named_parameters.size() == 0 || input.named_parameters.size() == 1);
string by_name {""}; // by_name defaults to false
if (input.named_parameters.size() == 1) {
auto it = input.named_parameters.find("by_name");
by_name = (it != input.named_parameters.end() && it->second.GetValue<bool>()) ? "BY NAME " : "";
}
// prepare the query
string union_all_clause = " UNION ALL " + by_name + "FROM ";
string query = "FROM " + input.inputs[0].ToString();
for (size_t i = 1; i < input.inputs.size(); ++i) {
query += union_all_clause + input.inputs[i].ToString();
static void union_tables_query(TableFunctionBindInput &input, string &query) {
auto it = input.named_parameters.find("by_name");
string by_name = (it != input.named_parameters.end() && it->second.GetValue<bool>())
? "BY NAME "
: ""; // 'by_name' variable defaults to false
if (input.inputs[0].type().id() == LogicalTypeId::VARCHAR) {
query += "FROM " + input.inputs[0].ToString();
} else if (input.inputs[0].type() == LogicalType::LIST(LogicalType::VARCHAR)) {
string union_all_clause = " UNION ALL " + by_name + "FROM ";
const auto &children = ListValue::GetChildren(input.inputs[0]);
if (children.empty()) {
throw InvalidInputException("Input list is empty");
}
query += "FROM " + children[0].ToString();
for (size_t i = 1; i < children.size(); ++i) {
auto child = children[i].ToString();
query += union_all_clause + child;
}
}
}

Parser parser(context.GetParserOptions());
parser.ParseQuery(query);
if (parser.statements.size() != 1 || parser.statements[0]->type != StatementType::SELECT_STATEMENT) {
throw ParserException("Expected a table as value");
static unique_ptr<TableRef> QueryBindReplace(ClientContext &context, TableFunctionBindInput &input) {
auto query = input.inputs[0].ToString();
auto subquery_ref = parse_subquery(query, context.GetParserOptions(), "Expected a single SELECT statement");
return subquery_ref;
}

static unique_ptr<TableRef> TableBindReplace(ClientContext &context, TableFunctionBindInput &input) {
D_ASSERT(input.named_parameters.size() == 0 || input.named_parameters.size() == 1); // expected only by_name arg
string err_msg = "Expected a table or a list with tables as input";
if (input.inputs.size() != 1) {
throw ParserException(err_msg);
}
auto select_stmt = unique_ptr_cast<SQLStatement, SelectStatement>(std::move(parser.statements[0]));
auto subquery_ref = duckdb::make_uniq<SubqueryRef>(std::move(select_stmt));
return std::move(subquery_ref);
string query;
union_tables_query(input, query);
auto subquery_ref = parse_subquery(query, context.GetParserOptions(), err_msg);
return subquery_ref;
}

void QueryTableFunction::RegisterFunction(BuiltinFunctions &set) {
TableFunction query("query", {LogicalType::VARCHAR}, nullptr, nullptr);
query.bind_replace = QueryBindReplace;
set.AddFunction(query);

TableFunction query_table("query_table", {LogicalType::VARCHAR}, nullptr, nullptr);
query_table.named_parameters["by_name"] = LogicalType::BOOLEAN;
query_table.bind_replace = TableBindReplace;
query_table.varargs = LogicalType::VARCHAR;
TableFunctionSet query_table("query_table");
TableFunction query_table_function({LogicalType::LIST(LogicalType::VARCHAR)}, nullptr, nullptr);
query_table_function.named_parameters["by_name"] = LogicalType::BOOLEAN;
query_table_function.bind_replace = TableBindReplace;
query_table.AddFunction(query_table_function);

query_table_function.arguments = {LogicalType::VARCHAR};
query_table.AddFunction(query_table_function);
set.AddFunction(query_table);
}

Expand Down
82 changes: 48 additions & 34 deletions test/sql/catalog/function/query_function.test
View file
Original file line number Diff line number Diff line change
Expand Up @@ -25,11 +25,6 @@ FROM query('SELECT abs(-42)');
----
42

query I
FROM query('FROM query("SELECT 1")');
----
1

query I
SELECT * FROM query('SELECT * FROM (SELECT 1 + 2)');
----
Expand All @@ -40,11 +35,6 @@ FROM query('SELECT 1, 2, 3');
----
1 2 3

query I
SELECT * FROM query('FROM query(''SELECT 17 + 25'')');
----
42

query I
FROM query('SELECT 42;;;--- hello;');
----
Expand All @@ -63,10 +53,10 @@ hello

# query from table
statement ok
CREATE TABLE tbl1 (a INT, b INT);
CREATE TABLE tbl (a INT, b INT);

statement ok
FROM query('SELECT *, 1 + 2 FROM tbl1');
FROM query('SELECT *, 1 + 2 FROM tbl');

statement ok
CREATE TABLE tbl2 (a INT, b INT, c INT);
Expand Down Expand Up @@ -134,61 +124,75 @@ Parser Error: Expected a single SELECT statement


# test query_table()
query III
from query_table(tbl2);
----
1 2 3
4 5 6
7 8 9

statement ok
create table t as select 42;
create table t_int as select 42;

statement ok
create table t2 as select 'duckdb';
create table t_varchar as select 'duckdb';

statement ok
create table t3 as select '';
create table t_empty as select '';

statement ok
create table t4 as select '1?ch@racter$';
create table t2_varchar as select '1?ch@racter$';

query I
from query_table('t');
from query_table('t_int');
----
42

query III
from query_table(tbl2);
----
1 2 3
4 5 6
7 8 9

query I
from query_table('t','t2','t4');
from query_table(['t_int']);
----
42
duckdb
1?ch@racter$

query I
from query_table(t, t2);
from query_table(['t_int', t2_varchar]);
----
42
duckdb
1?ch@racter$

query I
from query_table('(select 17 + 25)');
----
42

query I
from query_table('t', 't2', 't3', t4, '(select ''I am a subquery'')');
from query_table(['t_int', 't_varchar', 't_empty', t2_varchar, '(select ''I am a subquery'')']);
----
42
duckdb
(empty)
1?ch@racter$
I am a subquery

# test incorrect usage
statement error
from query_table([]);
----
Binder Error: No function matches the given name and argument types 'query_table(INTEGER[])'.

statement error
from query_table(['']);
----
Parser Error: syntax error at end of input

statement error
from query_table('t_int', 't_varchar', t2_varchar);
----
Binder Error: No function matches the given name and argument types 'query_table(VARCHAR, VARCHAR, VARCHAR)'.

statement error
from query_table(['t', 't2', 't3', t4, '(select ''I am a subquery'')']);
from query_table([t_int, tbl2]);
----
Binder Error: No function matches the given name and argument types 'query_table(VARCHAR[])'.
Binder Error: Set operations can only apply to expressions with the same number of result columns

statement error
from query_table(not_defined_tbl);
Expand All @@ -213,16 +217,26 @@ from query_table(tbl2, by_name=true);
4 5 6
7 8 9

query I
from query_table(['t_int', 't_varchar', 't_empty', t2_varchar, '(select ''I am a subquery'')'], by_name=false);
----
42
duckdb
(empty)
1?ch@racter$
I am a subquery

query IIIII
from query_table('t', 't2', 't3', t4, '(select ''I am a subquery'')', by_name=true);
from query_table(['t_int', 't_varchar', 't_empty', t2_varchar, '(select ''I am a subquery'')'], by_name=true);
----
42 NULL NULL NULL NULL
NULL duckdb NULL NULL NULL
NULL NULL (empty) NULL NULL
NULL NULL NULL 1?ch@racter$ NULL
NULL NULL NULL NULL I am a subquery

# test incorrect usage
statement error
from query_table('t', not_valid=true);
from query_table('t_int', not_valid=true);
----
Binder Error: Invalid named parameter "not_valid" for function query_table

0 comments on commit f5c2617

Please sign in to comment.