better copy local

bmsdna/lakeapi/api/api.py

@@ -11,7 +11,7 @@ class LakeApiStartInfo:
     config: Configs
 
 
-async def init_lakeapi(
+def init_lakeapi(
     app: FastAPI,
     use_basic_auth: bool,
     start_config: BasicConfig | None = None,
@@ -27,7 +27,7 @@ async def init_lakeapi(
         real_config = Configs.from_yamls(start_config, config)
     else:
         real_config = config
-    router = await init_routes(real_config, start_config)
+    router = init_routes(real_config, start_config)
     if use_basic_auth:
         from bmsdna.lakeapi.core.uservalidation import add_user_middlware
 

bmsdna/lakeapi/context/df_base.py

@@ -8,7 +8,7 @@
 import sqlglot.expressions as ex
 
 from bmsdna.lakeapi.utils.async_utils import _async
-
+from deltalake2db import FilterType
 
 from .source_uri import SourceUri
 
@@ -72,11 +72,17 @@ def __init__(self, chunk_size: int) -> None:
         super().__init__()
         self.chunk_size = chunk_size
 
-    async def columns(self):
-        return (await self.arrow_schema()).names
+    def __enter__(self):
+        return self
+
+    def __exit__(self, *args, **kwargs):
+        pass
+
+    def columns(self):
+        return (self.arrow_schema()).names
 
     @abstractmethod
-    async def arrow_schema(self) -> "pa.Schema": ...
+    def arrow_schema(self) -> "pa.Schema": ...
 
     @abstractmethod
     async def to_pylist(self) -> Sequence[dict]: ...
@@ -204,7 +210,7 @@ def get_pyarrow_dataset(
         self,
         uri: SourceUri,
         file_type: FileTypes,
-        filters: Optional[List[Tuple[str, OperatorType, Any]]],
+        filters: Optional[FilterType],
     ) -> "Optional[pas.Dataset | pa.Table]":
         spec_fs, spec_uri = uri.get_fs_spec()
         match file_type:
@@ -378,15 +384,16 @@ def register_datasource(
         source_table_name: Optional[str],
         uri: SourceUri,
         file_type: FileTypes,
-        filters: Optional[List[Tuple[str, OperatorType, Any]]],
+        filters: Optional[FilterType],
+        meta_only: bool = False,
     ):
         ds = self.get_pyarrow_dataset(uri, file_type, filters)
         self.modified_dates[target_name] = self.get_modified_date(uri, file_type)
         assert ds is not None
         self.register_arrow(target_name, ds)
 
     @abstractmethod
-    async def execute_sql(self, sql: Union[ex.Query, str]) -> ResultData: ...
+    def execute_sql(self, sql: Union[ex.Query, str]) -> ResultData: ...
 
     @abstractmethod
-    async def list_tables(self) -> ResultData: ...
+    def list_tables(self) -> ResultData: ...

bmsdna/lakeapi/context/df_duckdb.py

@@ -5,7 +5,11 @@
 from typing import List, Optional, Tuple, Any, Union, cast
 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
+from deltalake2db import (
+    duckdb_create_view_for_delta,
+    duckdb_apply_storage_options,
+    FilterType,
+)
 import duckdb
 import pyarrow.dataset
 import sqlglot.expressions as ex
@@ -46,8 +50,8 @@ def __init__(
         self._arrow_schema = None
         self._df = None
 
-    async def columns(self):
-        return (await self.arrow_schema()).names
+    def columns(self):
+        return (self.arrow_schema()).names
 
     def query_builder(self) -> ex.Select:
         if not isinstance(self.original_sql, str):
@@ -59,7 +63,7 @@ def query_builder(self) -> ex.Select:
                 ).subquery()
             )
 
-    async def arrow_schema(self) -> pa.Schema:
+    def arrow_schema(self) -> pa.Schema:
         if self._arrow_schema is not None:
             return self._arrow_schema
         query = get_sql(self.original_sql, limit=0, dialect="duckdb")
@@ -68,7 +72,7 @@ def _get_schema():
             self.con.execute(query)
             return self.con.arrow().schema
 
-        self._arrow_schema = await run_in_threadpool(_get_schema)
+        self._arrow_schema = _get_schema()
 
         return self._arrow_schema
 
@@ -207,7 +211,7 @@ def register_arrow(
     def close(self):
         self.con.close()
 
-    async def execute_sql(
+    def execute_sql(
         self,
         sql: Union[
             ex.Query,
@@ -337,7 +341,8 @@ def register_datasource(
         source_table_name: Optional[str],
         uri: SourceUri,
         file_type: FileTypes,
-        filters: List[Tuple[str, OperatorType, Any]] | None,
+        filters: Optional[FilterType],
+        meta_only: bool = False,
     ):
         self.modified_dates[target_name] = self.get_modified_date(uri, file_type)
 
@@ -378,15 +383,12 @@ def register_datasource(
             return
         if file_type == "delta" and uri.exists():
             ab_uri, uri_opts = uri.get_uri_options(flavor="original")
-            df_filter = (
-                {p[0]: p[2] for p in filters if p[1] == "="} if filters else None
-            )
             duckdb_create_view_for_delta(
                 self.con,
                 ab_uri,
                 target_name,
                 storage_options=uri_opts,
-                conditions=df_filter,
+                conditions=filters,
                 use_fsspec=os.getenv("DUCKDB_DELTA_USE_FSSPEC", "0") == "1",
             )
             return
@@ -411,8 +413,8 @@ def register_datasource(
             target_name, source_table_name, uri, file_type, filters
         )
 
-    async def list_tables(self) -> ResultData:
-        return await self.execute_sql(
+    def list_tables(self) -> ResultData:
+        return self.execute_sql(
             """SELECT table_name as name, table_type 
                from information_schema.tables where table_schema='main'
             """

bmsdna/lakeapi/context/df_odbc.py

@@ -78,7 +78,7 @@ def query_builder(self) -> ex.Select:
                 .as_("t")
             )
 
-    async def arrow_schema(self) -> pa.Schema:
+    def arrow_schema(self) -> pa.Schema:
         if self._arrow_schema is not None:
             return self._arrow_schema
         query = get_sql(self.original_sql, limit=0, dialect=self.dialect)
@@ -151,7 +151,7 @@ def dialect(self):
     def supports_view_creation(self) -> bool:
         return False
 
-    async def execute_sql(
+    def execute_sql(
         self,
         sql: Union[
             ex.Query,
@@ -184,14 +184,15 @@ def register_datasource(
         source_table_name: Optional[str],
         uri: SourceUri,
         file_type: FileTypes,
-        filters: List[Tuple[str, OperatorType, Any]] | None,
+        filters: Any,
+        meta_only: bool = False,
     ):
         assert file_type == "odbc"
         assert uri.account is None
         self.datasources[target_name] = uri.uri
 
-    async def list_tables(self) -> ResultData:
-        return await self.execute_sql(
+    def list_tables(self) -> ResultData:
+        return self.execute_sql(
             "SELECT table_schema, table_name as name, table_type from information_schema.tables"
         )
 

bmsdna/lakeapi/context/df_polars.py

@@ -11,6 +11,7 @@
 if TYPE_CHECKING:
     import polars as pl
 from bmsdna.lakeapi.core.types import FileTypes, OperatorType
+from deltalake2db import FilterType
 import pyarrow.dataset
 import sqlglot.expressions as ex
 from sqlglot import from_, parse_one
@@ -67,13 +68,24 @@ def __init__(
         self.random_name = "tbl_" + str(uuid4()).replace("-", "")
         self.registred_df = False
         self.to_register = to_register
+        self._ctx = None
+
+    def __exit__(self, *args, **kwargs):
+        if self._ctx is not None:
+            self._ctx.unregister(self._ctx.tables())
+        self.to_register = {}
+        super().__exit__(*args, **kwargs)
 
     def get_sql_context(self):
+        if self._ctx is not None:
+            self._ctx.unregister(self._ctx.tables())
+            self._ctx = None
         import polars as pl
 
         ctx = pl.SQLContext()
         for k, v in self.to_register.items():
             ctx.register(k, v)
+        self._ctx = ctx
         return ctx
 
     def get_df(self):
@@ -89,7 +101,7 @@ async def get_df_collected(self) -> "pl.DataFrame":
             self._df = _df
         return _df
 
-    async def columns(self):
+    def columns(self):
         if self._df is None:
             _df = pl.DataFrame(
                 [],
@@ -110,7 +122,7 @@ def query_builder(self) -> ex.Select:
                 )
             )
 
-    async def arrow_schema(self) -> pa.Schema:
+    def arrow_schema(self) -> pa.Schema:
         if self._df is None:
             _df = pl.DataFrame(
                 [],
@@ -217,7 +229,8 @@ def register_datasource(
         source_table_name: Optional[str],
         uri: SourceUri,
         file_type: FileTypes,
-        filters: Optional[List[Tuple[str, OperatorType, Any]]],
+        filters: Optional[FilterType],
+        meta_only: bool = False,
     ):
         import polars as pl
 
@@ -230,18 +243,17 @@ def register_datasource(
 
                 try:
                     db_uri, db_opts = uri.get_uri_options(flavor="original")
-                    from deltalake2db import polars_scan_delta
-
-                    df_filter = (
-                        {p[0]: p[2] for p in filters if p[1] == "="}
-                        if filters
-                        else None
-                    )
-                    df = polars_scan_delta(
-                        db_uri,
-                        storage_options=db_opts,
-                        conditions=df_filter if df_filter else None,
-                    )
+                    from deltalake2db import polars_scan_delta, get_polars_schema
+
+                    if meta_only:
+                        schema = get_polars_schema(db_uri, storage_options=db_opts)
+                        df = pl.DataFrame(schema=schema)
+                    else:
+                        df = polars_scan_delta(
+                            db_uri,
+                            storage_options=db_opts,
+                            conditions=filters,
+                        )
                 except DeltaProtocolError as de:
                     raise FileTypeNotSupportedError(
                         f"Delta table version {ab_uri} not supported"
@@ -280,7 +292,7 @@ def register_datasource(
 
         self._to_register[target_name] = df
 
-    async def execute_sql(
+    def execute_sql(
         self,
         sql: Union[
             ex.Query,
@@ -289,8 +301,8 @@ async def execute_sql(
     ) -> PolarsResultData:
         return PolarsResultData(sql, self.chunk_size, self._to_register.copy())
 
-    async def list_tables(self) -> ResultData:
-        return await self.execute_sql("SHOW TABLES")
+    def list_tables(self) -> ResultData:
+        return self.execute_sql("SHOW TABLES")
 
     def __enter__(self):
         return self

bmsdna/lakeapi/context/source_uri.py

@@ -100,33 +100,62 @@ def exists(self) -> bool:
         fs, fs_path = self.get_fs_spec()
         return fs.exists(fs_path)
 
-    def copy_to_local(self, local_path: str):
+    def copy_to_local(self, local_path: str, delta_table: bool):
         if self.account is None:
             raise ValueError("Cannot copy local files")
-        from deltalake2db.delta_meta_retrieval import get_meta, PolarsEngine
-
-        df_uri, df_opts = self.get_uri_options(flavor="object_store")
-        meta = get_meta(PolarsEngine(df_opts), df_uri)
-        vnr = meta.version
-        if local_versions.get(self.uri) == vnr:
+        if not delta_table:
+            fs, fs_path = self.get_fs_spec()
+            stat = fs.stat(fs_path)
+            lm = stat.get("last_modified", stat.get("Last-Modified"))
+            if not lm:
+                raise ValueError("Cannot determine last modified date", stat.keys())
+
+            if local_versions.get(self.uri) != lm:
+                fs.get_file(fs_path, local_path)
+                local_versions[self.uri] = lm
+            return SourceUri(
+                uri=local_path,
+                data_path=None,
+                account=None,
+                accounts=self.accounts,
+                token_retrieval_func=self.token_retrieval_func,
+            )
+        if delta_table:
+            from deltalake2db.delta_meta_retrieval import get_meta, PolarsEngine
+
+            df_uri, df_opts = self.get_uri_options(flavor="object_store")
+            meta = get_meta(PolarsEngine(df_opts), df_uri)
+            vnr = meta.version
+            if local_versions.get(self.uri) == vnr:
+                return SourceUri(
+                    uri=local_path,
+                    data_path=None,
+                    account=None,
+                    accounts=self.accounts,
+                    token_retrieval_func=self.token_retrieval_func,
+                )
+
+            os.makedirs(local_path + "/_delta_log", exist_ok=True)
+            fs, fs_path = self.get_fs_spec()
+            if fs.exists(fs_path.removesuffix("/") + "/_delta_log"):
+                fs.get(
+                    fs_path + "/_delta_log/",
+                    local_path + "/_delta_log/",
+                    recursive=True,
+                )
+                for path in meta.add_actions.keys():
+                    if not os.path.exists(local_path + "/" + path):
+                        fs.get_file(fs_path + "/" + path, local_path + "/" + path)
+            else:
+                fs.get(fs_path, local_path, recursive=True)
+            local_versions[self.uri] = vnr
             return SourceUri(
                 uri=local_path,
                 data_path=None,
                 account=None,
                 accounts=self.accounts,
                 token_retrieval_func=self.token_retrieval_func,
             )
-        os.makedirs(local_path, exist_ok=True)
-        fs, fs_path = self.get_fs_spec()
-        fs.get(fs_path + "/", local_path, recursive=True)
-        local_versions[self.uri] = vnr
-        return SourceUri(
-            uri=local_path,
-            data_path=None,
-            account=None,
-            accounts=self.accounts,
-            token_retrieval_func=self.token_retrieval_func,
-        )
 
     def __str__(self):
         return self.real_uri