From 54b338df0930a4e4c11e63cc5d14a1c8eab71da6 Mon Sep 17 00:00:00 2001 From: William Ayd Date: Wed, 10 May 2023 05:13:22 -0700 Subject: [PATCH] feat(c/driver/postgresql): Implement Postgres Get table types (#668) --- c/driver/postgresql/connection.cc | 63 ++++++++++++++++++++++++++ c/driver/postgresql/connection.h | 2 + c/driver/postgresql/postgresql.cc | 5 +- c/driver/postgresql/postgresql_test.cc | 2 - 4 files changed, 69 insertions(+), 3 deletions(-) diff --git a/c/driver/postgresql/connection.cc b/c/driver/postgresql/connection.cc index 0ae2836b07..893f3c9941 100644 --- 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 index c4218e73e1..5029fe6eb7 100644 --- 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 index d46e0b853a..c879a34c76 100644 --- 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*>(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 index a5613d7c8e..27d4c4904d 100644 --- 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"; }