feat(c/driver/postgresql): Implement Postgres Get table types (#668)

apache · May 10, 2023 · 54b338d · 54b338d
1 parent 6a032c7
commit 54b338d
Show file tree

Hide file tree

Showing 4 changed files with 69 additions and 3 deletions.
diff --git a/c/driver/postgresql/connection.cc b/c/driver/postgresql/connection.cc
@@ -225,6 +225,69 @@ AdbcStatusCode PostgresConnection::GetTableSchema(const char* catalog,
   return final_status;
 }
 
+AdbcStatusCode PostgresConnectionGetTableTypesImpl(struct ArrowSchema* schema,
+                                                   struct ArrowArray* array,
+                                                   struct AdbcError* error) {
+  // See 'relkind' in https://www.postgresql.org/docs/current/catalog-pg-class.html
+  auto uschema = nanoarrow::UniqueSchema();
+  ArrowSchemaInit(uschema.get());
+
+  CHECK_NA(INTERNAL, ArrowSchemaSetType(uschema.get(), NANOARROW_TYPE_STRUCT), error);
+  CHECK_NA(INTERNAL, ArrowSchemaAllocateChildren(uschema.get(), /*num_columns=*/1),
+           error);
+  ArrowSchemaInit(uschema.get()->children[0]);
+  CHECK_NA(INTERNAL,
+           ArrowSchemaSetType(uschema.get()->children[0], NANOARROW_TYPE_STRING), error);
+  CHECK_NA(INTERNAL, ArrowSchemaSetName(uschema.get()->children[0], "table_type"), error);
+  uschema.get()->children[0]->flags &= ~ARROW_FLAG_NULLABLE;
+
+  CHECK_NA(INTERNAL, ArrowArrayInitFromSchema(array, uschema.get(), NULL), error);
+  CHECK_NA(INTERNAL, ArrowArrayStartAppending(array), error);
+
+  CHECK_NA(INTERNAL, ArrowArrayAppendString(array->children[0], ArrowCharView("table")),
+           error);
+  CHECK_NA(INTERNAL, ArrowArrayFinishElement(array), error);
+  CHECK_NA(INTERNAL,
+           ArrowArrayAppendString(array->children[0], ArrowCharView("toast_table")),
+           error);
+  CHECK_NA(INTERNAL, ArrowArrayFinishElement(array), error);
+  CHECK_NA(INTERNAL, ArrowArrayAppendString(array->children[0], ArrowCharView("view")),
+           error);
+  CHECK_NA(INTERNAL, ArrowArrayFinishElement(array), error);
+  CHECK_NA(INTERNAL,
+           ArrowArrayAppendString(array->children[0], ArrowCharView("materialized_view")),
+           error);
+  CHECK_NA(INTERNAL, ArrowArrayFinishElement(array), error);
+  CHECK_NA(INTERNAL,
+           ArrowArrayAppendString(array->children[0], ArrowCharView("foreign_table")),
+           error);
+  CHECK_NA(INTERNAL, ArrowArrayFinishElement(array), error);
+  CHECK_NA(INTERNAL,
+           ArrowArrayAppendString(array->children[0], ArrowCharView("partitioned_table")),
+           error);
+  CHECK_NA(INTERNAL, ArrowArrayFinishElement(array), error);
+
+  CHECK_NA(INTERNAL, ArrowArrayFinishBuildingDefault(array, NULL), error);
+
+  uschema.move(schema);
+  return ADBC_STATUS_OK;
+}
+
+AdbcStatusCode PostgresConnection::GetTableTypes(struct AdbcConnection* connection,
+                                                 struct ArrowArrayStream* out,
+                                                 struct AdbcError* error) {
+  struct ArrowSchema schema = {0};
+  struct ArrowArray array = {0};
+
+  AdbcStatusCode status = PostgresConnectionGetTableTypesImpl(&schema, &array, error);
+  if (status != ADBC_STATUS_OK) {
+    if (schema.release) schema.release(&schema);
+    if (array.release) array.release(&array);
+    return status;
+  }
+  return BatchToArrayStream(&array, &schema, out, error);
+}
+
 AdbcStatusCode PostgresConnection::Init(struct AdbcDatabase* database,
                                         struct AdbcError* error) {
   if (!database || !database->private_data) {

diff --git a/c/driver/postgresql/connection.h b/c/driver/postgresql/connection.h
@@ -38,6 +38,8 @@ class PostgresConnection {
   AdbcStatusCode GetTableSchema(const char* catalog, const char* db_schema,
                                 const char* table_name, struct ArrowSchema* schema,
                                 struct AdbcError* error);
+  AdbcStatusCode GetTableTypes(struct AdbcConnection* connection,
+                               struct ArrowArrayStream* out, struct AdbcError* error);
   AdbcStatusCode Init(struct AdbcDatabase* database, struct AdbcError* error);
   AdbcStatusCode Release(struct AdbcError* error);
   AdbcStatusCode Rollback(struct AdbcError* error);

diff --git a/c/driver/postgresql/postgresql.cc b/c/driver/postgresql/postgresql.cc
@@ -150,7 +150,10 @@ AdbcStatusCode PostgresConnectionGetTableSchema(
 AdbcStatusCode PostgresConnectionGetTableTypes(struct AdbcConnection* connection,
                                                struct ArrowArrayStream* stream,
                                                struct AdbcError* error) {
-  return ADBC_STATUS_NOT_IMPLEMENTED;
+  if (!connection->private_data) return ADBC_STATUS_INVALID_STATE;
+  auto ptr =
+      reinterpret_cast<std::shared_ptr<PostgresConnection>*>(connection->private_data);
+  return (*ptr)->GetTableTypes(connection, stream, error);
 }
 
 AdbcStatusCode PostgresConnectionInit(struct AdbcConnection* connection,

diff --git a/c/driver/postgresql/postgresql_test.cc b/c/driver/postgresql/postgresql_test.cc
@@ -86,8 +86,6 @@ class PostgresConnectionTest : public ::testing::Test,
   void SetUp() override { ASSERT_NO_FATAL_FAILURE(SetUpTest()); }
   void TearDown() override { ASSERT_NO_FATAL_FAILURE(TearDownTest()); }
 
-  void TestMetadataGetTableTypes() { GTEST_SKIP() << "Not yet implemented"; }
-
   void TestMetadataGetObjectsCatalogs() { GTEST_SKIP() << "Not yet implemented"; }
   void TestMetadataGetObjectsDbSchemas() { GTEST_SKIP() << "Not yet implemented"; }
   void TestMetadataGetObjectsTables() { GTEST_SKIP() << "Not yet implemented"; }