Add support for passing tokens in Apache Arrow Flight Tester

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

@@ -67,14 +67,20 @@ message Configuration {
   // The number of seconds between each transfer of data to the remote object store.
   optional uint64 transfer_time_in_seconds = 5;
 
+  // The approximate maximum size, in bytes, of a single WAL segment file before it is closed and a new one is started.
+  uint64 segment_size_threshold_in_bytes = 6;
+
   // Number of threads to allocate for converting multivariate time series to univariate time series.
-  uint32 ingestion_threads = 6;
+  uint32 ingestion_threads = 7;
 
   // Number of threads to allocate for compressing univariate time series to segments.
-  uint32 compression_threads = 7;
+  uint32 compression_threads = 8;
 
   // Number of threads to allocate for writing segments to a local and/or remote data folder.
-  uint32 writer_threads = 8;
+  uint32 writer_threads = 9;
+
+  // Whether the write-ahead log is enabled.
+  bool wal_enabled = 10;
 }
 
 // Request to update the configuration of a ModelarDB node.
@@ -85,6 +91,7 @@ message UpdateConfiguration {
     COMPRESSED_RESERVED_MEMORY_IN_BYTES = 2;
     TRANSFER_BATCH_SIZE_IN_BYTES = 3;
     TRANSFER_TIME_IN_SECONDS = 4;
+    SEGMENT_SIZE_THRESHOLD_IN_BYTES = 5;
   }
 
   // Setting to update in the configuration.

Apache-Arrow-Flight-Tester/protobuf/protocol_pb2.pyi

@@ -49,24 +49,28 @@ class TableMetadata(_message.Message):
     def __init__(self, normal_table: _Optional[_Union[TableMetadata.NormalTableMetadata, _Mapping]] = ..., time_series_table: _Optional[_Union[TableMetadata.TimeSeriesTableMetadata, _Mapping]] = ...) -> None: ...
 
 class Configuration(_message.Message):
-    __slots__ = ("multivariate_reserved_memory_in_bytes", "uncompressed_reserved_memory_in_bytes", "compressed_reserved_memory_in_bytes", "transfer_batch_size_in_bytes", "transfer_time_in_seconds", "ingestion_threads", "compression_threads", "writer_threads")
+    __slots__ = ("multivariate_reserved_memory_in_bytes", "uncompressed_reserved_memory_in_bytes", "compressed_reserved_memory_in_bytes", "transfer_batch_size_in_bytes", "transfer_time_in_seconds", "segment_size_threshold_in_bytes", "ingestion_threads", "compression_threads", "writer_threads", "wal_enabled")
     MULTIVARIATE_RESERVED_MEMORY_IN_BYTES_FIELD_NUMBER: _ClassVar[int]
     UNCOMPRESSED_RESERVED_MEMORY_IN_BYTES_FIELD_NUMBER: _ClassVar[int]
     COMPRESSED_RESERVED_MEMORY_IN_BYTES_FIELD_NUMBER: _ClassVar[int]
     TRANSFER_BATCH_SIZE_IN_BYTES_FIELD_NUMBER: _ClassVar[int]
     TRANSFER_TIME_IN_SECONDS_FIELD_NUMBER: _ClassVar[int]
+    SEGMENT_SIZE_THRESHOLD_IN_BYTES_FIELD_NUMBER: _ClassVar[int]
     INGESTION_THREADS_FIELD_NUMBER: _ClassVar[int]
     COMPRESSION_THREADS_FIELD_NUMBER: _ClassVar[int]
     WRITER_THREADS_FIELD_NUMBER: _ClassVar[int]
+    WAL_ENABLED_FIELD_NUMBER: _ClassVar[int]
     multivariate_reserved_memory_in_bytes: int
     uncompressed_reserved_memory_in_bytes: int
     compressed_reserved_memory_in_bytes: int
     transfer_batch_size_in_bytes: int
     transfer_time_in_seconds: int
+    segment_size_threshold_in_bytes: int
     ingestion_threads: int
     compression_threads: int
     writer_threads: int
-    def __init__(self, multivariate_reserved_memory_in_bytes: _Optional[int] = ..., uncompressed_reserved_memory_in_bytes: _Optional[int] = ..., compressed_reserved_memory_in_bytes: _Optional[int] = ..., transfer_batch_size_in_bytes: _Optional[int] = ..., transfer_time_in_seconds: _Optional[int] = ..., ingestion_threads: _Optional[int] = ..., compression_threads: _Optional[int] = ..., writer_threads: _Optional[int] = ...) -> None: ...
+    wal_enabled: bool
+    def __init__(self, multivariate_reserved_memory_in_bytes: _Optional[int] = ..., uncompressed_reserved_memory_in_bytes: _Optional[int] = ..., compressed_reserved_memory_in_bytes: _Optional[int] = ..., transfer_batch_size_in_bytes: _Optional[int] = ..., transfer_time_in_seconds: _Optional[int] = ..., segment_size_threshold_in_bytes: _Optional[int] = ..., ingestion_threads: _Optional[int] = ..., compression_threads: _Optional[int] = ..., writer_threads: _Optional[int] = ..., wal_enabled: bool = ...) -> None: ...
 
 class UpdateConfiguration(_message.Message):
     __slots__ = ("setting", "new_value")
@@ -77,11 +81,13 @@ class UpdateConfiguration(_message.Message):
         COMPRESSED_RESERVED_MEMORY_IN_BYTES: _ClassVar[UpdateConfiguration.Setting]
         TRANSFER_BATCH_SIZE_IN_BYTES: _ClassVar[UpdateConfiguration.Setting]
         TRANSFER_TIME_IN_SECONDS: _ClassVar[UpdateConfiguration.Setting]
+        SEGMENT_SIZE_THRESHOLD_IN_BYTES: _ClassVar[UpdateConfiguration.Setting]
     MULTIVARIATE_RESERVED_MEMORY_IN_BYTES: UpdateConfiguration.Setting
     UNCOMPRESSED_RESERVED_MEMORY_IN_BYTES: UpdateConfiguration.Setting
     COMPRESSED_RESERVED_MEMORY_IN_BYTES: UpdateConfiguration.Setting
     TRANSFER_BATCH_SIZE_IN_BYTES: UpdateConfiguration.Setting
     TRANSFER_TIME_IN_SECONDS: UpdateConfiguration.Setting
+    SEGMENT_SIZE_THRESHOLD_IN_BYTES: UpdateConfiguration.Setting
     SETTING_FIELD_NUMBER: _ClassVar[int]
     NEW_VALUE_FIELD_NUMBER: _ClassVar[int]
     setting: UpdateConfiguration.Setting

Apache-Arrow-Flight-Tester/server.py

@@ -1,3 +1,5 @@
+import os
+
 import pyarrow
 from pyarrow import flight
 from pyarrow._flight import Result, Ticket
@@ -28,7 +30,7 @@ def workload_balanced_query(self, query: str) -> None:
         cloud_node_url = endpoint.locations[0]
 
         print(f"Executing query on {cloud_node_url}...")
-        cloud_client = ModelarDBServerFlightClient(cloud_node_url)
+        cloud_client = ModelarDBServerFlightClient(cloud_node_url, token=self._token)
         cloud_client.do_get(endpoint.ticket)
 
     def create_table(self, table_name: str, columns: list[tuple[str, str]], time_series_table=False) -> None:
@@ -55,6 +57,10 @@ def vacuum_tables(self, table_names: list[str]) -> None:
         """Vacuum the given tables in the server."""
         self.do_get(Ticket(f"VACUUM {', '.join(table_names)}"))
 
+    def optimize_tables(self, table_names: list[str]) -> None:
+        """Optimize the given tables in the server."""
+        self.do_get(Ticket(f"OPTIMIZE {', '.join(table_names)}"))
+
     def create_normal_table_from_metadata(self, table_name: str, schema: pyarrow.Schema) -> None:
         """Create a normal table using the table name and schema."""
         normal_table_metadata = protocol_pb2.TableMetadata.NormalTableMetadata()
@@ -107,7 +113,9 @@ def node_type(self) -> str:
 
 
 if __name__ == "__main__":
-    server_client = ModelarDBServerFlightClient("grpc://127.0.0.1:9999")
+    token = os.environ.get("MODELARDB_TOKEN")
+    server_client = ModelarDBServerFlightClient("grpc://127.0.0.1:9999", token=token)
+
     print(f"Node type: {server_client.node_type()}\n")
 
     util.create_test_tables(server_client)

Apache-Arrow-Flight-Tester/util.py

@@ -42,9 +42,9 @@ def create_record_batch(num_rows: int) -> pyarrow.RecordBatch:
 def get_time_series_table_schema() -> pyarrow.Schema:
     """Return a schema for a time series table with one timestamp column, three tag columns, and three field columns."""
     return pyarrow.schema([
-        ("location", pyarrow.utf8()),
-        ("install_year", pyarrow.utf8()),
-        ("model", pyarrow.utf8()),
+        ("location", pyarrow.string_view()),
+        ("install_year", pyarrow.string_view()),
+        ("model", pyarrow.string_view()),
         ("timestamp", pyarrow.timestamp("us")),
         ("power_output", pyarrow.float32()),
         ("wind_speed", pyarrow.float32()),
@@ -93,7 +93,7 @@ def create_test_tables_from_metadata(server_client: ModelarDBServerFlightClient)
     normal_table_schema = pyarrow.schema([
         ("timestamp", pyarrow.timestamp("us")),
         ("values", pyarrow.float32()),
-        ("metadata", pyarrow.utf8())
+        ("metadata", pyarrow.string_view())
     ])
 
     server_client.create_normal_table_from_metadata("test_table_1", normal_table_schema)

Apache-Arrow-Flight-Tester/wrapper.py

@@ -5,11 +5,36 @@
 from pyarrow._flight import FlightInfo, ActionType, Result, Ticket
 
 
+class _BearerTokenMiddlewareFactory(flight.ClientMiddlewareFactory):
+    """Creates middleware that attaches a Bearer token to every outgoing call."""
+
+    def __init__(self, token: str, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self._token = token
+
+    def start_call(self, _info) -> flight.ClientMiddleware:
+        return _BearerTokenMiddleware(self._token)
+
+
+class _BearerTokenMiddleware(flight.ClientMiddleware):
+    """Adds 'authorization: Bearer <token>' to the outgoing request headers."""
+
+    def __init__(self, token: str, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self._token = token
+
+    def sending_headers(self) -> dict[str, str]:
+        return {"authorization": f"Bearer {self._token}"}
+
+
 class FlightClientWrapper:
     """Wrapper around the FlightClient class to simplify interaction with an Apache Arrow Flight server."""
 
-    def __init__(self, location: str):
-        self.flight_client = flight.FlightClient(location)
+    def __init__(self, location: str, token: str | None = None):
+        self._token = token
+
+        middleware = [_BearerTokenMiddlewareFactory(token)] if token else []
+        self.flight_client = flight.FlightClient(location, middleware=middleware)
 
     def list_flights(self) -> list[FlightInfo]:
         """Wrapper around the list_flights method of the FlightClient class."""