refactor(c/driver/postgresql): Factor out Postgres type abstraction and test it independently of the driver

.github/workflows/native-unix.yml

@@ -501,7 +501,6 @@ jobs:
 
       - uses: r-lib/actions/setup-r-dependencies@v2
         with:
-          pak-version: devel
           extra-packages: any::rcmdcheck, local::../adbcdrivermanager
           needs: check
           working-directory: r/${{ matrix.config.pkg }}

c/driver/postgresql/CMakeLists.txt

@@ -43,7 +43,6 @@ add_arrow_lib(adbc_driver_postgresql
               database.cc
               postgresql.cc
               statement.cc
-              type.cc
               OUTPUTS
               ADBC_LIBRARIES
               CMAKE_PACKAGE_NAME
@@ -77,6 +76,7 @@ if(ADBC_BUILD_TESTS)
                 PREFIX
                 adbc
                 SOURCES
+                postgres_type_test.cc
                 postgresql_test.cc
                 ../../validation/adbc_validation.cc
                 ../../validation/adbc_validation_util.cc

c/driver/postgresql/CMakeUserPresets.example.json

@@ -0,0 +1,33 @@
+{
+    "version": 3,
+    "cmakeMinimumRequired": {
+        "major": 3,
+        "minor": 21,
+        "patch": 0
+    },
+    "configurePresets": [
+        {
+            "name": "user-local",
+            "displayName": "(user) local build",
+            "cacheVariables": {
+                "CMAKE_BUILD_TYPE": "Debug",
+                "ADBC_BUILD_TESTS": "ON"
+            },
+            "environment": {
+                "PKG_CONFIG_PATH": ""
+            }
+        }
+    ],
+    "testPresets": [
+        {
+            "name": "user-test-preset",
+            "description": "",
+            "displayName": "(user) test preset",
+            "configurePreset": "user-local",
+            "environment": {
+              "CTEST_OUTPUT_ON_FAILURE": "1",
+              "ADBC_POSTGRESQL_TEST_URI": "postgresql://localhost:5432/postgres?user=postgres&password=password"
+            }
+        }
+    ]
+}

c/driver/postgresql/README.md

@@ -50,10 +50,23 @@ $ docker run -it --rm \
     postgres
 ```
 
+Alternatively use the `docker compose` provided by ADBC to manage the test
+database container.
+
+```shell
+$ docker compose up postgres_test
+# When finished:
+# docker compose down postgres_test
+```
+
 Then, to run the tests, set the environment variable specifying the
 PostgreSQL URI before running tests:
 
 ```shell
 $ export ADBC_POSTGRESQL_TEST_URI=postgresql://localhost:5432/postgres?user=postgres&password=password
 $ ctest
 ```
+
+Users of VSCode can use the CMake extension with the supplied CMakeUserPresets.json
+example to supply the required CMake and environment variables required to build and
+run tests.

c/driver/postgresql/connection.cc

@@ -58,7 +58,7 @@ AdbcStatusCode PostgresConnection::Init(struct AdbcDatabase* database,
   }
   database_ =
       *reinterpret_cast<std::shared_ptr<PostgresDatabase>*>(database->private_data);
-  type_mapping_ = database_->type_mapping();
+  type_resolver_ = database_->type_resolver();
   return database_->Connect(&conn_, error);
 }
 

c/driver/postgresql/connection.h

@@ -23,7 +23,7 @@
 #include <adbc.h>
 #include <libpq-fe.h>
 
-#include "type.h"
+#include "postgres_type.h"
 
 namespace adbcpq {
 class PostgresDatabase;
@@ -41,11 +41,13 @@ class PostgresConnection {
   AdbcStatusCode SetOption(const char* key, const char* value, struct AdbcError* error);
 
   PGconn* conn() const { return conn_; }
-  const std::shared_ptr<TypeMapping>& type_mapping() const { return type_mapping_; }
+  const std::shared_ptr<PostgresTypeResolver>& type_resolver() const {
+    return type_resolver_;
+  }
 
  private:
   std::shared_ptr<PostgresDatabase> database_;
-  std::shared_ptr<TypeMapping> type_mapping_;
+  std::shared_ptr<PostgresTypeResolver> type_resolver_;
   PGconn* conn_;
   bool autocommit_;
 };

c/driver/postgresql/database.cc

@@ -19,6 +19,8 @@
 
 #include <cstring>
 #include <memory>
+#include <utility>
+#include <vector>
 
 #include <adbc.h>
 #include <libpq-fe.h>
@@ -29,52 +31,13 @@
 namespace adbcpq {
 
 PostgresDatabase::PostgresDatabase() : open_connections_(0) {
-  type_mapping_ = std::make_shared<TypeMapping>();
+  type_resolver_ = std::make_shared<PostgresTypeResolver>();
 }
 PostgresDatabase::~PostgresDatabase() = default;
 
 AdbcStatusCode PostgresDatabase::Init(struct AdbcError* error) {
   // Connect to validate the parameters.
-  PGconn* conn = nullptr;
-  AdbcStatusCode final_status = Connect(&conn, error);
-  if (final_status != ADBC_STATUS_OK) {
-    return final_status;
-  }
-
-  // Build the type mapping table.
-  const std::string kTypeQuery = R"(
-SELECT
-    oid,
-    typname,
-    typreceive
-FROM
-    pg_catalog.pg_type
-)";
-
-  pg_result* result = PQexec(conn, kTypeQuery.c_str());
-  ExecStatusType pq_status = PQresultStatus(result);
-  if (pq_status == PGRES_TUPLES_OK) {
-    int num_rows = PQntuples(result);
-    for (int row = 0; row < num_rows; row++) {
-      const uint32_t oid = static_cast<uint32_t>(
-          std::strtol(PQgetvalue(result, row, 0), /*str_end=*/nullptr, /*base=*/10));
-      const char* typname = PQgetvalue(result, row, 1);
-      const char* typreceive = PQgetvalue(result, row, 2);
-
-      type_mapping_->Insert(oid, typname, typreceive);
-    }
-  } else {
-    SetError(error, "Failed to build type mapping table: ", PQerrorMessage(conn));
-    final_status = ADBC_STATUS_IO;
-  }
-  PQclear(result);
-
-  // Disconnect since PostgreSQL connections can be heavy.
-  {
-    AdbcStatusCode status = Disconnect(&conn, error);
-    if (status != ADBC_STATUS_OK) final_status = status;
-  }
-  return final_status;
+  return RebuildTypeResolver(error);
 }
 
 AdbcStatusCode PostgresDatabase::Release(struct AdbcError* error) {
@@ -121,4 +84,179 @@ AdbcStatusCode PostgresDatabase::Disconnect(PGconn** conn, struct AdbcError* err
   }
   return ADBC_STATUS_OK;
 }
+
+// Helpers for building the type resolver from queries
+static inline int32_t InsertPgAttributeResult(
+    pg_result* result, const std::shared_ptr<PostgresTypeResolver>& resolver);
+
+static inline int32_t InsertPgTypeResult(
+    pg_result* result, const std::shared_ptr<PostgresTypeResolver>& resolver);
+
+AdbcStatusCode PostgresDatabase::RebuildTypeResolver(struct AdbcError* error) {
+  PGconn* conn = nullptr;
+  AdbcStatusCode final_status = Connect(&conn, error);
+  if (final_status != ADBC_STATUS_OK) {
+    return final_status;
+  }
+
+  // We need a few queries to build the resolver. The current strategy might
+  // fail for some recursive definitions (e.g., arrays of records of arrays).
+  // First, one on the pg_attribute table to resolve column names/oids for
+  // record types.
+  const std::string kColumnsQuery = R"(
+SELECT
+    attrelid,
+    attname,
+    atttypid
+FROM
+    pg_catalog.pg_attribute
+ORDER BY
+    attrelid, attnum
+)";
+
+  // Second, a query of the pg_type table. This query may need a few attempts to handle
+  // recursive definitions (e.g., record types with array column). This currently won't
+  // handle range types because those rows don't have child OID information. Arrays types
+  // are inserted after a successful insert of the element type.
+  const std::string kTypeQuery = R"(
+SELECT
+    oid,
+    typname,
+    typreceive,
+    typbasetype,
+    typarray,
+    typrelid
+FROM
+    pg_catalog.pg_type
+WHERE
+    (typreceive != 0 OR typname = 'aclitem') AND typtype != 'r' AND typreceive::TEXT != 'array_recv'
+ORDER BY
+    oid
+)";
+
+  // Create a new type resolver (this instance's type_resolver_ member
+  // will be updated at the end if this succeeds).
+  auto resolver = std::make_shared<PostgresTypeResolver>();
+
+  // Insert record type definitions (this includes table schemas)
+  pg_result* result = PQexec(conn, kColumnsQuery.c_str());
+  ExecStatusType pq_status = PQresultStatus(result);
+  if (pq_status == PGRES_TUPLES_OK) {
+    InsertPgAttributeResult(result, resolver);
+  } else {
+    SetError(error, "Failed to build type mapping table: ", PQerrorMessage(conn));
+    final_status = ADBC_STATUS_IO;
+  }
+
+  PQclear(result);
+
+  // Attempt filling the resolver a few times to handle recursive definitions.
+  int32_t max_attempts = 3;
+  for (int32_t i = 0; i < max_attempts; i++) {
+    result = PQexec(conn, kTypeQuery.c_str());
+    ExecStatusType pq_status = PQresultStatus(result);
+    if (pq_status == PGRES_TUPLES_OK) {
+      InsertPgTypeResult(result, resolver);
+    } else {
+      SetError(error, "Failed to build type mapping table: ", PQerrorMessage(conn));
+      final_status = ADBC_STATUS_IO;
+    }
+
+    PQclear(result);
+    if (final_status != ADBC_STATUS_OK) {
+      break;
+    }
+  }
+
+  // Disconnect since PostgreSQL connections can be heavy.
+  {
+    AdbcStatusCode status = Disconnect(&conn, error);
+    if (status != ADBC_STATUS_OK) final_status = status;
+  }
+
+  if (final_status == ADBC_STATUS_OK) {
+    type_resolver_ = std::move(resolver);
+  }
+
+  return final_status;
+}
+
+static inline int32_t InsertPgAttributeResult(
+    pg_result* result, const std::shared_ptr<PostgresTypeResolver>& resolver) {
+  int num_rows = PQntuples(result);
+  std::vector<std::pair<std::string, uint32_t>> columns;
+  uint32_t current_type_oid = 0;
+  int32_t n_added = 0;
+
+  for (int row = 0; row < num_rows; row++) {
+    const uint32_t type_oid = static_cast<uint32_t>(
+        std::strtol(PQgetvalue(result, row, 0), /*str_end=*/nullptr, /*base=*/10));
+    const char* col_name = PQgetvalue(result, row, 1);
+    const uint32_t col_oid = static_cast<uint32_t>(
+        std::strtol(PQgetvalue(result, row, 2), /*str_end=*/nullptr, /*base=*/10));
+
+    if (type_oid != current_type_oid && !columns.empty()) {
+      resolver->InsertClass(current_type_oid, columns);
+      columns.clear();
+      current_type_oid = type_oid;
+      n_added++;
+    }
+
+    columns.push_back({col_name, col_oid});
+  }
+
+  if (!columns.empty()) {
+    resolver->InsertClass(current_type_oid, columns);
+    n_added++;
+  }
+
+  return n_added;
+}
+
+static inline int32_t InsertPgTypeResult(
+    pg_result* result, const std::shared_ptr<PostgresTypeResolver>& resolver) {
+  int num_rows = PQntuples(result);
+  PostgresTypeResolver::Item item;
+  int32_t n_added = 0;
+
+  for (int row = 0; row < num_rows; row++) {
+    const uint32_t oid = static_cast<uint32_t>(
+        std::strtol(PQgetvalue(result, row, 0), /*str_end=*/nullptr, /*base=*/10));
+    const char* typname = PQgetvalue(result, row, 1);
+    const char* typreceive = PQgetvalue(result, row, 2);
+    const uint32_t typbasetype = static_cast<uint32_t>(
+        std::strtol(PQgetvalue(result, row, 3), /*str_end=*/nullptr, /*base=*/10));
+    const uint32_t typarray = static_cast<uint32_t>(
+        std::strtol(PQgetvalue(result, row, 4), /*str_end=*/nullptr, /*base=*/10));
+    const uint32_t typrelid = static_cast<uint32_t>(
+        std::strtol(PQgetvalue(result, row, 5), /*str_end=*/nullptr, /*base=*/10));
+
+    // Special case the aclitem because it shows up in a bunch of internal tables
+    if (strcmp(typname, "aclitem") == 0) {
+      typreceive = "aclitem_recv";
+    }
+
+    item.oid = oid;
+    item.typname = typname;
+    item.typreceive = typreceive;
+    item.class_oid = typrelid;
+    item.base_oid = typbasetype;
+
+    int result = resolver->Insert(item, nullptr);
+
+    // If there's an array type and the insert succeeded, add that now too
+    if (result == NANOARROW_OK && typarray != 0) {
+      std::string array_typname = StringBuilder("_", typname);
+      item.oid = typarray;
+      item.typname = array_typname.c_str();
+      item.typreceive = "array_recv";
+      item.child_oid = oid;
+
+      resolver->Insert(item, nullptr);
+    }
+  }
+
+  return n_added;
+}
+
 }  // namespace adbcpq

c/driver/postgresql/database.h

@@ -24,7 +24,7 @@
 #include <adbc.h>
 #include <libpq-fe.h>
 
-#include "type.h"
+#include "postgres_type.h"
 
 namespace adbcpq {
 class PostgresDatabase {
@@ -42,12 +42,15 @@ class PostgresDatabase {
 
   AdbcStatusCode Connect(PGconn** conn, struct AdbcError* error);
   AdbcStatusCode Disconnect(PGconn** conn, struct AdbcError* error);
+  const std::shared_ptr<PostgresTypeResolver>& type_resolver() const {
+    return type_resolver_;
+  }
 
-  const std::shared_ptr<TypeMapping>& type_mapping() const { return type_mapping_; }
+  AdbcStatusCode RebuildTypeResolver(struct AdbcError* error);
 
  private:
   int32_t open_connections_;
   std::string uri_;
-  std::shared_ptr<TypeMapping> type_mapping_;
+  std::shared_ptr<PostgresTypeResolver> type_resolver_;
 };
 }  // namespace adbcpq