Merge pull request #53250 from evillique/postgresql-single-quotes

Fix character escaping in the PostgreSQL engine

src/IO/WriteHelpers.h

@@ -304,9 +304,10 @@ inline void writeJSONString(const char * begin, const char * end, WriteBuffer &
 /** Will escape quote_character and a list of special characters('\b', '\f', '\n', '\r', '\t', '\0', '\\').
  *   - when escape_quote_with_quote is true, use backslash to escape list of special characters,
  *      and use quote_character to escape quote_character. such as: 'hello''world'
- *   - otherwise use backslash to escape list of special characters and quote_character
+ *     otherwise use backslash to escape list of special characters and quote_character
+ *   - when escape_backslash_with_backslash is true, backslash is escaped with another backslash
  */
-template <char quote_character, bool escape_quote_with_quote = false>
+template <char quote_character, bool escape_quote_with_quote = false, bool escape_backslash_with_backslash = true>
 void writeAnyEscapedString(const char * begin, const char * end, WriteBuffer & buf)
 {
     const char * pos = begin;
@@ -360,7 +361,8 @@ void writeAnyEscapedString(const char * begin, const char * end, WriteBuffer & b
                     writeChar('0', buf);
                     break;
                 case '\\':
-                    writeChar('\\', buf);
+                    if constexpr (escape_backslash_with_backslash)
+                        writeChar('\\', buf);
                     writeChar('\\', buf);
                     break;
                 default:
@@ -466,6 +468,13 @@ inline void writeQuotedString(std::string_view ref, WriteBuffer & buf)
     writeAnyQuotedString<'\''>(ref.data(), ref.data() + ref.size(), buf);
 }
 
+inline void writeQuotedStringPostgreSQL(std::string_view ref, WriteBuffer & buf)
+{
+    writeChar('\'', buf);
+    writeAnyEscapedString<'\'', true, false>(ref.data(), ref.data() + ref.size(), buf);
+    writeChar('\'', buf);
+}
+
 inline void writeDoubleQuotedString(const String & s, WriteBuffer & buf)
 {
     writeAnyQuotedString<'"'>(s, buf);

src/Parsers/ASTLiteral.cpp

@@ -93,7 +93,7 @@ void ASTLiteral::appendColumnNameImpl(WriteBuffer & ostr) const
 
 void ASTLiteral::appendColumnNameImplLegacy(WriteBuffer & ostr) const
 {
-     /// 100 - just arbitrary value.
+    /// 100 - just arbitrary value.
     constexpr auto min_elements_for_hashing = 100;
 
     /// Special case for very large arrays. Instead of listing all elements, will use hash of them.
@@ -118,9 +118,31 @@ void ASTLiteral::appendColumnNameImplLegacy(WriteBuffer & ostr) const
     }
 }
 
+/// Use different rules for escaping backslashes and quotes
+class FieldVisitorToStringPostgreSQL : public StaticVisitor<String>
+{
+public:
+    template<typename T>
+    String operator() (const T & x) const { return visitor(x); }
+
+private:
+    FieldVisitorToString visitor;
+};
+
+template<>
+String FieldVisitorToStringPostgreSQL::operator() (const String & x) const
+{
+    WriteBufferFromOwnString wb;
+    writeQuotedStringPostgreSQL(x, wb);
+    return wb.str();
+}
+
 void ASTLiteral::formatImplWithoutAlias(const FormatSettings & settings, IAST::FormatState &, IAST::FormatStateStacked) const
 {
-    settings.ostr << applyVisitor(FieldVisitorToString(), value);
+    if (settings.literal_escaping_style == LiteralEscapingStyle::Regular)
+        settings.ostr << applyVisitor(FieldVisitorToString(), value);
+    else
+        settings.ostr << applyVisitor(FieldVisitorToStringPostgreSQL(), value);
 }
 
 }

src/Parsers/IAST.h

@@ -3,6 +3,7 @@
 #include <base/types.h>
 #include <Parsers/IAST_fwd.h>
 #include <Parsers/IdentifierQuotingStyle.h>
+#include <Parsers/LiteralEscapingStyle.h>
 #include <Common/Exception.h>
 #include <Common/TypePromotion.h>
 #include <IO/WriteBufferFromString.h>
@@ -197,6 +198,7 @@ public:
         IdentifierQuotingStyle identifier_quoting_style;
         bool show_secrets; /// Show secret parts of the AST (e.g. passwords, encryption keys).
         char nl_or_ws; /// Newline or whitespace.
+        LiteralEscapingStyle literal_escaping_style;
 
         explicit FormatSettings(
             WriteBuffer & ostr_,
@@ -204,7 +206,8 @@ public:
             bool hilite_ = false,
             bool always_quote_identifiers_ = false,
             IdentifierQuotingStyle identifier_quoting_style_ = IdentifierQuotingStyle::Backticks,
-            bool show_secrets_ = true)
+            bool show_secrets_ = true,
+            LiteralEscapingStyle literal_escaping_style_ = LiteralEscapingStyle::Regular)
             : ostr(ostr_)
             , one_line(one_line_)
             , hilite(hilite_)
@@ -212,6 +215,7 @@ public:
             , identifier_quoting_style(identifier_quoting_style_)
             , show_secrets(show_secrets_)
             , nl_or_ws(one_line ? ' ' : '\n')
+            , literal_escaping_style(literal_escaping_style_)
         {
         }
 
@@ -223,6 +227,7 @@ public:
             , identifier_quoting_style(other.identifier_quoting_style)
             , show_secrets(other.show_secrets)
             , nl_or_ws(other.nl_or_ws)
+            , literal_escaping_style(other.literal_escaping_style)
         {
         }
 

src/Parsers/LiteralEscapingStyle.h

@@ -0,0 +1,14 @@
+#pragma once
+
+
+namespace DB
+{
+
+/// Method to escape single quotes.
+enum class LiteralEscapingStyle
+{
+    Regular,         /// Escape backslashes with backslash (\\) and quotes with backslash (\')
+    PostgreSQL,      /// Do not escape backslashes (\), escape quotes with quote ('')
+};
+
+}

src/Storages/StorageMySQL.cpp

@@ -104,6 +104,7 @@ Pipe StorageMySQL::read(
         column_names_,
         storage_snapshot->metadata->getColumns().getOrdinary(),
         IdentifierQuotingStyle::BackticksMySQL,
+        LiteralEscapingStyle::Regular,
         remote_database_name,
         remote_table_name,
         context_);

src/Storages/StoragePostgreSQL.cpp

@@ -122,7 +122,7 @@ Pipe StoragePostgreSQL::read(
         query_info_,
         column_names_,
         storage_snapshot->metadata->getColumns().getOrdinary(),
-        IdentifierQuotingStyle::DoubleQuotes, remote_table_schema, remote_table_name, context_);
+        IdentifierQuotingStyle::DoubleQuotes, LiteralEscapingStyle::PostgreSQL, remote_table_schema, remote_table_name, context_);
     LOG_TRACE(log, "Query: {}", query);
 
     Block sample_block;

src/Storages/StorageSQLite.cpp

@@ -91,6 +91,7 @@ Pipe StorageSQLite::read(
         column_names,
         storage_snapshot->metadata->getColumns().getOrdinary(),
         IdentifierQuotingStyle::DoubleQuotes,
+        LiteralEscapingStyle::Regular,
         "",
         remote_table_name,
         context_);

src/Storages/StorageXDBC.cpp

@@ -79,6 +79,7 @@ std::function<void(std::ostream &)> StorageXDBC::getReadPOSTDataCallback(
         column_names,
         columns_description.getOrdinary(),
         bridge_helper->getIdentifierQuotingStyle(),
+        LiteralEscapingStyle::Regular,
         remote_database_name,
         remote_table_name,
         local_context);

src/Storages/tests/gtest_transform_query_for_external_database.cpp

@@ -127,7 +127,8 @@ static void checkOld(
     std::string transformed_query = transformQueryForExternalDatabase(
         query_info,
         query_info.syntax_analyzer_result->requiredSourceColumns(),
-        state.getColumns(0), IdentifierQuotingStyle::DoubleQuotes, "test", "table", state.context);
+        state.getColumns(0), IdentifierQuotingStyle::DoubleQuotes,
+        LiteralEscapingStyle::Regular, "test", "table", state.context);
 
     EXPECT_EQ(transformed_query, expected) << query;
 }
@@ -180,7 +181,8 @@ static void checkNewAnalyzer(
     query_info.table_expression = findTableExpression(query_node->getJoinTree(), "table");
 
     std::string transformed_query = transformQueryForExternalDatabase(
-        query_info, column_names, state.getColumns(0), IdentifierQuotingStyle::DoubleQuotes, "test", "table", state.context);
+        query_info, column_names, state.getColumns(0), IdentifierQuotingStyle::DoubleQuotes,
+        LiteralEscapingStyle::Regular, "test", "table", state.context);
 
     EXPECT_EQ(transformed_query, expected) << query;
 }

src/Storages/transformQueryForExternalDatabase.cpp

@@ -258,6 +258,7 @@ String transformQueryForExternalDatabaseImpl(
     Names used_columns,
     const NamesAndTypesList & available_columns,
     IdentifierQuotingStyle identifier_quoting_style,
+    LiteralEscapingStyle literal_escaping_style,
     const String & database,
     const String & table,
     ContextPtr context)
@@ -337,7 +338,8 @@ String transformQueryForExternalDatabaseImpl(
     IAST::FormatSettings settings(
             out, /*one_line*/ true, /*hilite*/ false,
             /*always_quote_identifiers*/ identifier_quoting_style != IdentifierQuotingStyle::None,
-            /*identifier_quoting_style*/ identifier_quoting_style);
+            /*identifier_quoting_style*/ identifier_quoting_style, /*show_secrets_*/ true,
+            /*literal_escaping_style*/ literal_escaping_style);
 
     select->format(settings);
 
@@ -351,6 +353,7 @@ String transformQueryForExternalDatabase(
     const Names & column_names,
     const NamesAndTypesList & available_columns,
     IdentifierQuotingStyle identifier_quoting_style,
+    LiteralEscapingStyle literal_escaping_style,
     const String & database,
     const String & table,
     ContextPtr context)
@@ -375,6 +378,7 @@ String transformQueryForExternalDatabase(
             column_names,
             available_columns,
             identifier_quoting_style,
+            literal_escaping_style,
             database,
             table,
             context);
@@ -386,6 +390,7 @@ String transformQueryForExternalDatabase(
         query_info.syntax_analyzer_result->requiredSourceColumns(),
         available_columns,
         identifier_quoting_style,
+        literal_escaping_style,
         database,
         table,
         context);

src/Storages/transformQueryForExternalDatabase.h

@@ -31,6 +31,7 @@ String transformQueryForExternalDatabase(
     const Names & column_names,
     const NamesAndTypesList & available_columns,
     IdentifierQuotingStyle identifier_quoting_style,
+    LiteralEscapingStyle literal_escaping_style,
     const String & database,
     const String & table,
     ContextPtr context);

tests/integration/test_storage_postgresql/test.py

@@ -726,6 +726,22 @@ def test_auto_close_connection(started_cluster):
     assert count == 2
 
 
+def test_literal_escaping(started_cluster):
+    cursor = started_cluster.postgres_conn.cursor()
+    cursor.execute(f"DROP TABLE IF EXISTS escaping")
+    cursor.execute(f"CREATE TABLE escaping(text varchar(255))")
+    node1.query(
+        "CREATE TABLE default.escaping (text String) ENGINE = PostgreSQL('postgres1:5432', 'postgres', 'escaping', 'postgres', 'mysecretpassword')"
+    )
+    node1.query("SELECT * FROM escaping WHERE text = ''''")  # ' -> ''
+    node1.query("SELECT * FROM escaping WHERE text = '\\''")  # ' -> ''
+    node1.query("SELECT * FROM escaping WHERE text = '\\\\\\''")  # \' -> \''
+    node1.query("SELECT * FROM escaping WHERE text = '\\\\\\''")  # \' -> \''
+    node1.query("SELECT * FROM escaping WHERE text like '%a''a%'")  # %a'a% -> %a''a%
+    node1.query("SELECT * FROM escaping WHERE text like '%a\\'a%'")  # %a'a% -> %a''a%
+    cursor.execute(f"DROP TABLE escaping")
+
+
 if __name__ == "__main__":
     cluster.start()
     input("Cluster created, press any key to destroy...")