Support different query quoting styles

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,

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<duckdb::LogicalOperator> &op);

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

src/include/dbconnector/query/query_writer_exception.hpp

@@ -0,0 +1,24 @@
+#pragma once
+
+#include <stdexcept>
+
+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

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);

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<string> select_fragments;
 	vector<string> group_names;
 	vector<LogicalType> 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,17 +133,15 @@ 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<BoundColumnRefExpression>();
 			TracedBindingColumn traced_binding =
 			    OptimizerUtil::TraceBindingToColumn(child_ref.BindingMutable(), aggr_child, get);
 			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;
 			}

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<BoundOrderByNode> &orders,

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<int>(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";
 

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<unique_ptr<Expression>> &filters, const std::string &op) {
 	vector<std::string> 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);
 }

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'");
+	}
 }