From f9519991b5f1f9009279f2f06020eba03e922475 Mon Sep 17 00:00:00 2001 From: Geoffrey Yu Date: Mon, 22 Apr 2024 00:42:02 -0400 Subject: [PATCH] Avoid creating estimator in stub mode, parse flags passed to the client --- cpp/flight_sql_example_client.cc | 1 + src/brad/front_end/session.py | 5 +++-- 2 files changed, 4 insertions(+), 2 deletions(-) diff --git a/cpp/flight_sql_example_client.cc b/cpp/flight_sql_example_client.cc index 43f81d19..39922147 100644 --- a/cpp/flight_sql_example_client.cc +++ b/cpp/flight_sql_example_client.cc @@ -92,6 +92,7 @@ arrow::Status Main() { } int main(int argc, char** argv) { + gflags::ParseCommandLineFlags(&argc, &argv, /*remove_flags=*/true); auto status = Main(); if (!status.ok()) { std::cerr << status.ToString() << std::endl; diff --git a/src/brad/front_end/session.py b/src/brad/front_end/session.py index 2edd0a99..09ae5311 100644 --- a/src/brad/front_end/session.py +++ b/src/brad/front_end/session.py @@ -70,7 +70,8 @@ def txn_start_timestamp(self) -> datetime: async def close(self): self._closed = True await self._engines.close() - await self._estimator.close() + if self._estimator is not None: + await self._estimator.close() class SessionManager: @@ -116,7 +117,7 @@ async def create_new_session(self) -> Tuple[SessionId, Session]: routing_policy_override == RoutingPolicy.ForestTableSelectivity or routing_policy_override == RoutingPolicy.Default ): - if self._config.stub_mode_path() is not None: + if self._config.stub_mode_path() is None: estimator: Optional[Estimator] = await PostgresEstimator.connect( self._schema_name, self._config )