Merge pull request apache#160 from rafael-telles/flight-sql-cpp-table…

…types [C++] Implement CommandGetTableTypes on server example
rafael-telles · Oct 8, 2021 · a8188fd · a8188fd
2 parents 732d519 + 63cd5b5
commit a8188fd
Show file tree

Hide file tree

Showing 5 changed files with 54 additions and 0 deletions.
diff --git a/cpp/src/arrow/flight/flight-sql/example/sqlite_server.cc b/cpp/src/arrow/flight/flight-sql/example/sqlite_server.cc
@@ -287,6 +287,30 @@ Status SQLiteFlightSqlServer::DoPutCommandStatementUpdate(
   return Status::OK();
 }
 
+Status SQLiteFlightSqlServer::GetFlightInfoTableTypes(const ServerCallContext& context,
+                                                      const FlightDescriptor& descriptor,
+                                                      std::unique_ptr<FlightInfo>* info) {
+  pb::sql::CommandGetTableTypes command;
+  return GetFlightInfoForCommand(descriptor, info, command,
+                                 SqlSchema::GetTableTypesSchema());
+}
+
+Status SQLiteFlightSqlServer::DoGetTableTypes(const ServerCallContext& context,
+                                              std::unique_ptr<FlightDataStream>* result) {
+  std::string query = "SELECT DISTINCT type as table_type FROM sqlite_master";
+
+  std::shared_ptr<SqliteStatement> statement;
+  ARROW_RETURN_NOT_OK(SqliteStatement::Create(db_, query, &statement));
+
+  std::shared_ptr<SqliteStatementBatchReader> reader;
+  ARROW_RETURN_NOT_OK(SqliteStatementBatchReader::Create(
+      statement, SqlSchema::GetTableTypesSchema(), &reader));
+
+  *result = std::unique_ptr<FlightDataStream>(new RecordBatchStream(reader));
+
+  return Status::OK();
+}
+
 }  // namespace example
 }  // namespace sql
 }  // namespace flight

diff --git a/cpp/src/arrow/flight/flight-sql/example/sqlite_server.h b/cpp/src/arrow/flight/flight-sql/example/sqlite_server.h
@@ -74,6 +74,11 @@ class SQLiteFlightSqlServer : public FlightSqlServerBase {
   Status DoGetTables(const pb::sql::CommandGetTables& command,
                      const ServerCallContext& context,
                      std::unique_ptr<FlightDataStream>* result) override;
+  Status GetFlightInfoTableTypes(const ServerCallContext &context,
+                                 const FlightDescriptor &descriptor,
+                                 std::unique_ptr<FlightInfo> *info) override;
+  Status DoGetTableTypes(const ServerCallContext &context,
+                         std::unique_ptr<FlightDataStream> *result) override;
 
  private:
   sqlite3* db_;

diff --git a/cpp/src/arrow/flight/flight-sql/sql_server.cpp b/cpp/src/arrow/flight/flight-sql/sql_server.cpp
@@ -290,6 +290,10 @@ std::shared_ptr<Schema> SqlSchema::GetTablesSchemaWithIncludedSchema() {
                         field("table_schema", binary())});
 }
 
+std::shared_ptr<Schema> SqlSchema::GetTableTypesSchema() {
+  return arrow::schema({field("table_type", utf8())});
+}
+
 }  // namespace sql
 }  // namespace flight
 }  // namespace arrow
diff --git a/cpp/src/arrow/flight/flight-sql/sql_server.h b/cpp/src/arrow/flight/flight-sql/sql_server.h
@@ -294,6 +294,10 @@ class SqlSchema {
   /// flags is set to true.
   /// \return The default schema template.
   static std::shared_ptr<Schema> GetTablesSchemaWithIncludedSchema();
+
+  /// \brief Gets the Schema used on CommandGetTableTypes response.
+  /// \return The default schema template.
+  static std::shared_ptr<Schema> GetTableTypesSchema();
 };
 }  // namespace sql
 }  // namespace flight

diff --git a/cpp/src/arrow/flight/flight-sql/sql_server_test.cc b/cpp/src/arrow/flight/flight-sql/sql_server_test.cc
@@ -285,6 +285,23 @@ TEST(TestFlightSqlServer, TestCommandGetSchemas) {
   ASSERT_EQ(0, table->num_rows());
 }
 
+TEST(TestFlightSqlServer, TestCommandGetTableTypes) {
+  std::unique_ptr<FlightInfo> flight_info;
+  ASSERT_OK(sql_client->GetTableTypes({}, &flight_info));
+
+  std::unique_ptr<FlightStreamReader> stream;
+  ASSERT_OK(sql_client->DoGet({}, flight_info->endpoints()[0].ticket, &stream));
+
+  std::shared_ptr<Table> table;
+  ASSERT_OK(stream->ReadAll(&table));
+
+  DECLARE_ARRAY(table_type, String, ({"table"}));
+
+  const std::shared_ptr<Table>& expected_table =
+      Table::Make(SqlSchema::GetTableTypesSchema(), {table_type});
+  ASSERT_TRUE(expected_table->Equals(*table));
+}
+
 TEST(TestFlightSqlServer, TestCommandStatementUpdate) {
   int64_t result;
   ASSERT_OK(sql_client->ExecuteUpdate(