AVRO-4026: [C++] Allow non-string custom attributes (#3344)

Author: wgtmac

lang/c++/impl/Compiler.cc

@@ -291,7 +291,8 @@ static void getCustomAttributes(const Object &m, CustomAttributes &customAttribu
     const std::unordered_set<std::string> &kKnownFields = getKnownFields();
     for (const auto &entry : m) {
         if (kKnownFields.find(entry.first) == kKnownFields.end()) {
-            customAttributes.addAttribute(entry.first, entry.second.toLiteralString());
+            bool addQuotes = entry.second.type() == json::EntityType::String;
+            customAttributes.addAttribute(entry.first, entry.second.toLiteralString(), addQuotes);
         }
     }
 }

lang/c++/impl/CustomAttributes.cc

@@ -35,19 +35,28 @@ std::optional<std::string> CustomAttributes::getAttribute(const std::string &nam
 }
 
 void CustomAttributes::addAttribute(const std::string &name,
-                                    const std::string &value) {
+                                    const std::string &value,
+                                    bool addQuotes) {
     auto iter_and_find =
         attributes_.insert(std::pair<std::string, std::string>(name, value));
     if (!iter_and_find.second) {
         throw Exception(name + " already exists and cannot be added");
     }
+    if (addQuotes) {
+        keysNeedQuotes_.insert(name);
+    }
 }
 
 void CustomAttributes::printJson(std::ostream &os,
                                  const std::string &name) const {
-    if (attributes().find(name) == attributes().end()) {
+    auto iter = attributes_.find(name);
+    if (iter == attributes_.cend()) {
         throw Exception(name + " doesn't exist");
     }
-    os << "\"" << name << "\": \"" << attributes().at(name) << "\"";
+    if (keysNeedQuotes_.find(name) != keysNeedQuotes_.cend()) {
+        os << "\"" << name << "\": \"" << iter->second << "\"";
+    } else {
+        os << "\"" << name << "\": " << iter->second;
+    }
 }
 } // namespace avro

lang/c++/include/avro/CustomAttributes.hh

@@ -24,6 +24,7 @@
 #include <map>
 #include <optional>
 #include <string>
+#include <unordered_set>
 
 namespace avro {
 
@@ -37,7 +38,10 @@ public:
     std::optional<std::string> getAttribute(const std::string &name) const;
 
     // Adds a custom attribute. If the attribute already exists, throw an exception.
-    void addAttribute(const std::string &name, const std::string &value);
+    //
+    // If `addQuotes` is true, the `value` will be wrapped in double quotes in the
+    // json serialization; otherwise, the `value` will be serialized as is.
+    void addAttribute(const std::string &name, const std::string &value, bool addQuotes = true);
 
     // Provides a way to iterate over the custom attributes or check attribute size.
     const std::map<std::string, std::string> &attributes() const {
@@ -49,6 +53,7 @@ public:
 
 private:
     std::map<std::string, std::string> attributes_;
+    std::unordered_set<std::string> keysNeedQuotes_;
 };
 
 } // namespace avro

lang/c++/test/SchemaTests.cc

@@ -856,6 +856,60 @@ static void testAddCustomAttributes() {
     BOOST_CHECK_EQUAL(removeWhitespaceFromSchema(json), removeWhitespaceFromSchema(expected));
 }
 
+static void testCustomAttributesJson2Schema2Json() {
+    const std::string schema = R"({
+        "type": "record",
+        "name": "my_record",
+        "fields": [
+            { "name": "long_field", "type": "long", "int_key": 1, "str_key": "1" }
+        ]
+    })";
+    ValidSchema compiledSchema = compileJsonSchemaFromString(schema);
+
+    // Verify custom attributes from parsed schema
+    auto customAttributes = compiledSchema.root()->customAttributesAt(0);
+    BOOST_CHECK_EQUAL(customAttributes.getAttribute("int_key").value(), "1");
+    BOOST_CHECK_EQUAL(customAttributes.getAttribute("str_key").value(), "1");
+
+    // Verify custom attributes from json result
+    std::string json = compiledSchema.toJson();
+    BOOST_CHECK_EQUAL(removeWhitespaceFromSchema(json), removeWhitespaceFromSchema(schema));
+}
+
+static void testCustomAttributesSchema2Json2Schema() {
+    const std::string expected = R"({
+        "type": "record",
+        "name": "my_record",
+        "fields": [
+            { "name": "long_field", "type": "long", "int_key": 1, "str_key": "1" }
+        ]
+    })";
+
+    auto recordNode = std::make_shared<NodeRecord>();
+    {
+        CustomAttributes customAttributes;
+        customAttributes.addAttribute("int_key", "1", /*addQuotes=*/false);
+        customAttributes.addAttribute("str_key", "1", /*addQuotes=*/true);
+        recordNode->addCustomAttributesForField(customAttributes);
+        recordNode->addLeaf(std::make_shared<NodePrimitive>(AVRO_LONG));
+        recordNode->addName("long_field");
+        recordNode->setName(Name("my_record"));
+    }
+
+    // Verify custom attributes from json result
+    ValidSchema schema(recordNode);
+    std::string json = schema.toJson();
+    BOOST_CHECK_EQUAL(removeWhitespaceFromSchema(json), removeWhitespaceFromSchema(expected));
+
+    // Verify custom attributes from parsed schema
+    {
+        auto parsedSchema = compileJsonSchemaFromString(json);
+        auto customAttributes = parsedSchema.root()->customAttributesAt(0);
+        BOOST_CHECK_EQUAL(customAttributes.getAttribute("int_key").value(), "1");
+        BOOST_CHECK_EQUAL(customAttributes.getAttribute("str_key").value(), "1");
+    }
+}
+
 } // namespace schema
 } // namespace avro
 
@@ -882,5 +936,7 @@ init_unit_test_suite(int /*argc*/, char * /*argv*/[]) {
     ts->add(BOOST_TEST_CASE(&avro::schema::testCustomLogicalType));
     ts->add(BOOST_TEST_CASE(&avro::schema::testParseCustomAttributes));
     ts->add(BOOST_TEST_CASE(&avro::schema::testAddCustomAttributes));
+    ts->add(BOOST_TEST_CASE(&avro::schema::testCustomAttributesJson2Schema2Json));
+    ts->add(BOOST_TEST_CASE(&avro::schema::testCustomAttributesSchema2Json2Schema));
     return ts;
 }