fix(c/driver/postgresql): add postgres type to cols created for numer…

…ic (#1516) - introduce constant with key name - introduce separate private method to add field-level method Fixes #1515
apache · Feb 6, 2024 · 51abae1 · 51abae1
1 parent 3828e0b
commit 51abae1
Show file tree

Hide file tree

Showing 2 changed files with 50 additions and 13 deletions.
diff --git a/c/driver/postgresql/postgres_type.h b/c/driver/postgresql/postgres_type.h
@@ -218,6 +218,8 @@ class PostgresType {
       // ---- Numeric/Decimal-------------------
       case PostgresTypeId::kNumeric:
         NANOARROW_RETURN_NOT_OK(ArrowSchemaSetType(schema, NANOARROW_TYPE_STRING));
+        NANOARROW_RETURN_NOT_OK(AddPostgresTypeMetadata(schema));
+
         break;
 
       // ---- Binary/string --------------------
@@ -279,20 +281,13 @@ class PostgresType {
         break;
 
       case PostgresTypeId::kUserDefined:
-      default: {
+      default:
         // For user-defined types or types we don't explicitly know how to deal with, we
         // can still return the bytes postgres gives us and attach the type name as
         // metadata
         NANOARROW_RETURN_NOT_OK(ArrowSchemaSetType(schema, NANOARROW_TYPE_BINARY));
-        nanoarrow::UniqueBuffer buffer;
-        ArrowMetadataBuilderInit(buffer.get(), nullptr);
-        NANOARROW_RETURN_NOT_OK(ArrowMetadataBuilderAppend(
-            buffer.get(), ArrowCharView("ADBC:postgresql:typname"),
-            ArrowCharView(typname_.c_str())));
-        NANOARROW_RETURN_NOT_OK(
-            ArrowSchemaSetMetadata(schema, reinterpret_cast<char*>(buffer->data)));
+        NANOARROW_RETURN_NOT_OK(AddPostgresTypeMetadata(schema));
         break;
-      }
     }
 
     NANOARROW_RETURN_NOT_OK(ArrowSchemaSetName(schema, field_name_.c_str()));
@@ -309,6 +304,25 @@ class PostgresType {
   std::string typname_;
   std::string field_name_;
   std::vector<PostgresType> children_;
+
+  static constexpr const char* kPostgresTypeKey = "ADBC:postgresql:typname";
+
+  ArrowErrorCode AddPostgresTypeMetadata(ArrowSchema* schema) const {
+    // the typname_ may not always be set: an instance of this class can be
+    // created with just the type id. That's why there is this here fallback to
+    // resolve the type name of built-in types.
+    const char* typname =
+        !typname_.empty() ? typname_.c_str() : PostgresTypname(type_id_);
+    nanoarrow::UniqueBuffer buffer;
+
+    ArrowMetadataBuilderInit(buffer.get(), nullptr);
+    NANOARROW_RETURN_NOT_OK(ArrowMetadataBuilderAppend(
+        buffer.get(), ArrowCharView(kPostgresTypeKey), ArrowCharView(typname)));
+    NANOARROW_RETURN_NOT_OK(
+        ArrowSchemaSetMetadata(schema, reinterpret_cast<char*>(buffer->data)));
+
+    return NANOARROW_OK;
+  }
 };
 
 // Because type information is stored in a database's pg_type table, it can't

diff --git a/c/driver/postgresql/postgres_type_test.cc b/c/driver/postgresql/postgres_type_test.cc
@@ -124,6 +124,7 @@ TEST(PostgresTypeTest, PostgresTypeBasic) {
 
 TEST(PostgresTypeTest, PostgresTypeSetSchema) {
   nanoarrow::UniqueSchema schema;
+  ArrowStringView typnameMetadataValue = ArrowCharView("<not found>");
 
   ArrowSchemaInit(schema.get());
   EXPECT_EQ(PostgresType(PostgresTypeId::kBool).SetSchema(schema.get()), NANOARROW_OK);
@@ -165,6 +166,28 @@ TEST(PostgresTypeTest, PostgresTypeSetSchema) {
   EXPECT_STREQ(schema->format, "z");
   schema.reset();
 
+  ArrowSchemaInit(schema.get());
+  EXPECT_EQ(PostgresType(PostgresTypeId::kNumeric).SetSchema(schema.get()), NANOARROW_OK);
+  EXPECT_STREQ(schema->format, "u");
+  typnameMetadataValue = ArrowCharView("<not found>");
+  ArrowMetadataGetValue(schema->metadata, ArrowCharView("ADBC:postgresql:typname"),
+                        &typnameMetadataValue);
+  EXPECT_EQ(std::string(typnameMetadataValue.data, typnameMetadataValue.size_bytes),
+            "numeric");
+  schema.reset();
+
+  ArrowSchemaInit(schema.get());
+  EXPECT_EQ(PostgresType(PostgresTypeId::kNumeric).Array().SetSchema(schema.get()),
+            NANOARROW_OK);
+  EXPECT_STREQ(schema->format, "+l");
+  EXPECT_STREQ(schema->children[0]->format, "u");
+  typnameMetadataValue = ArrowCharView("<not found>");
+  ArrowMetadataGetValue(schema->children[0]->metadata,
+                        ArrowCharView("ADBC:postgresql:typname"), &typnameMetadataValue);
+  EXPECT_EQ(std::string(typnameMetadataValue.data, typnameMetadataValue.size_bytes),
+            "numeric");
+  schema.reset();
+
   ArrowSchemaInit(schema.get());
   EXPECT_EQ(PostgresType(PostgresTypeId::kBool).Array().SetSchema(schema.get()),
             NANOARROW_OK);
@@ -184,11 +207,11 @@ TEST(PostgresTypeTest, PostgresTypeSetSchema) {
   PostgresType unknown(PostgresTypeId::kBrinMinmaxMultiSummary);
   EXPECT_EQ(unknown.WithPgTypeInfo(0, "some_name").SetSchema(schema.get()), NANOARROW_OK);
   EXPECT_STREQ(schema->format, "z");
-
-  ArrowStringView value = ArrowCharView("<not found>");
+  typnameMetadataValue = ArrowCharView("<not found>");
   ArrowMetadataGetValue(schema->metadata, ArrowCharView("ADBC:postgresql:typname"),
-                        &value);
-  EXPECT_EQ(std::string(value.data, value.size_bytes), "some_name");
+                        &typnameMetadataValue);
+  EXPECT_EQ(std::string(typnameMetadataValue.data, typnameMetadataValue.size_bytes),
+            "some_name");
   schema.reset();
 }