fix: Force reading longer batch if type inference failed

Author: neomatamune

cosmotech/coal/cli/commands/store/dump_to_postgresql.py

@@ -8,6 +8,7 @@
 from time import perf_counter
 
 from adbc_driver_postgresql import dbapi
+import pyarrow as pa
 
 from cosmotech.coal.cli.utils.click import click
 from cosmotech.coal.cli.utils.decorators import web_help

cosmotech/coal/store/store.py

@@ -52,10 +52,17 @@ def add_table(self, table_name: str, data=pyarrow.Table, replace: bool = False):
                 LOGGER.debug(f"Inserted {rows} rows in table {table_name}")
 
     def execute_query(self, sql_query: str) -> pyarrow.Table:
-        with dbapi.connect(self._database, autocommit=True) as conn:
-            with conn.cursor() as curs:
-                curs.execute(sql_query)
-                return curs.fetch_arrow_table()
+        batch_size = 1024
+        batch_size_increment = 1024
+        while True:
+            try:
+                with dbapi.connect(self._database, autocommit=True) as conn:
+                    with conn.cursor() as curs:
+                        curs.adbc_statement.set_options(**{"adbc.sqlite.query.batch_rows":str(batch_size)})
+                        curs.execute(sql_query)
+                        return curs.fetch_arrow_table()
+            except OSError:
+                batch_size += batch_size_increment
 
     def list_tables(self) -> list[str]:
         with dbapi.connect(self._database) as conn:

requirements.txt

@@ -7,7 +7,7 @@ azure-kusto-ingest~=4.4.1
 tenacity~=8.3.0
 
 # Keycloak connection
-python-keycloak~=4.3.0
+python-keycloak~=4.7.3
 
 # Modelops requirements
 redis==4.4.4