diff --git a/src/include/dbconnector/optimizer/aggregate_optimizer.hpp b/src/include/dbconnector/optimizer/aggregate_optimizer.hpp index df96bc0..447f9eb 100644 --- a/src/include/dbconnector/optimizer/aggregate_optimizer.hpp +++ b/src/include/dbconnector/optimizer/aggregate_optimizer.hpp @@ -5,6 +5,8 @@ #include "duckdb/main/client_context.hpp" #include "duckdb/optimizer/optimizer_extension.hpp" +#include "dbconnector/query/query_writer.hpp" + namespace dbconnector { namespace optimizer { @@ -15,13 +17,13 @@ class AggregateOptimizer { struct Config { bool enabled = false; char identifier_quote = '"'; + query::QuoteEscapeStyle escape_style = query::QuoteEscapeStyle::DOUBLE_QUOTE; std::string table_scan_name; should_push_aggregate_t should_push_aggregate = nullptr; }; - static Config CreateConfig(duckdb::ClientContext &ctx, const std::string &enabled_option, - - char identifier_quote, std::string table_scan_name, + static Config CreateConfig(duckdb::ClientContext &ctx, const std::string &enabled_option, char identifier_quote, + query::QuoteEscapeStyle escape_style, std::string table_scan_name, should_push_aggregate_t should_push_aggregate = nullptr); static void Optimize(const Config &config, duckdb::OptimizerExtensionInput &input, diff --git a/src/include/dbconnector/optimizer/order_by_and_limit_optimizer.hpp b/src/include/dbconnector/optimizer/order_by_and_limit_optimizer.hpp index b963f52..fafeff9 100644 --- a/src/include/dbconnector/optimizer/order_by_and_limit_optimizer.hpp +++ b/src/include/dbconnector/optimizer/order_by_and_limit_optimizer.hpp @@ -5,6 +5,8 @@ #include "duckdb/main/client_context.hpp" #include "duckdb/optimizer/optimizer_extension.hpp" +#include "dbconnector/query/query_writer.hpp" + namespace dbconnector { namespace optimizer { @@ -13,11 +15,12 @@ class OrderByAndLimitOptimizer { struct Config { bool enabled = false; char identifier_quote = '"'; + query::QuoteEscapeStyle escape_style = query::QuoteEscapeStyle::DOUBLE_QUOTE; std::string table_scan_name; }; static Config CreateConfig(duckdb::ClientContext &ctx, const std::string &enabled_option, char identifier_quote, - std::string table_scan_name); + query::QuoteEscapeStyle escape_style, std::string table_scan_name); static void Optimize(const Config &config, duckdb::OptimizerExtensionInput &input, duckdb::unique_ptr &op); diff --git a/src/include/dbconnector/query/query_writer.hpp b/src/include/dbconnector/query/query_writer.hpp index 7c6cf76..f1a3f2e 100644 --- a/src/include/dbconnector/query/query_writer.hpp +++ b/src/include/dbconnector/query/query_writer.hpp @@ -4,19 +4,26 @@ #include "duckdb/common/types/value.hpp" +#include "dbconnector/query/query_writer_exception.hpp" + namespace dbconnector { namespace query { +enum class QuoteEscapeStyle { BACKSLASH, DOUBLE_QUOTE }; + class QueryWriter { public: - static std::string WriteIdentifier(const std::string &identifier, char identifier_quote); - static std::string WriteLiteral(const std::string &identifier); + struct Config { + char quote = '"'; + QuoteEscapeStyle escape_style = QuoteEscapeStyle::DOUBLE_QUOTE; + }; + + static Config CreateConfig(char quote, QuoteEscapeStyle escape_style); + static std::string WriteQuotedAndEscaped(const QueryWriter::Config &config, const std::string &text); static std::string WriteConstant(const duckdb::Value &val); private: static std::string EncodeBlob(const std::string &val); - static std::string EscapeQuotes(const std::string &text, char quote); - static std::string WriteQuoted(const std::string &text, char quote); }; } // namespace query diff --git a/src/include/dbconnector/query/query_writer_exception.hpp b/src/include/dbconnector/query/query_writer_exception.hpp new file mode 100644 index 0000000..a58513e --- /dev/null +++ b/src/include/dbconnector/query/query_writer_exception.hpp @@ -0,0 +1,24 @@ +#pragma once + +#include + +namespace dbconnector { +namespace query { + +class QueryWriterException : public std::exception { +protected: + std::string message; + +public: + QueryWriterException() = default; + + QueryWriterException(const std::string &message) : message(message.data(), message.length()) { + } + + virtual const char *what() const noexcept { + return message.c_str(); + } +}; + +} // namespace query +} // namespace dbconnector diff --git a/src/include/dbconnector/table_scan/filter_pushdown.hpp b/src/include/dbconnector/table_scan/filter_pushdown.hpp index 72c8d9e..1ee8f80 100644 --- a/src/include/dbconnector/table_scan/filter_pushdown.hpp +++ b/src/include/dbconnector/table_scan/filter_pushdown.hpp @@ -5,16 +5,19 @@ #include "duckdb/common/enums/expression_type.hpp" #include "duckdb/planner/table_filter.hpp" +#include "dbconnector/query/query_writer.hpp" + namespace dbconnector { namespace table_scan { class FilterPushdown { struct Config { char identifier_quote = '"'; + query::QuoteEscapeStyle escape_style = query::QuoteEscapeStyle::DOUBLE_QUOTE; }; public: - static Config CreateConfig(char identifier_quote); + static Config CreateConfig(char identifier_quote, query::QuoteEscapeStyle escape_style); static std::string TransformFilter(const Config &config, const std::string &column_name, const duckdb::TableFilter &filter); diff --git a/src/optimizer/aggregate_optimizer.cpp b/src/optimizer/aggregate_optimizer.cpp index 0332bcc..057f8c8 100644 --- a/src/optimizer/aggregate_optimizer.cpp +++ b/src/optimizer/aggregate_optimizer.cpp @@ -26,8 +26,8 @@ namespace optimizer { using namespace duckdb; AggregateOptimizer::Config AggregateOptimizer::CreateConfig(ClientContext &ctx, const std::string &enabled_option, - - char identifier_quote, std::string table_scan_name, + char identifier_quote, query::QuoteEscapeStyle escape_style, + std::string table_scan_name, should_push_aggregate_t should_push_aggregate) { Config res; @@ -38,6 +38,7 @@ AggregateOptimizer::Config AggregateOptimizer::CreateConfig(ClientContext &ctx, } res.identifier_quote = identifier_quote; + res.escape_style = escape_style; res.table_scan_name = std::move(table_scan_name); res.should_push_aggregate = should_push_aggregate; @@ -105,6 +106,7 @@ static PushedAggregate TryPushAggregateToMySQL(const AggregateOptimizer::Config return res; } + auto query_config = query::QueryWriter::CreateConfig(config.identifier_quote, config.escape_style); vector select_fragments; vector group_names; vector new_types; @@ -117,8 +119,7 @@ static PushedAggregate TryPushAggregateToMySQL(const AggregateOptimizer::Config if (!traced_binding.Found()) { return res; } - string quoted = - dbconnector::query::QueryWriter::WriteIdentifier(traced_binding.col_name, config.identifier_quote); + string quoted = query::QueryWriter::WriteQuotedAndEscaped(query_config, traced_binding.col_name); select_fragments.push_back(quoted); group_names.push_back(quoted); new_types.push_back(traced_binding.col_type); @@ -132,8 +133,7 @@ static PushedAggregate TryPushAggregateToMySQL(const AggregateOptimizer::Config string alias = "_agg_" + to_string(agg_idx); if (agg_expr.FunctionMutable().GetName() == "count_star") { - fragment = - "COUNT(*) AS " + dbconnector::query::QueryWriter::WriteIdentifier(alias, config.identifier_quote); + fragment = "COUNT(*) AS " + query::QueryWriter::WriteQuotedAndEscaped(query_config, alias); } else { auto &child_ref = agg_expr.GetChildrenMutable()[0]->Cast(); TracedBindingColumn traced_binding = @@ -141,8 +141,7 @@ static PushedAggregate TryPushAggregateToMySQL(const AggregateOptimizer::Config if (!traced_binding.Found()) { return res; } - string quoted_col = - dbconnector::query::QueryWriter::WriteIdentifier(traced_binding.col_name, config.identifier_quote); + string quoted_col = query::QueryWriter::WriteQuotedAndEscaped(query_config, traced_binding.col_name); string func_upper; if (agg_expr.FunctionMutable().GetName() == "count") { func_upper = "COUNT"; @@ -167,8 +166,7 @@ static PushedAggregate TryPushAggregateToMySQL(const AggregateOptimizer::Config if (func_upper == "AVG") { agg_sql = "CAST(" + agg_sql + " AS DOUBLE)"; } - fragment = - agg_sql + " AS " + dbconnector::query::QueryWriter::WriteIdentifier(alias, config.identifier_quote); + fragment = agg_sql + " AS " + query::QueryWriter::WriteQuotedAndEscaped(query_config, alias); } select_fragments.push_back(fragment); @@ -193,9 +191,8 @@ static PushedAggregate TryPushAggregateToMySQL(const AggregateOptimizer::Config return res; } auto column_name = get.names[table_col_idx]; - auto config = dbconnector::table_scan::FilterPushdown::CreateConfig('`'); - auto new_filter = - dbconnector::table_scan::FilterPushdown::TransformFilter(config, column_name, entry.Filter()); + auto scan_config = table_scan::FilterPushdown::CreateConfig('`', config.escape_style); + auto new_filter = table_scan::FilterPushdown::TransformFilter(scan_config, column_name, entry.Filter()); if (new_filter.empty()) { return res; } diff --git a/src/optimizer/order_by_and_limit_optimizer.cpp b/src/optimizer/order_by_and_limit_optimizer.cpp index 43b6729..3f1096a 100644 --- a/src/optimizer/order_by_and_limit_optimizer.cpp +++ b/src/optimizer/order_by_and_limit_optimizer.cpp @@ -18,10 +18,9 @@ namespace optimizer { using namespace duckdb; -OrderByAndLimitOptimizer::Config OrderByAndLimitOptimizer::CreateConfig(ClientContext &ctx, - const std::string &enabled_option, - char identifier_quote, - std::string table_scan_name) { +OrderByAndLimitOptimizer::Config +OrderByAndLimitOptimizer::CreateConfig(ClientContext &ctx, const std::string &enabled_option, char identifier_quote, + query::QuoteEscapeStyle escape_style, std::string table_scan_name) { Config res; res.enabled = false; @@ -31,6 +30,7 @@ OrderByAndLimitOptimizer::Config OrderByAndLimitOptimizer::CreateConfig(ClientCo } res.identifier_quote = identifier_quote; + res.escape_style = escape_style; res.table_scan_name = std::move(table_scan_name); return res; @@ -84,7 +84,8 @@ static string TraceColumnToGet(const OrderByAndLimitOptimizer::Config &config, E if (actual_col_idx >= get.names.size()) { return std::string(); } - return dbconnector::query::QueryWriter::WriteIdentifier(get.names[actual_col_idx], config.identifier_quote); + auto query_config = query::QueryWriter::CreateConfig(config.identifier_quote, config.escape_style); + return query::QueryWriter::WriteQuotedAndEscaped(query_config, get.names[actual_col_idx]); } static string TryBuildOrderByClause(const OrderByAndLimitOptimizer::Config &config, vector &orders, diff --git a/src/query/query_writer.cpp b/src/query/query_writer.cpp index 55e1937..6bbe9f2 100644 --- a/src/query/query_writer.cpp +++ b/src/query/query_writer.cpp @@ -5,35 +5,42 @@ namespace dbconnector { namespace query { -std::string QueryWriter::EscapeQuotes(const std::string &text, char quote) { +QueryWriter::Config QueryWriter::CreateConfig(char quote, QuoteEscapeStyle escape_style) { + Config res; + res.quote = quote; + res.escape_style = escape_style; + return res; +} + +std::string QueryWriter::WriteQuotedAndEscaped(const QueryWriter::Config &config, const std::string &text) { std::string result; + result.reserve(text.length() + 2); + result.push_back(config.quote); for (auto c : text) { - if (c == quote) { - result += "\\"; - result += quote; - } else if (c == '\\') { - result += "\\\\"; + if (c == config.quote) { + switch (config.escape_style) { + case QuoteEscapeStyle::BACKSLASH: + result.push_back('\\'); + break; + case QuoteEscapeStyle::DOUBLE_QUOTE: + result.push_back(config.quote); + break; + default: + throw QueryWriterException("Invalid unsupported escape style specified: " + + std::to_string(static_cast(config.escape_style))); + } + result.push_back(config.quote); + } else if (c == '\\' && config.escape_style == QuoteEscapeStyle::BACKSLASH) { + result.push_back('\\'); + result.push_back('\\'); } else { - result += c; + result.push_back(c); } } + result.push_back(config.quote); return result; } -std::string QueryWriter::WriteQuoted(const std::string &text, char quote) { - // 1. Escapes all occurences of 'quote' by escaping them with a backslash - // 2. Adds quotes around the string - return std::string(1, quote) + EscapeQuotes(text, quote) + std::string(1, quote); -} - -std::string QueryWriter::WriteIdentifier(const std::string &identifier, char identifier_quote) { - return WriteQuoted(identifier, identifier_quote); -} - -std::string QueryWriter::WriteLiteral(const std::string &identifier) { - return QueryWriter::WriteQuoted(identifier, '\''); -} - std::string QueryWriter::EncodeBlob(const std::string &val) { char const HEX_DIGITS[] = "0123456789ABCDEF"; diff --git a/src/table_scan/filter_pushdown.cpp b/src/table_scan/filter_pushdown.cpp index e18b74d..fd40ed7 100644 --- a/src/table_scan/filter_pushdown.cpp +++ b/src/table_scan/filter_pushdown.cpp @@ -19,6 +19,13 @@ namespace table_scan { using namespace duckdb; +FilterPushdown::Config FilterPushdown::CreateConfig(char identifier_quote, query::QuoteEscapeStyle escape_style) { + Config res; + res.identifier_quote = identifier_quote; + res.escape_style = escape_style; + return res; +} + std::string FilterPushdown::CreateExpression(const std::string &column_name, const vector> &filters, const std::string &op) { vector filter_entries; @@ -150,15 +157,10 @@ std::string FilterPushdown::TransformExpression(const std::string &column_name, } } -FilterPushdown::Config FilterPushdown::CreateConfig(char identifier_quote) { - Config res; - res.identifier_quote = identifier_quote; - return res; -} - std::string FilterPushdown::TransformFilter(const FilterPushdown::Config &config, const std::string &column_name, const TableFilter &filter) { - std::string column_name_quoted = query::QueryWriter::WriteIdentifier(column_name, config.identifier_quote); + auto query_config = query::QueryWriter::CreateConfig(config.identifier_quote, config.escape_style); + std::string column_name_quoted = query::QueryWriter::WriteQuotedAndEscaped(query_config, column_name); auto &expr = FilterUtil::GetExpression(filter, "FilterPushdown::TransformFilter"); return TransformExpression(column_name_quoted, expr); } diff --git a/test/extension/test_query.cpp b/test/extension/test_query.cpp index e0d1a57..11a85f5 100644 --- a/test/extension/test_query.cpp +++ b/test/extension/test_query.cpp @@ -4,7 +4,19 @@ static const std::string group_name = "[query]"; +using namespace dbconnector::query; + TEST_CASE("Test query writer identifier", group_name) { - REQUIRE(dbconnector::query::QueryWriter::WriteIdentifier("foobar", '"') == "\"foobar\""); - REQUIRE(dbconnector::query::QueryWriter::WriteIdentifier("foo\"bar", '"') == "\"foo\\\"bar\""); + { + auto config = QueryWriter::CreateConfig('"', QuoteEscapeStyle::BACKSLASH); + REQUIRE(QueryWriter::WriteQuotedAndEscaped(config, "foobar") == "\"foobar\""); + REQUIRE(QueryWriter::WriteQuotedAndEscaped(config, "foo\"bar\"baz") == "\"foo\\\"bar\\\"baz\""); + REQUIRE(QueryWriter::WriteQuotedAndEscaped(config, "foo\\bar\\baz") == "\"foo\\\\bar\\\\baz\""); + } + { + auto config = QueryWriter::CreateConfig('\'', QuoteEscapeStyle::DOUBLE_QUOTE); + REQUIRE(QueryWriter::WriteQuotedAndEscaped(config, "foobar") == "'foobar'"); + REQUIRE(QueryWriter::WriteQuotedAndEscaped(config, "foo'bar'baz") == "'foo''bar''baz'"); + REQUIRE(QueryWriter::WriteQuotedAndEscaped(config, "foo\\bar\\baz") == "'foo\\bar\\baz'"); + } }