Merge pull request #19011 from ClickHouse/fix_operators_escaping

Fix manipulators for common string types

src/IO/Operators.h

@@ -61,7 +61,11 @@ template <typename T> WriteBuffer & operator<< (QuoteManipWriteBuffer buf,
 template <typename T> WriteBuffer & operator<< (DoubleQuoteManipWriteBuffer buf,   const T & x) { writeDoubleQuoted(x, buf.get()); return buf; }
 template <typename T> WriteBuffer & operator<< (BinaryManipWriteBuffer buf,        const T & x) { writeBinary(x, buf.get());       return buf; }
 
-inline WriteBuffer & operator<< (EscapeManipWriteBuffer buf,      const char * x) { writeAnyEscapedString<'\''>(x, x + strlen(x), buf.get()); return buf; }
+inline  WriteBuffer & operator<< (EscapeManipWriteBuffer buf, const String & x)           { writeEscapedString(x, buf); return buf; }
+inline  WriteBuffer & operator<< (EscapeManipWriteBuffer buf, const std::string_view & x) { writeEscapedString(x, buf); return buf; }
+inline  WriteBuffer & operator<< (EscapeManipWriteBuffer buf, const StringRef & x)        { writeEscapedString(x, buf); return buf; }
+inline  WriteBuffer & operator<< (EscapeManipWriteBuffer buf, const char * x)             { writeEscapedString(x, strlen(x), buf); return buf; }
+
 inline WriteBuffer & operator<< (QuoteManipWriteBuffer buf,       const char * x) { writeAnyQuotedString<'\''>(x, x + strlen(x), buf.get()); return buf; }
 inline WriteBuffer & operator<< (DoubleQuoteManipWriteBuffer buf, const char * x) { writeAnyQuotedString<'"'>(x, x + strlen(x), buf.get()); return buf; }
 inline WriteBuffer & operator<< (BinaryManipWriteBuffer buf,      const char * x) { writeStringBinary(x, buf.get()); return buf; }

src/IO/WriteHelpers.h

@@ -483,6 +483,10 @@ inline void writeEscapedString(const StringRef & ref, WriteBuffer & buf)
     writeEscapedString(ref.data, ref.size, buf);
 }
 
+inline void writeEscapedString(const std::string_view & ref, WriteBuffer & buf)
+{
+    writeEscapedString(ref.data(), ref.size(), buf);
+}
 
 template <char quote_character>
 void writeAnyQuotedString(const char * begin, const char * end, WriteBuffer & buf)
@@ -512,17 +516,31 @@ inline void writeQuotedString(const String & s, WriteBuffer & buf)
     writeAnyQuotedString<'\''>(s, buf);
 }
 
-
 inline void writeQuotedString(const StringRef & ref, WriteBuffer & buf)
 {
     writeAnyQuotedString<'\''>(ref, buf);
 }
 
+inline void writeQuotedString(const std::string_view & ref, WriteBuffer & buf)
+{
+    writeAnyQuotedString<'\''>(ref.data(), ref.data() + ref.size(), buf);
+}
+
+inline void writeDoubleQuotedString(const String & s, WriteBuffer & buf)
+{
+    writeAnyQuotedString<'"'>(s, buf);
+}
+
 inline void writeDoubleQuotedString(const StringRef & s, WriteBuffer & buf)
 {
     writeAnyQuotedString<'"'>(s, buf);
 }
 
+inline void writeDoubleQuotedString(const std::string_view & s, WriteBuffer & buf)
+{
+    writeAnyQuotedString<'"'>(s.data(), s.data() + s.size(), buf);
+}
+
 /// Outputs a string in backquotes.
 inline void writeBackQuotedString(const StringRef & s, WriteBuffer & buf)
 {
@@ -901,6 +919,7 @@ writeBinary(const T & x, WriteBuffer & buf) { writePODBinary(x, buf); }
 
 inline void writeBinary(const String & x, WriteBuffer & buf) { writeStringBinary(x, buf); }
 inline void writeBinary(const StringRef & x, WriteBuffer & buf) { writeStringBinary(x, buf); }
+inline void writeBinary(const std::string_view & x, WriteBuffer & buf) { writeStringBinary(x, buf); }
 inline void writeBinary(const Int128 & x, WriteBuffer & buf) { writePODBinary(x, buf); }
 inline void writeBinary(const UInt128 & x, WriteBuffer & buf) { writePODBinary(x, buf); }
 inline void writeBinary(const DummyUInt256 & x, WriteBuffer & buf) { writePODBinary(x, buf); }
@@ -1001,6 +1020,10 @@ writeQuoted(const T & x, WriteBuffer & buf) { writeText(x, buf); }
 
 inline void writeQuoted(const String & x, WriteBuffer & buf) { writeQuotedString(x, buf); }
 
+inline void writeQuoted(const std::string_view & x, WriteBuffer & buf) { writeQuotedString(x, buf); }
+
+inline void writeQuoted(const StringRef & x, WriteBuffer & buf) { writeQuotedString(x, buf); }
+
 inline void writeQuoted(const LocalDate & x, WriteBuffer & buf)
 {
     writeChar('\'', buf);
@@ -1043,6 +1066,10 @@ writeDoubleQuoted(const T & x, WriteBuffer & buf) { writeText(x, buf); }
 
 inline void writeDoubleQuoted(const String & x, WriteBuffer & buf) { writeDoubleQuotedString(x, buf); }
 
+inline void writeDoubleQuoted(const std::string_view & x, WriteBuffer & buf) { writeDoubleQuotedString(x, buf); }
+
+inline void writeDoubleQuoted(const StringRef & x, WriteBuffer & buf) { writeDoubleQuotedString(x, buf); }
+
 inline void writeDoubleQuoted(const LocalDate & x, WriteBuffer & buf)
 {
     writeChar('"', buf);

src/IO/tests/gtest_manip.cpp

@@ -0,0 +1,82 @@
+#include <gtest/gtest.h>
+
+#include <string>
+#include <type_traits>
+#include <common/StringRef.h>
+#include <IO/Operators.h>
+#include <IO/WriteHelpers.h>
+#include <IO/WriteBufferFromString.h>
+
+using namespace DB;
+
+template <typename T, typename U>
+void checkString(const T & str, U manip, const std::string & expected)
+{
+    WriteBufferFromOwnString buf;
+
+    buf << manip << str;
+    EXPECT_EQ(expected, buf.str()) << "str type:" << typeid(str).name();
+}
+
+TEST(OperatorsManipTest, EscapingTest)
+{
+    checkString("Hello 'world'", escape, "Hello \\'world\\'");
+    checkString("Hello \\world\\", escape, "Hello \\\\world\\\\"); // NOLINT
+
+    std::string s1 = "Hello 'world'";
+    checkString(s1, escape, "Hello \\'world\\'");
+    std::string s2 = "Hello \\world\\";
+    checkString(s2, escape, "Hello \\\\world\\\\"); // NOLINT
+
+    std::string_view sv1 = s1;
+    checkString(sv1, escape, "Hello \\'world\\'");
+    std::string_view sv2 = s2;
+    checkString(sv2, escape, "Hello \\\\world\\\\"); // NOLINT
+
+    StringRef sr1 = s1;
+    checkString(sr1, escape, "Hello \\'world\\'");
+    StringRef sr2 = s2;
+    checkString(sr2, escape, "Hello \\\\world\\\\"); // NOLINT
+}
+
+TEST(OperatorsManipTest, QuouteTest)
+{
+    checkString("Hello 'world'", quote, "'Hello \\'world\\''");
+
+    std::string s1 = "Hello 'world'";
+    checkString(s1, quote, "'Hello \\'world\\''");
+
+    std::string_view sv1 = s1;
+    checkString(sv1, quote, "'Hello \\'world\\''");
+
+    StringRef sr1 = s1;
+    checkString(sr1, quote, "'Hello \\'world\\''");
+}
+
+TEST(OperatorsManipTest, DoubleQuouteTest)
+{
+    checkString("Hello 'world'", double_quote, "\"Hello 'world'\"");
+
+    std::string s1 = "Hello 'world'";
+    checkString(s1, double_quote, "\"Hello 'world'\"");
+
+    std::string_view sv1 = s1;
+    checkString(sv1, double_quote, "\"Hello 'world'\"");
+
+    StringRef sr1 = s1;
+    checkString(sr1, double_quote, "\"Hello 'world'\"");
+}
+
+TEST(OperatorsManipTest, binary)
+{
+    checkString("Hello", binary, "\x5Hello");
+
+    std::string s1 = "Hello";
+    checkString(s1, binary, "\x5Hello");
+
+    std::string_view sv1 = s1;
+    checkString(sv1, binary, "\x5Hello");
+
+    StringRef sr1 = s1;
+    checkString(sr1, binary, "\x5Hello");
+}