diff --git a/bmsdna/lakeapi/api/api.py b/bmsdna/lakeapi/api/api.py index d10b883..f248f75 100644 --- a/bmsdna/lakeapi/api/api.py +++ b/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 diff --git a/bmsdna/lakeapi/context/df_base.py b/bmsdna/lakeapi/context/df_base.py index 1dcaa8f..bc87b13 100644 --- a/bmsdna/lakeapi/context/df_base.py +++ b/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,7 +384,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, ): ds = self.get_pyarrow_dataset(uri, file_type, filters) self.modified_dates[target_name] = self.get_modified_date(uri, file_type) @@ -386,7 +393,7 @@ def register_datasource( 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: ... diff --git a/bmsdna/lakeapi/context/df_duckdb.py b/bmsdna/lakeapi/context/df_duckdb.py index 34b840b..dda7bbb 100644 --- a/bmsdna/lakeapi/context/df_duckdb.py +++ b/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' """ diff --git a/bmsdna/lakeapi/context/df_odbc.py b/bmsdna/lakeapi/context/df_odbc.py index 121f266..20e699c 100644 --- a/bmsdna/lakeapi/context/df_odbc.py +++ b/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" ) diff --git a/bmsdna/lakeapi/context/df_polars.py b/bmsdna/lakeapi/context/df_polars.py index ea6b226..0789278 100644 --- a/bmsdna/lakeapi/context/df_polars.py +++ b/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 diff --git a/bmsdna/lakeapi/context/source_uri.py b/bmsdna/lakeapi/context/source_uri.py index f47817c..e3dc333 100644 --- a/bmsdna/lakeapi/context/source_uri.py +++ b/bmsdna/lakeapi/context/source_uri.py @@ -100,15 +100,55 @@ 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, @@ -116,17 +156,6 @@ def copy_to_local(self, local_path: str): 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 diff --git a/bmsdna/lakeapi/core/datasource.py b/bmsdna/lakeapi/core/datasource.py index 78ec63e..4dbd0e0 100644 --- a/bmsdna/lakeapi/core/datasource.py +++ b/bmsdna/lakeapi/core/datasource.py @@ -38,6 +38,7 @@ MetaState as DeltaMetadata, ) from bmsdna.lakeapi.utils.async_utils import _async +from deltalake2db.filter_by_meta import FilterType, Operator as SupportedOperator logger = get_logger(__name__) @@ -138,19 +139,24 @@ def __init__( def __str__(self) -> str: return str(self.source_uri) + def __enter__(self): + return self + + def __exit__(self, *args): + if self.df: + self.df.__exit__(*args) + @property def execution_uri(self): if not self.copy_local: return self.source_uri - assert self.config.file_type == "delta", ( - "only delta is supported for copy_local" - ) if self._execution_uri is None: self._execution_uri = self.source_uri.copy_to_local( os.path.join( self.basic_config.local_data_cache_path, hashlib.md5(self.source_uri.uri.encode("utf-8")).hexdigest(), - ) + ), + self.config.file_type == "delta", ) return self._execution_uri @@ -204,7 +210,7 @@ def get_table_name(self, force_unique_name: bool) -> ex.Table: ) return ex.table_(tname) - async def get_schema(self) -> pa.Schema: + def get_schema(self) -> pa.Schema: schema: pa.Schema | None = None if self.config.file_type == "delta" and self.file_exists(): dt = self.get_delta_table(schema_only=True) @@ -223,11 +229,11 @@ async def get_schema(self) -> pa.Schema: ] return pyarrow.schema(fields) return schema - return await (await self.get_df(endpoint="meta")).arrow_schema() + return (self.get_df(endpoint="meta")).arrow_schema() - async def get_df( + def get_df( self, - filters: Optional[List[Tuple[str, OperatorType, Any]]] = None, + filters: Optional[FilterType] = None, endpoint: endpoints = "request", ) -> ResultData: if self.df is None: @@ -249,8 +255,9 @@ async def get_df( self.source_uri if endpoint == "meta" else self.execution_uri, self.config.file_type, filters=filters, + meta_only=endpoint == "meta", ) - self.df = await _async(self.sql_context.execute_sql(self.query)) + self.df = self.sql_context.execute_sql(self.query) return self.df # type: ignore @@ -409,12 +416,13 @@ def _sql_value(value: str | datetime | date | None, engine: str): return value -def get_filter_dict( +def get_filters( params: dict[str, Any], param_def: list[Union[Param, str]], columns: Optional[Sequence[str]], -) -> dict[str, Any]: - result = {} +) -> tuple[FilterType, bool]: + result: list[tuple[str, SupportedOperator, Any]] = [] + covered_all = True for key, value in params.items(): if not key or not value or key in ("limit", "offset"): continue # can that happen? I don't know @@ -425,21 +433,18 @@ def get_filter_dict( colname = prmdef.colname or prmdef.name if prmdef.combi: + covered_all = False continue elif columns and colname not in columns: + covered_all = False pass else: - match op: - case "==": - result[colname] = value if value is not None else None - case "=": - result[colname] = value if value is not None else None - case "in": - lsv = cast(list[str], value) - if len(lsv) == 1: - result[colname] = lsv[0] + if op not in get_args(SupportedOperator): + covered_all = False + else: + result.append((colname, cast(SupportedOperator, op), value)) - return result + return result, covered_all def filter_df_based_on_params( @@ -505,14 +510,6 @@ def filter_df_based_on_params( if value is not None else ~ex.column(colname, quoted=True).is_(ex.convert(None)) ) - case "==": - exprs.append( - ex.column(colname, quoted=True).eq( - _sql_value(value, engine=context.engine_name) - ) - if value is not None - else ex.column(colname, quoted=True).is_(ex.convert(None)) - ) case "=": exprs.append( ex.column(colname, quoted=True).eq( diff --git a/bmsdna/lakeapi/core/response.py b/bmsdna/lakeapi/core/response.py index 830f68e..611f5ca 100644 --- a/bmsdna/lakeapi/core/response.py +++ b/bmsdna/lakeapi/core/response.py @@ -288,17 +288,19 @@ async def create_response( ) if format == OutputFormats.JSON: - return Response( - content=await _async((await _async(context.execute_sql(sql))).to_json()), - headers=headers, - media_type=(media_type or "application/json") + "; charset=" + charset, - ) + with context.execute_sql(sql) as res: + return Response( + content=await _async(res.to_json()), + headers=headers, + media_type=(media_type or "application/json") + "; charset=" + charset, + ) if format == OutputFormats.ND_JSON: - return Response( - content=await _async((await _async(context.execute_sql(sql))).to_ndjson()), - headers=headers, - media_type="application/json-nd; charset=" + charset, - ) + with context.execute_sql(sql) as res: + return Response( + content=await _async(res.to_ndjson()), + headers=headers, + media_type="application/json-nd; charset=" + charset, + ) if format in [ OutputFormats.JSON, @@ -314,16 +316,16 @@ async def create_response( async def response_stream(context: ExecutionContext, sql, url, format): chunk_size = 64 * 1024 - content = await _async(context.execute_sql(sql)) - additional_files = await _write_frame( - url, - content, - format, - temp_file.name, - basic_config, - query_params.get("$csv_separator", None), - charset, - ) + with context.execute_sql(sql) as content: + additional_files = await _write_frame( + url, + content, + format, + temp_file.name, + basic_config, + query_params.get("$csv_separator", None), + charset, + ) async with await anyio.open_file(temp_file.name, mode="rb") as file: more_body = True diff --git a/bmsdna/lakeapi/core/route.py b/bmsdna/lakeapi/core/route.py index 2bde24a..5438c41 100644 --- a/bmsdna/lakeapi/core/route.py +++ b/bmsdna/lakeapi/core/route.py @@ -12,7 +12,7 @@ all_lake_api_routers: list[Tuple[BasicConfig, Configs]] = [] -async def init_routes(configs: Configs, basic_config: BasicConfig): +def init_routes(configs: Configs, basic_config: BasicConfig): from bmsdna.lakeapi.endpoint.endpoint import ( get_response_model, create_config_endpoint, @@ -38,7 +38,7 @@ async def init_routes(configs: Configs, basic_config: BasicConfig): from bmsdna.lakeapi.core.datasource import Datasource assert config.datasource is not None - realdataframe = Datasource( + with Datasource( config.version_str, config.tag, config.name, @@ -46,21 +46,23 @@ async def init_routes(configs: Configs, basic_config: BasicConfig): sql_context=mgr.get_context(config.engine), basic_config=basic_config, accounts=configs.accounts, - ) - try: - schema = await get_schema_cached( - basic_config, realdataframe, config.datasource.get_unique_hash() - ) - if schema is None: + ) as realdataframe: + try: + schema = get_schema_cached( + basic_config, + realdataframe, + config.datasource.get_unique_hash(), + ) + if schema is None: + logger.warning( + f"Could not get response type for {config.route}. Path does not exist:{realdataframe}" + ) + except Exception as err: logger.warning( - f"Could not get response type for {config.route}. Path does not exist:{realdataframe}" + f"Could not get schema for {config.route}. Error:{err}", + exc_info=err, ) - except Exception as err: - logger.warning( - f"Could not get schema for {config.route}. Error:{err}", - exc_info=err, - ) - schema = None + schema = None metadata.append( { diff --git a/bmsdna/lakeapi/core/schema_cache.py b/bmsdna/lakeapi/core/schema_cache.py index 73a9586..5147050 100644 --- a/bmsdna/lakeapi/core/schema_cache.py +++ b/bmsdna/lakeapi/core/schema_cache.py @@ -3,7 +3,7 @@ from bmsdna.lakeapi.utils.async_utils import _async -async def get_schema_cached(cfg: BasicConfig, datasource: Datasource, key: str): +def get_schema_cached(cfg: BasicConfig, datasource: Datasource, key: str): if cfg.schema_cache_ttl is not None: import os import time @@ -21,7 +21,7 @@ async def get_schema_cached(cfg: BasicConfig, datasource: Datasource, key: str): if not datasource.file_exists(): schema = None else: - schema = await datasource.get_schema() + schema = datasource.get_schema() pq.write_metadata(schema, schema_cache_file) else: schema = pq.read_schema(schema_cache_file) @@ -29,4 +29,4 @@ async def get_schema_cached(cfg: BasicConfig, datasource: Datasource, key: str): else: if not datasource.file_exists(): return None - return await _async(datasource.get_schema()) + return datasource.get_schema() diff --git a/bmsdna/lakeapi/endpoint/detail_endpoint.py b/bmsdna/lakeapi/endpoint/detail_endpoint.py index 53a28fe..f62883d 100644 --- a/bmsdna/lakeapi/endpoint/detail_endpoint.py +++ b/bmsdna/lakeapi/endpoint/detail_endpoint.py @@ -60,7 +60,7 @@ async def get_detailed_metadata( basic_config.default_chunk_size, ) as context: assert config.datasource is not None - realdataframe = Datasource( + with Datasource( config.version_str, config.tag, config.name, @@ -68,138 +68,145 @@ async def get_detailed_metadata( sql_context=context, basic_config=basic_config, accounts=configs.accounts, - ) - - if not realdataframe.file_exists(): - raise HTTPException(404) - partition_columns = [] - partition_values = None - delta_tbl = None - df = await realdataframe.get_df(None) - if config.datasource.file_type == "delta": - delta_tbl = realdataframe.get_delta_table(schema_only=True) - assert delta_tbl is not None - assert delta_tbl.last_metadata is not None - partition_columns = delta_tbl.last_metadata.get("partitionColumns", []) - partition_columns = [ - c - for c in partition_columns - if not basic_config.should_hide_col_name(c) - ] # also hide those from metadata detail - if len(partition_columns) > 0: - qb = cast( - ex.Select, - df.query_builder().select( - *[ex.column(c, quoted=True) for c in partition_columns], - append=False, - ), - ).distinct() - partition_values = await (await context.execute_sql(qb)).to_pylist() - schema = await df.arrow_schema() - str_cols = [ - name - for name in schema.names - if ( - pa.types.is_string(schema.field(name).type) - or pa.types.is_large_string(schema.field(name).type) - ) - and not basic_config.should_hide_col_name(name) - ] - complex_str_cols = ( - [ + ) as realdataframe: + if not realdataframe.file_exists(): + raise HTTPException(404) + partition_columns = [] + partition_values = None + delta_tbl = None + df = realdataframe.get_df(None) + if config.datasource.file_type == "delta": + delta_tbl = realdataframe.get_delta_table(schema_only=True) + assert delta_tbl is not None + assert delta_tbl.last_metadata is not None + partition_columns = delta_tbl.last_metadata.get( + "partitionColumns", [] + ) + partition_columns = [ + c + for c in partition_columns + if not basic_config.should_hide_col_name(c) + ] # also hide those from metadata detail + if len(partition_columns) > 0: + qb = cast( + ex.Select, + df.query_builder().select( + *[ex.column(c, quoted=True) for c in partition_columns], + append=False, + ), + ).distinct() + with context.execute_sql(qb) as res: + partition_values = await res.to_pylist() + schema = df.arrow_schema() + str_cols = [ name for name in schema.names if ( - pa.types.is_struct(schema.field(name).type) - or pa.types.is_list(schema.field(name).type) + pa.types.is_string(schema.field(name).type) + or pa.types.is_large_string(schema.field(name).type) ) and not basic_config.should_hide_col_name(name) ] - if jsonify_complex - else [] - ) - if include_str_lengths: - str_lengths_query = ( - select( - *[ - ex.func( - "MAX", - ex.func(context.len_func, ex.column(sc, quoted=True)), - ).as_(sc) - for sc in str_cols + complex_str_cols - ] - ) - .from_("strcols") - .with_( - "strcols", - as_=context.jsonify_complex( - df.query_builder(), - complex_str_cols, - str_cols + complex_str_cols, - ), - ) + complex_str_cols = ( + [ + name + for name in schema.names + if ( + pa.types.is_struct(schema.field(name).type) + or pa.types.is_list(schema.field(name).type) + ) + and not basic_config.should_hide_col_name(name) + ] + if jsonify_complex + else [] ) - str_lengths_df = ( - (await (await context.execute_sql(str_lengths_query)).to_pylist()) - if len(str_cols) > 0 or len(complex_str_cols) > 0 - else [{}] - ) - str_lengths = str_lengths_df[0] - else: - str_lengths = {} + if include_str_lengths: + str_lengths_query = ( + select( + *[ + ex.func( + "MAX", + ex.func( + context.len_func, ex.column(sc, quoted=True) + ), + ).as_(sc) + for sc in str_cols + complex_str_cols + ] + ) + .from_("strcols") + .with_( + "strcols", + as_=context.jsonify_complex( + df.query_builder(), + complex_str_cols, + str_cols + complex_str_cols, + ), + ) + ) + with context.execute_sql(str_lengths_query) as res: + str_lengths_df = ( + (await res.to_pylist()) + if len(str_cols) > 0 or len(complex_str_cols) > 0 + else [{}] + ) + str_lengths = str_lengths_df[0] + else: + str_lengths = {} - def _recursive_get_type(t: pa.DataType) -> MetadataSchemaFieldType: - is_complex = pa.types.is_nested(t) + def _recursive_get_type(t: pa.DataType) -> MetadataSchemaFieldType: + is_complex = pa.types.is_nested(t) - return MetadataSchemaFieldType( - type_str=str(pa.string()) - if is_complex and jsonify_complex - else str(t), - orig_type_str=str(t), - fields=( - [ - MetadataSchemaField( - name=f.name, type=_recursive_get_type(f.type) - ) - for f in [ - t.field(find) - for find in range(0, cast(pa.StructType, t).num_fields) + return MetadataSchemaFieldType( + type_str=str(pa.string()) + if is_complex and jsonify_complex + else str(t), + orig_type_str=str(t), + fields=( + [ + MetadataSchemaField( + name=f.name, type=_recursive_get_type(f.type) + ) + for f in [ + t.field(find) + for find in range( + 0, cast(pa.StructType, t).num_fields + ) + ] ] - ] - if pa.types.is_struct(t) and not jsonify_complex - else None - ), - inner=( - _recursive_get_type(cast(pa.ListType, t).value_type) - if pa.types.is_list(t) - or pa.types.is_large_list(t) - or pa.types.is_fixed_size_list(t) - and cast(pa.ListType, t).value_type is not None - and not jsonify_complex - else None - ), - ) - - schema = await df.arrow_schema() - mdt = realdataframe.sql_context.get_modified_date( - realdataframe.source_uri, realdataframe.config.file_type - ) - return MetadataDetailResult( - partition_values=partition_values, - partition_columns=partition_columns, - max_string_lengths=str_lengths, - data_schema=[ - MetadataSchemaField( - name=n, - type=_recursive_get_type(schema.field(n).type), - max_str_length=str_lengths.get(n, None), + if pa.types.is_struct(t) and not jsonify_complex + else None + ), + inner=( + _recursive_get_type(cast(pa.ListType, t).value_type) + if pa.types.is_list(t) + or pa.types.is_large_list(t) + or pa.types.is_fixed_size_list(t) + and cast(pa.ListType, t).value_type is not None + and not jsonify_complex + else None + ), ) - for n in schema.names - if not basic_config.should_hide_col_name(n) - ], - delta_meta=delta_tbl.last_metadata if delta_tbl else None, - delta_schema=delta_tbl.schema if delta_tbl else None, - parameters=config.params, # type: ignore - search=config.search, - modified_date=mdt, - ) + + schema = df.arrow_schema() + mdt = realdataframe.sql_context.get_modified_date( + realdataframe.source_uri, realdataframe.config.file_type + ) + return MetadataDetailResult( + partition_values=partition_values, + partition_columns=partition_columns, + max_string_lengths=str_lengths, + data_schema=[ + MetadataSchemaField( + name=n, + type=_recursive_get_type(schema.field(n).type), + max_str_length=str_lengths.get(n, None), + ) + for n in schema.names + if not basic_config.should_hide_col_name(n) + ], + delta_meta=delta_tbl.last_metadata if delta_tbl else None, + delta_schema=delta_tbl.schema if delta_tbl else None, + parameters=config.params, # type: ignore + search=config.search, + modified_date=mdt, + ) diff --git a/bmsdna/lakeapi/endpoint/endpoint.py b/bmsdna/lakeapi/endpoint/endpoint.py index 8a4f14b..5a02b92 100644 --- a/bmsdna/lakeapi/endpoint/endpoint.py +++ b/bmsdna/lakeapi/endpoint/endpoint.py @@ -15,7 +15,7 @@ from bmsdna.lakeapi.core.datasource import ( Datasource, filter_df_based_on_params, - get_filter_dict, + get_filters, filter_partitions_based_on_params, ) from bmsdna.lakeapi.core.log import get_logger @@ -35,36 +35,6 @@ logger = get_logger(__name__) -async def get_partitions( - datasource: Datasource, - uri: SourceUri, - params: BaseModel, - config: Config, -) -> Optional[list[tuple[str, OperatorType, Any]]]: - try: - df_uri, df_opts = uri.get_uri_options(flavor="object_store") - - def _schema_meta(): - meta = get_meta(PolarsEngine(df_opts), df_uri) - assert meta.last_metadata is not None - return meta - - meta = await run_in_threadpool(_schema_meta) - parts = ( - filter_partitions_based_on_params( - meta, - params.model_dump(exclude_unset=True) if params else {}, - config.params or [], - ) - if not config.datasource or config.datasource.file_type == "delta" - else None - ) - except Exception as err: - logger.warning(f"Could not get partitions for {uri}", exc_info=err) - parts = None - return parts - - def remove_search_nearby( prm_dict: dict, config: Config, @@ -242,7 +212,7 @@ async def data( chunk_size=real_chunk_size, ) assert config.datasource is not None - realdataframe = Datasource( + with Datasource( config.version_str, config.tag, config.name, @@ -250,106 +220,92 @@ async def data( sql_context=context, basic_config=basic_config, accounts=configs.accounts, - ) - if config.datasource.file_type == "delta": - pre_filter = await get_partitions( - realdataframe, realdataframe.execution_uri, params, config - ) - else: - pre_filter = None - pre_filter_p2 = ( - get_filter_dict( - params.model_dump(exclude_unset=True), - config.params if config.params else [], - None, + ) as realdataframe: + if params: + pre_filter, _ = get_filters( + params.model_dump(exclude_unset=True), + config.params if config.params else [], + None, + ) + else: + pre_filter = None + df = realdataframe.get_df(filters=pre_filter) + df_cols = df.columns() + expr = get_params_filter_expr( # this supports all kinds of filters, while the prefilter only supports equality + context, + df_cols, + config, + params, ) - if params - else None - ) - if pre_filter or pre_filter_p2: - pre_filter = pre_filter or [] - if pre_filter_p2: - for k, v in pre_filter_p2.items(): - if k not in pre_filter: - pre_filter.append((k, "=", v)) - - df = await realdataframe.get_df(filters=pre_filter) - df_cols = await _async(df.columns()) - expr = get_params_filter_expr( # this supports all kinds of filters, while the prefilter only supports equality - context, - df_cols, - config, - params, - ) - new_query = df.query_builder() - new_query = new_query.where(expr) if expr is not None else new_query - columns = exclude_cols(df_cols, basic_config) - if select: - columns = [ - c for c in columns if c in split_csv(select) - ] # split , is a bit naive, we might want to support real CSV with quotes here - if config.datasource.exclude and len(config.datasource.exclude) > 0: - columns = [c for c in columns if c not in config.datasource.exclude] - if config.datasource.sortby: - for s in config.datasource.sortby: - new_query = cast(ex.Select, new_query).order_by( - ex.column(s.by, quoted=True).desc() - if s.direction and s.direction.lower() == "desc" - else ex.column(s.by, quoted=True), - copy=False, + new_query = df.query_builder() + new_query = new_query.where(expr) if expr is not None else new_query + columns = exclude_cols(df_cols, basic_config) + if select: + columns = [ + c for c in columns if c in split_csv(select) + ] # split , is a bit naive, we might want to support real CSV with quotes here + if config.datasource.exclude and len(config.datasource.exclude) > 0: + columns = [c for c in columns if c not in config.datasource.exclude] + if config.datasource.sortby: + for s in config.datasource.sortby: + new_query = cast(ex.Select, new_query).order_by( + ex.column(s.by, quoted=True).desc() + if s.direction and s.direction.lower() == "desc" + else ex.column(s.by, quoted=True), + copy=False, + ) + if has_complex and format in ["csv", "excel", "scsv", "csv4excel"]: + jsonify_complex = True + if jsonify_complex: + base_schema = df.arrow_schema() + + complex_cols = [c for c in columns if is_complex_type(base_schema, c)] + + new_query = context.jsonify_complex(new_query, complex_cols, columns) + else: + new_query = new_query.select( + *[ex.column(c, quoted=True) for c in columns], append=False ) - if has_complex and format in ["csv", "excel", "scsv", "csv4excel"]: - jsonify_complex = True - if jsonify_complex: - base_schema = await df.arrow_schema() - - complex_cols = [c for c in columns if is_complex_type(base_schema, c)] - new_query = context.jsonify_complex(new_query, complex_cols, columns) - else: - new_query = new_query.select( - *[ex.column(c, quoted=True) for c in columns], append=False + if distinct: + assert len(columns) <= 3 # reduce complexity here + new_query = cast(ex.Select, new_query).distinct() + + if not (limit == -1 and config.allow_get_all_pages): + limit = (1000 if limit == -1 else limit) or 1000 + new_query = new_query.limit(limit) + if offset: + new_query.offset(offset, copy=False) + + new_query = handle_search_request( + context, + config, + params, + basic_config, + source_view=realdataframe.tablename, + query=new_query, ) - - if distinct: - assert len(columns) <= 3 # reduce complexity here - new_query = cast(ex.Select, new_query).distinct() - - if not (limit == -1 and config.allow_get_all_pages): - limit = (1000 if limit == -1 else limit) or 1000 - new_query = new_query.limit(limit) - if offset: - new_query.offset(offset, copy=False) - - new_query = handle_search_request( - context, - config, - params, - basic_config, - source_view=realdataframe.tablename, - query=new_query, - ) - new_query = handle_nearby_request( - context, - config, - params, - basic_config, - source_view=realdataframe.tablename, - query=new_query, - ) - logger.debug(f"Query: {get_sql(new_query, dialect='duckdb')}") - - try: - return await create_response( - request.url, - request.query_params, - format or request.headers["Accept"], - context=context, - sql=new_query, - basic_config=basic_config, - close_context=True, - charset=request.query_params.get("$encoding"), + new_query = handle_nearby_request( + context, + config, + params, + basic_config, + source_view=realdataframe.tablename, + query=new_query, ) - except Exception as err: - logger.error("Error in creating response", exc_info=err) - raise HTTPException(status_code=500) + logger.debug(f"Query: {get_sql(new_query, dialect='duckdb')}") + + try: + return await create_response( + request.url, + request.query_params, + format or request.headers["Accept"], + context=context, + sql=new_query, + basic_config=basic_config, + close_context=True, + charset=request.query_params.get("$encoding"), + ) + except Exception as err: + logger.error("Error in creating response", exc_info=err) + raise HTTPException(status_code=500) diff --git a/bmsdna/lakeapi/endpoint/sql_endpoint.py b/bmsdna/lakeapi/endpoint/sql_endpoint.py index e3438eb..c85af64 100644 --- a/bmsdna/lakeapi/endpoint/sql_endpoint.py +++ b/bmsdna/lakeapi/endpoint/sql_endpoint.py @@ -21,7 +21,7 @@ def init_duck_con( ): for cfg in configs: assert cfg.datasource is not None - df = Datasource( + with Datasource( cfg.version_str, cfg.tag, cfg.name, @@ -29,19 +29,18 @@ def init_duck_con( sql_context=con, accounts=configs.accounts, basic_config=basic_config, - ) - - if cfg.engine != "odbc" and df.file_exists(): - try: - con.register_datasource( - df.unique_table_name, - df.tablename, - df.execution_uri, - df.config.file_type, - None, - ) - except (FileTypeNotSupportedError, FileNotFoundError): - logger.warning(f"Cannot query {df.tablename}") + ) as df: + if cfg.engine != "odbc" and df.file_exists(): + try: + con.register_datasource( + df.unique_table_name, + df.tablename, + df.execution_uri, + df.config.file_type, + None, + ) + except (FileTypeNotSupportedError, FileNotFoundError): + logger.warning(f"Cannot query {df.tablename}") def _get_sql_context( @@ -92,7 +91,7 @@ async def get_sql_tables( ), ): con = _get_sql_context(engine, basic_config, configs) - return await (await con.list_tables()).to_pylist() + return await con.list_tables().to_pylist() @router.post( "/api/sql", diff --git a/bmsdna/lakeapi/standalone/__init__.py b/bmsdna/lakeapi/standalone/__init__.py index f6dd14c..f02bc83 100644 --- a/bmsdna/lakeapi/standalone/__init__.py +++ b/bmsdna/lakeapi/standalone/__init__.py @@ -6,10 +6,7 @@ def run_fastapi(): app = FastAPI() - async def _init(): - await init_lakeapi(app, use_basic_auth=True) - - asyncio.run(_init()) + init_lakeapi(app, use_basic_auth=True) @app.get("/") async def root(req: Request): diff --git a/config_test.yml b/config_test.yml index 085006d..97de11a 100644 --- a/config_test.yml +++ b/config_test.yml @@ -504,6 +504,15 @@ tables: file_type: parquet account: "test_account" + - name: blob_test_copied + tag: blobb + api_method: get + datasource: + uri: "az://testlake/td/faker.parquet" + file_type: parquet + copy_local: true + account: "test_account" + - name: "*" tag: blobb api_method: get diff --git a/pyproject.toml b/pyproject.toml index 4f950ee..24d3127 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,6 +1,6 @@ [project] name = "bmsdna-lakeapi" -version = "0.24.3" +version = "0.25.0" description = "" authors = [{ name = "DWH Team", email = "you@example.com" }] dependencies = [ @@ -19,7 +19,7 @@ dependencies = [ "pandas >=2.1.0,<3", "fsspec >=2024.2.0,<2025", "adlfs >=2024.2.0,<2025", - "deltalake2db >=1.0.7", + "deltalake2db >=1.1.1", ] requires-python = "~=3.10" diff --git a/start_test_instance.py b/start_test_instance.py index a55f2fd..cd5a078 100644 --- a/start_test_instance.py +++ b/start_test_instance.py @@ -11,14 +11,14 @@ app = fastapi.FastAPI() -async def init(): +def init(): def_cfg = bmsdna.lakeapi.get_default_config() # Get default startup config cfg = dataclasses.replace( def_cfg, enable_sql_endpoint=True, data_path="tests/data" ) # Use dataclasses.replace to set the properties you want - sti = await bmsdna.lakeapi.init_lakeapi( + sti = bmsdna.lakeapi.init_lakeapi( app, True, cfg, "config_test.yml" ) # Enable it. The first parameter is the FastAPI instance, the 2nd one is the basic config and the third one the config of the tables -asyncio.create_task(init()) +init() diff --git a/tests/test_blobb.py b/tests/test_blobb.py index c52c8e7..35f4ea0 100644 --- a/tests/test_blobb.py +++ b/tests/test_blobb.py @@ -14,6 +14,23 @@ def test_parquet(engine, client: TestClient): assert len(fakedt) == 50 +@pytest.mark.parametrize("engine", engines) +def test_parquet_copied(engine, client: TestClient): + response = client.get( + f"/api/v1/blobb/blob_test_copied?format=json&limit=50&$engine={engine}", + ) + assert response.status_code == 200 + fakedt = response.json() + assert len(fakedt) == 50 + + response2 = client.get( + f"/api/v1/blobb/blob_test_copied?format=json&limit=50&$engine={engine}", + ) + assert response2.status_code == 200 + fakedt2 = response2.json() + assert len(fakedt2) == 50 + + @pytest.mark.parametrize("engine", engines) def test_delta(engine, client: TestClient): response = client.get( diff --git a/tests/utils.py b/tests/utils.py index c5712a5..373f90d 100644 --- a/tests/utils.py +++ b/tests/utils.py @@ -1,4 +1,3 @@ -import asyncio from fastapi import FastAPI, Request, status import dataclasses from faker import Faker @@ -22,10 +21,7 @@ def get_app(default_engine="duckdb"): default_engine=default_engine, ) - async def _init(): - await bmsdna.lakeapi.init_lakeapi(app, True, cfg, "config_test.yml") - - sti = asyncio.run(_init()) + bmsdna.lakeapi.init_lakeapi(app, True, cfg, "config_test.yml") @app.exception_handler(RequestValidationError) async def validation_exception_handler( diff --git a/uv.lock b/uv.lock index 32e73fa..21543b4 100644 --- a/uv.lock +++ b/uv.lock @@ -432,7 +432,7 @@ wheels = [ [[package]] name = "bmsdna-lakeapi" -version = "0.24.2" +version = "0.24.3" source = { editable = "." } dependencies = [ { name = "adlfs" }, @@ -501,7 +501,7 @@ requires-dist = [ { name = "argon2-cffi", marker = "extra == 'auth'" }, { name = "arrow-odbc", specifier = ">=5.0.0" }, { name = "arrow-odbc", marker = "extra == 'odbc'" }, - { name = "deltalake2db", specifier = ">=1.0.7" }, + { name = "deltalake2db", specifier = ">=1.1.1" }, { name = "duckdb", specifier = ">=1.1.0,<2" }, { name = "expandvars", specifier = ">=0.12.0" }, { name = "fastapi", specifier = ">=0.110.0" }, @@ -912,15 +912,15 @@ wheels = [ [[package]] name = "deltalake2db" -version = "1.0.7" +version = "1.1.1" source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "polars" }, { name = "sqlglot" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/14/49/4ad39cc14b0cfdf609b3d66ce08893a2c54bb933a5256c90345cec39f124/deltalake2db-1.0.7.tar.gz", hash = "sha256:ee064597d19d0d95c325134fc466357eac854ed708b0a7ce5623c1ac7ce2f70b", size = 244965, upload-time = "2025-09-23T12:32:30.782Z" } +sdist = { url = "https://files.pythonhosted.org/packages/8b/c3/0dac5043069d8b49ad846c97ba5ce3443eaf6bac248ec27f3e0a8bbdbbe9/deltalake2db-1.1.1.tar.gz", hash = "sha256:9aa1a62b3f1b0f2d79e8a028d0e43beb5bf3ea89ad599e14dbd622affe6b6994", size = 246492, upload-time = "2025-09-24T16:11:37.51Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/b2/cb/1d3e1ba7b642c0f8956a4f3bd62124a274e6a1ff9f8590804d7054ee80ad/deltalake2db-1.0.7-py3-none-any.whl", hash = "sha256:49335e5d56ab3efaad256cd0250d2c8a4907289bcb0f481bb485111e82998e53", size = 17865, upload-time = "2025-09-23T12:32:28.674Z" }, + { url = "https://files.pythonhosted.org/packages/73/8a/a2da1c17994cd3d543303ea16aa3c3b33889f24b6856df0b3718cd1de0c4/deltalake2db-1.1.1-py3-none-any.whl", hash = "sha256:6f17caef7630d3b217d2898f30984755d824c6aaf44d6b7b0fd6daa9c91198a4", size = 19137, upload-time = "2025-09-24T16:11:36.128Z" }, ] [[package]]