Use Protocol buffers for Apache Arrow Flight testers (#50)

* Remove functionality related to collect metrics

* Add protobuf definition used in Apache Flight Interface

* Add compiled protobuf file for Python

* Add Python stub file to allow syntax highlighting

* Use protocol buffer in initialize_database()

* Use protocol buffers in register_node()

* Use protocol buffers in remove_node()

* Use protocol buffers for get_configuration()

* Simplify interface for register_node()

* Use protocol buffers in update_configuration()

* Remove deprecated method to encode arguments

* Remove unsued import

* Add back original copyright header

Author: CGodiksen

Apache-Arrow-Flight-Tester/common.py

@@ -49,9 +49,7 @@ def list_table_names(self) -> list[str]:
         flights = self.list_flights()
         return [table_name.decode("utf-8") for table_name in flights[0].descriptor.path]
 
-    def create_table(
-        self, table_name: str, columns: list[tuple[str, str]], time_series_table=False
-    ) -> None:
+    def create_table(self, table_name: str, columns: list[tuple[str, str]], time_series_table=False) -> None:
         """
         Create a table in the server or manager with the given name and columns. Each pair in columns should have the
         format (column_name, column_type).
@@ -102,9 +100,7 @@ def truncate_table(self, table_name: str) -> None:
         """Truncate the table with the given name in the server or manager."""
         self.do_get(Ticket(f"TRUNCATE TABLE {table_name}"))
 
-    def clean_up_tables(
-        self, tables: list[str], operation: Literal["drop", "truncate"]
-    ) -> None:
+    def clean_up_tables(self, tables: list[str], operation: Literal["drop", "truncate"]) -> None:
         """
         Clean up the given tables by either dropping them or truncating them. If no tables are given, all tables
         are dropped or truncated.
@@ -125,11 +121,3 @@ def node_type(self) -> str:
         """Return the type of the node."""
         node_type = self.do_action("NodeType", b"")
         return node_type[0].body.to_pybytes().decode("utf-8")
-
-
-def encode_argument(argument: str) -> bytes:
-    """Encode the given argument as bytes and prepend the size of the argument as a 2-byte integer."""
-    argument_bytes = str.encode(argument)
-    argument_size = len(argument_bytes).to_bytes(2, byteorder="big")
-
-    return argument_size + argument_bytes

Apache-Arrow-Flight-Tester/manager.py

@@ -1,44 +1,49 @@
-from typing import Literal
-
 from pyarrow import flight
 from pyarrow._flight import Result
 
-from common import ModelarDBFlightClient, encode_argument
+from common import ModelarDBFlightClient
+from protobuf import protocol_pb2
 from server import ModelarDBServerFlightClient
 
 
 class ModelarDBManagerFlightClient(ModelarDBFlightClient):
     """Functionality for interacting with a ModelarDB manager using Apache Arrow Flight."""
 
-    def initialize_database(self, existing_tables: list[str]) -> list[str]:
+    def initialize_database(self, existing_tables: list[str]) -> protocol_pb2.TableMetadata:
         """
-        Retrieve the SQL statements required to initialize the database with the tables that are not included in the
+        Retrieve the table metadata required to initialize the database with the tables that are not included in the
         given list of tables. Throws an error if a table in the given list does not exist in the database.
         """
-        result = self.do_action(
-            "InitializeDatabase", str.encode(",".join(existing_tables))
-        )[0]
-        decoded_result = result.body.to_pybytes().decode("utf-8")
+        database_metadata = protocol_pb2.DatabaseMetadata()
+        database_metadata.table_names.extend(existing_tables)
+
+        response = self.do_action("InitializeDatabase", database_metadata.SerializeToString())
 
-        return decoded_result.split(";")
+        table_metadata = protocol_pb2.TableMetadata()
+        table_metadata.ParseFromString(response[0].body.to_pybytes())
 
-    def register_node(
-        self, node_url: str, node_mode: Literal["cloud", "edge"]
-    ) -> list[Result]:
+        return table_metadata
+
+    def register_node(self, node_url: str,
+                      node_mode: protocol_pb2.NodeMetadata.ServerMode) -> protocol_pb2.ManagerMetadata:
         """Register a node with the given URL and mode in the manager."""
-        encoded_node_url = encode_argument(node_url)
-        encoded_node_mode = encode_argument(node_mode)
+        node_metadata = protocol_pb2.NodeMetadata()
+        node_metadata.url = node_url
+        node_metadata.server_mode = node_mode
+
+        response = self.do_action("RegisterNode", node_metadata.SerializeToString())
 
-        action_body = encoded_node_url + encoded_node_mode
-        response = self.do_action("RegisterNode", action_body)
+        manager_metadata = protocol_pb2.ManagerMetadata()
+        manager_metadata.ParseFromString(response[0].body.to_pybytes())
 
-        return response[0].body.to_pybytes()
+        return manager_metadata
 
     def remove_node(self, node_url: str) -> list[Result]:
         """Remove the node with the given URL from the manager."""
-        encoded_node_url = encode_argument(node_url)
+        node_metadata = protocol_pb2.NodeMetadata()
+        node_metadata.url = node_url
 
-        return self.do_action("RemoveNode", encoded_node_url)
+        return self.do_action("RemoveNode", node_metadata.SerializeToString())
 
     def query(self, query: str) -> None:
         """
@@ -65,7 +70,7 @@ def query(self, query: str) -> None:
 
     print(manager_client.initialize_database(["test_table_1"]))
 
-    print(manager_client.register_node("grpc://127.0.0.1:9999", "edge"))
+    print(manager_client.register_node("grpc://127.0.0.1:9999", protocol_pb2.NodeMetadata.ServerMode.EDGE))
     print(manager_client.remove_node("grpc://127.0.0.1:9999"))
 
     manager_client.query("SELECT * FROM test_time_series_table_1 LIMIT 5")

Apache-Arrow-Flight-Tester/protobuf/protocol.proto

@@ -0,0 +1,141 @@
+/* Copyright 2025 The ModelarDB Contributors
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+syntax = "proto3";
+
+package modelardb.flight.protocol;
+
+// Metadata for the ModelarDB cluster manager, including its unique key and storage configuration.
+message ManagerMetadata {
+  // key used to uniquely identify the cluster manager.
+  string key = 1;
+
+  // Storage configuration used to connect to an S3 object store.
+  message S3Configuration {
+    string endpoint = 1;
+    string bucket_name = 2;
+    string access_key_id = 3;
+    string secret_access_key = 4;
+  }
+
+  // Storage configuration used to connect to an Azure Blob Storage object store.
+  message AzureConfiguration {
+    string account_name = 1;
+    string access_key = 2;
+    string container_name = 3;
+  }
+
+  // Storage configuration used by the cluster manager.
+  oneof storage_configuration {
+    S3Configuration s3_configuration = 2;
+    AzureConfiguration azure_configuration = 3;
+  }
+}
+
+// Metadata for a node in the ModelarDB cluster, including its URL and server mode.
+message NodeMetadata {
+  enum ServerMode {
+    CLOUD = 0;
+    EDGE = 1;
+  }
+
+  // gRPC URL of the node.
+  string url = 1;
+
+  // Mode indicating whether the node is a cloud or edge server.
+  ServerMode server_mode = 2;
+}
+
+// Metadata for multiple normal tables and time series tables.
+message TableMetadata {
+  // Metadata for a normal table, including its name and schema.
+  message NormalTableMetadata {
+    string name = 1;
+    bytes schema = 2;
+  }
+
+  // Metadata for a time series table, including its name, schema, error bounds, and generated column expressions.
+  message TimeSeriesTableMetadata {
+    message ErrorBound {
+      enum Type {
+        ABSOLUTE = 0;
+        RELATIVE = 1;
+      }
+      Type type = 1;
+      float value = 2;
+    }
+
+    string name = 1;
+    bytes schema = 2;
+    repeated ErrorBound error_bounds = 3;
+    repeated bytes generated_column_expressions = 4;
+  }
+
+  // Normal tables included in the table metadata.
+  repeated NormalTableMetadata normal_tables = 1;
+
+  // Time series tables included in the table metadata.
+  repeated TimeSeriesTableMetadata time_series_tables = 2;
+}
+
+// Configuration of a ModelarDB node.
+message Configuration {
+  // Amount of memory to reserve for storing multivariate time series.
+  uint64 multivariate_reserved_memory_in_bytes = 1;
+
+  // Amount of memory to reserve for storing uncompressed data buffers.
+  uint64 uncompressed_reserved_memory_in_bytes = 2;
+
+  // Amount of memory to reserve for storing compressed data buffers.
+  uint64 compressed_reserved_memory_in_bytes = 3;
+
+  // The number of bytes that are required before transferring a batch of data to the remote object store.
+  optional uint64 transfer_batch_size_in_bytes = 4;
+
+  // The number of seconds between each transfer of data to the remote object store.
+  optional uint64 transfer_time_in_seconds = 5;
+
+  // Number of threads to allocate for converting multivariate time series to univariate time series.
+  uint32 ingestion_threads = 6;
+
+  // Number of threads to allocate for compressing univariate time series to segments.
+  uint32 compression_threads = 7;
+
+  // Number of threads to allocate for writing segments to a local and/or remote data folder.
+  uint32 writer_threads = 8;
+}
+
+// Request to update the configuration of a ModelarDB node.
+message UpdateConfiguration {
+  enum Setting {
+    MULTIVARIATE_RESERVED_MEMORY_IN_BYTES = 0;
+    UNCOMPRESSED_RESERVED_MEMORY_IN_BYTES = 1;
+    COMPRESSED_RESERVED_MEMORY_IN_BYTES = 2;
+    TRANSFER_BATCH_SIZE_IN_BYTES = 3;
+    TRANSFER_TIME_IN_SECONDS = 4;
+  }
+
+  // Setting to update in the configuration.
+  Setting setting = 1;
+
+  // New value for the setting.
+  optional uint64 new_value = 2;
+}
+
+// Metadata for a ModelarDB database instance.
+message DatabaseMetadata {
+  // Names of the tables in the database.
+  repeated string table_names = 1;
+}