Merge pull request #217 from bmsuisse/updates

bmsdna/lakeapi/context/df_base.py

@@ -2,7 +2,7 @@
 from datetime import datetime, timezone
 
 from sqlglot import Dialect
-from bmsdna.lakeapi.core.types import FileTypes
+from bmsdna.lakeapi.core.types import FileTypes, OperatorType
 from typing import Literal, Optional, List, Tuple, Any, TYPE_CHECKING, Union
 import pyarrow as pa
 from deltalake import DeltaTable
@@ -211,7 +211,7 @@ def get_pyarrow_dataset(
         self,
         uri: SourceUri,
         file_type: FileTypes,
-        partitions: Optional[List[Tuple[str, str, Any]]],
+        partitions: Optional[List[Tuple[str, OperatorType, Any]]],
     ) -> "Optional[pas.Dataset | pa.Table]":
         spec_fs, spec_uri = uri.get_fs_spec()
         match file_type:
@@ -260,7 +260,7 @@ def get_pyarrow_dataset(
                         "Delta table protocol version not supported, use DuckDB or Polars"
                     )
                 return dt.to_pyarrow_dataset(
-                    partitions=partitions,
+                    partitions=partitions,  # type: ignore
                     parquet_read_options={"coerce_int96_timestamp_unit": "us"},
                 )
             case _:
@@ -387,7 +387,7 @@ def register_datasource(
         source_table_name: Optional[str],
         uri: SourceUri,
         file_type: FileTypes,
-        partitions: Optional[List[Tuple[str, str, Any]]],
+        partitions: Optional[List[Tuple[str, OperatorType, Any]]],
     ):
         ds = self.get_pyarrow_dataset(uri, file_type, partitions)
         self.modified_dates[target_name] = self.get_modified_date(uri, file_type)

bmsdna/lakeapi/context/df_duckdb.py

@@ -2,7 +2,7 @@
 
 import pyarrow as pa
 from typing import List, Optional, Tuple, Any, Union, cast
-from bmsdna.lakeapi.core.types import FileTypes
+from bmsdna.lakeapi.core.types import FileTypes, OperatorType
 from bmsdna.lakeapi.context.df_base import ExecutionContext, ResultData, get_sql
 from deltalake2db import duckdb_create_view_for_delta, duckdb_apply_storage_options
 import duckdb
@@ -48,7 +48,7 @@ def __init__(
     def columns(self):
         return self.arrow_schema().names
 
-    def query_builder(self) -> ex.Query:
+    def query_builder(self) -> ex.Select:
         if not isinstance(self.original_sql, str):
             return from_(self.original_sql.subquery())
         else:
@@ -315,7 +315,7 @@ def register_datasource(
         source_table_name: Optional[str],
         uri: SourceUri,
         file_type: FileTypes,
-        partitions: List[Tuple[str, str, Any]] | None,
+        partitions: List[Tuple[str, OperatorType, Any]] | None,
     ):
         self.modified_dates[target_name] = self.get_modified_date(uri, file_type)
 

bmsdna/lakeapi/context/df_odbc.py

@@ -2,7 +2,7 @@
 
 import pyarrow as pa
 from typing import List, Optional, Tuple, Any, Union, cast
-from bmsdna.lakeapi.core.types import FileTypes
+from bmsdna.lakeapi.core.types import FileTypes, OperatorType
 from bmsdna.lakeapi.context.df_base import (
     FLAVORS,
     ExecutionContext,
@@ -70,7 +70,7 @@ def __init__(
     def columns(self):
         return self.arrow_schema().names
 
-    def query_builder(self) -> ex.Query:
+    def query_builder(self) -> ex.Select:
         if not isinstance(self.original_sql, str):
             return from_(self.original_sql.subquery().as_("t"))
         else:
@@ -182,7 +182,7 @@ def register_datasource(
         source_table_name: Optional[str],
         uri: SourceUri,
         file_type: FileTypes,
-        partitions: List[Tuple[str, str, Any]] | None,
+        partitions: List[Tuple[str, OperatorType, Any]] | None,
     ):
         assert file_type == "odbc"
         assert uri.account is None

bmsdna/lakeapi/context/df_polars.py

@@ -93,7 +93,7 @@ def columns(self):
             return _df.columns
         return self._df.columns
 
-    def query_builder(self) -> ex.Query:
+    def query_builder(self) -> ex.Select:
         if not isinstance(self.sql, str):
             return from_(self.sql.subquery(alias="s1"))
         else:

bmsdna/lakeapi/core/datasource.py

@@ -33,7 +33,7 @@
 )
 from bmsdna.lakeapi.core.log import get_logger
 from bmsdna.lakeapi.core.model import get_param_def
-from bmsdna.lakeapi.core.types import DeltaOperatorTypes
+from bmsdna.lakeapi.core.types import DeltaOperatorTypes, OperatorType
 
 logger = get_logger(__name__)
 
@@ -176,7 +176,7 @@ def get_schema(self) -> pa.Schema:
 
     def get_df(
         self,
-        partitions: Optional[List[Tuple[str, str, Any]]] = None,
+        partitions: Optional[List[Tuple[str, OperatorType, Any]]] = None,
         endpoint: endpoints = "request",
     ) -> ResultData:
         if self.df is None:

bmsdna/lakeapi/core/response.py

@@ -199,9 +199,9 @@ def __init__(
         **kwargs,
     ):
         if isinstance(content, typing.AsyncIterable):
-            self.body_iterator = content
+            self.body_iterator = content  # type: ignore
         else:
-            self.body_iterator = iterate_in_threadpool(content)
+            self.body_iterator = iterate_in_threadpool(content)  # type: ignore
 
         # taking over from FileResponse
         self.status_code = status_code