adapta 2.11.9__py3-none-any.whl → 3.5.13__py3-none-any.whl

adapta/storage/query_enabled_store/_models.py

@@ -2,7 +2,7 @@
  Query Enabled Store Connection interface.
 """
 
-#  Copyright (c) 2023-2024. ECCO Sneaks & Data
+#  Copyright (c) 2023-2026. ECCO Data & AI and other project contributors.
 #
 #  Licensed under the Apache License, Version 2.0 (the "License");
 #  you may not use this file except in compliance with the License.
@@ -21,12 +21,13 @@ import re
 from abc import ABC, abstractmethod
 from enum import Enum
 from pydoc import locate
-from typing import TypeVar, Generic, Type, Iterator, Union, final, Optional
-
-from pandas import DataFrame
+from typing import TypeVar, Generic, final
+from collections.abc import Iterator
 
 from adapta.storage.models.base import DataPath
 from adapta.storage.models.filter_expression import Expression
+from adapta.storage.models.enum import QueryEnabledStoreOptions
+from adapta.utils.metaframe import MetaFrame
 
 TCredential = TypeVar("TCredential")  # pylint: disable=C0103
 TSettings = TypeVar("TSettings")  # pylint: disable=C0103
@@ -42,6 +43,8 @@ class BundledQueryEnabledStores(Enum):
 
     DELTA = "adapta.storage.query_enabled_store.DeltaQueryEnabledStore"
     ASTRA = "adapta.storage.query_enabled_store.AstraQueryEnabledStore"
+    LOCAL = "adapta.storage.query_enabled_store.LocalQueryEnabledStore"
+    TRINO = "adapta.storage.query_enabled_store.TrinoQueryEnabledStore"
 
 
 BUNDLED_STORES = {store.name: store.value for store in BundledQueryEnabledStores}
@@ -84,16 +87,21 @@ class QueryEnabledStore(Generic[TCredential, TSettings], ABC):
 
     @abstractmethod
     def _apply_filter(
-        self, path: DataPath, filter_expression: Expression, columns: list[str]
-    ) -> Union[DataFrame, Iterator[DataFrame]]:
+        self,
+        path: DataPath,
+        filter_expression: Expression,
+        columns: list[str],
+        options: dict[QueryEnabledStoreOptions, any] | None = None,
+        limit: int | None = None,
+    ) -> MetaFrame | Iterator[MetaFrame]:
         """
-        Applies the provided filter expression to this Store and returns the result in a pandas DataFrame
+        Applies the provided filter expression to this Store and returns the result in a MetaFrame
         """
 
     @abstractmethod
-    def _apply_query(self, query: str) -> Union[DataFrame, Iterator[DataFrame]]:
+    def _apply_query(self, query: str) -> MetaFrame | Iterator[MetaFrame]:
         """
-        Applies a plaintext query to this Store and returns the result in a pandas DataFrame
+        Applies a plaintext query to this Store and returns the result in a MetaFrame
         """
 
     @classmethod
@@ -117,7 +125,7 @@ class QueryEnabledStore(Generic[TCredential, TSettings], ABC):
         :param: lazy_init: Whether to set this instance QES for querying eagerly or lazily.
         """
 
-        def get_qes_class(name: str) -> Type[QueryEnabledStore[TCredential, TSettings]]:
+        def get_qes_class(name: str) -> type[QueryEnabledStore[TCredential, TSettings]]:
             return locate(BUNDLED_STORES.get(name, name))
 
         class_name, _, _ = re.findall(re.compile(CONNECTION_STRING_REGEX), connection_string)[0]
@@ -138,8 +146,10 @@ class QueryConfigurationBuilder:
     def __init__(self, store: QueryEnabledStore, path: DataPath):
         self._store = store
         self._path = path
-        self._filter_expression: Optional[Expression] = None
+        self._filter_expression: Expression | None = None
         self._columns: list[str] = []
+        self._options: dict[QueryEnabledStoreOptions, any] = {}
+        self._limit = None
 
     def filter(self, filter_expression: Expression) -> "QueryConfigurationBuilder":
         """
@@ -157,10 +167,29 @@ class QueryConfigurationBuilder:
         self._columns = list(columns)
         return self
 
-    def read(self) -> Union[DataFrame, Iterator[DataFrame]]:
+    def add_options(self, option_key: QueryEnabledStoreOptions, option_value: any) -> "QueryConfigurationBuilder":
+        """
+        Use the provided options when querying the underlying storage.
+        """
+
+        self._options[option_key] = option_value
+        return self
+
+    def limit(self, limit: int | None) -> "QueryConfigurationBuilder":
+        """
+        Limit the number of results returned by the underlying store.
+        """
+        self._limit = limit
+        return self
+
+    def read(self) -> MetaFrame | Iterator[MetaFrame]:
         """
         Execute the query on the underlying store.
         """
         return self._store._apply_filter(
-            path=self._path, filter_expression=self._filter_expression, columns=self._columns
+            path=self._path,
+            filter_expression=self._filter_expression,
+            columns=self._columns,
+            options=self._options,
+            limit=self._limit,
         )

adapta/storage/query_enabled_store/_qes_astra.py

@@ -4,19 +4,23 @@
 import os
 import re
 from dataclasses import dataclass
-from typing import final, Optional, Union, Iterator
-
-from pandas import DataFrame
+from typing import final
+from collections.abc import Iterator
 
 from dataclasses_json import DataClassJsonMixin
 
 from adapta._version import __version__
-from adapta.storage.distributed_object_store.datastax_astra.astra_client import AstraClient
+from adapta.storage.distributed_object_store.v3.datastax_astra import AstraClient
 from adapta.storage.models.astra import AstraPath
 from adapta.storage.models.base import DataPath
 from adapta.storage.models.filter_expression import Expression
 
-from adapta.storage.query_enabled_store._models import QueryEnabledStore, CONNECTION_STRING_REGEX
+from adapta.storage.query_enabled_store._models import (
+    QueryEnabledStore,
+    CONNECTION_STRING_REGEX,
+)
+from adapta.storage.models.enum import QueryEnabledStoreOptions
+from adapta.utils.metaframe import MetaFrame
 
 
 @dataclass
@@ -25,9 +29,9 @@ class AstraCredential(DataClassJsonMixin):
     Astra DB credential helper for QES.
     """
 
-    secret_connection_bundle_bytes: Optional[str] = None
-    client_id: Optional[str] = None
-    client_secret: Optional[str] = None
+    secret_connection_bundle_bytes: str | None = None
+    client_id: str | None = None
+    client_secret: str | None = None
 
     def __post_init__(self):
         self.secret_connection_bundle_bytes = self.secret_connection_bundle_bytes or os.getenv(
@@ -36,7 +40,7 @@ class AstraCredential(DataClassJsonMixin):
         self.client_id = self.client_id or os.getenv("PROTEUS__ASTRA_CLIENT_ID")
         self.client_secret = self.client_secret or os.getenv("PROTEUS__ASTRA_CLIENT_SECRET")
 
-        if not all([self.secret_connection_bundle_bytes, self.client_secret, self.client_secret]):
+        if not all([self.secret_connection_bundle_bytes, self.client_id, self.client_secret]):
             raise RuntimeError(
                 "Authentication information provided is insufficient. Please verify you are supplying bundle bytes, client id and secret either via connection string or via environment variables."
             )
@@ -48,8 +52,8 @@ class AstraSettings(DataClassJsonMixin):
     Astra DB connection settings for QES.
     """
 
-    client_name: Optional[str] = None
-    keyspace: Optional[str] = None
+    client_name: str | None = None
+    keyspace: str | None = None
 
     def __post_init__(self):
         self.client_name = self.client_name or f"Adapta Client {__version__}"
@@ -79,8 +83,13 @@ class AstraQueryEnabledStore(QueryEnabledStore[AstraCredential, AstraSettings]):
             self._astra_client.connect()
 
     def _apply_filter(
-        self, path: DataPath, filter_expression: Expression, columns: list[str]
-    ) -> Union[DataFrame, Iterator[DataFrame]]:
+        self,
+        path: DataPath,
+        filter_expression: Expression,
+        columns: list[str],
+        options: dict[QueryEnabledStoreOptions, any] | None = None,
+        limit: int | None = 10000,
+    ) -> MetaFrame | Iterator[MetaFrame]:
         assert isinstance(path, AstraPath)
         astra_path: AstraPath = path
         if self._lazy:
@@ -92,6 +101,8 @@ class AstraQueryEnabledStore(QueryEnabledStore[AstraCredential, AstraSettings]):
                     table_name=astra_path.table,
                     select_columns=columns,
                     num_threads=-1,  # auto-infer, see method documentation
+                    options=options,
+                    limit=limit,
                 )
 
         return self._astra_client.filter_entities(
@@ -101,9 +112,11 @@ class AstraQueryEnabledStore(QueryEnabledStore[AstraCredential, AstraSettings]):
             table_name=astra_path.table,
             select_columns=columns,
             num_threads=-1,  # auto-infer, see method documentation
+            options=options,
+            limit=limit,
         )
 
-    def _apply_query(self, query: str) -> Union[DataFrame, Iterator[DataFrame]]:
+    def _apply_query(self, query: str) -> MetaFrame | Iterator[MetaFrame]:
         if self._lazy:
             with self._astra_client as astra_client:
                 return astra_client.get_entities_raw(query)

adapta/storage/query_enabled_store/_qes_delta.py

@@ -4,15 +4,21 @@
 import re
 from dataclasses import dataclass
 from pydoc import locate
-from typing import final, Union, Iterator
+from typing import final
+from collections.abc import Iterator
 
-from pandas import DataFrame
 from dataclasses_json import DataClassJsonMixin
 
-from adapta.storage.delta_lake import load
+from adapta.security.clients import AuthenticationClient
+from adapta.storage.delta_lake.v3 import load
 from adapta.storage.models.base import DataPath
 from adapta.storage.models.filter_expression import Expression
-from adapta.storage.query_enabled_store._models import QueryEnabledStore, CONNECTION_STRING_REGEX
+from adapta.storage.query_enabled_store._models import (
+    QueryEnabledStore,
+    CONNECTION_STRING_REGEX,
+)
+from adapta.storage.models.enum import QueryEnabledStoreOptions
+from adapta.utils.metaframe import MetaFrame
 
 
 @dataclass
@@ -22,16 +28,25 @@ class DeltaCredential(DataClassJsonMixin):
     """
 
     auth_client_class: str
+    auth_client_credentials_class: str | None = None
+
+    auth_client: AuthenticationClient | None = None
+    auth_client_credentials: type | None = None
 
     def __post_init__(self):
         if not self.auth_client_class:
             raise ValueError("Authentication plugin class name not provided but is required")
 
-        if locate(self.auth_client_class) is None:
+        self.auth_client = locate(self.auth_client_class)
+
+        if self.auth_client is None:
             raise ModuleNotFoundError(
                 "Authentication plugin class name cannot be loaded. Please check the spelling and make sure your application can resolve the import"
             )
 
+        if self.auth_client_credentials_class:
+            self.auth_client_credentials = locate(self.auth_client_credentials_class)
+
 
 @dataclass
 class DeltaSettings(DataClassJsonMixin):
@@ -57,14 +72,21 @@ class DeltaQueryEnabledStore(QueryEnabledStore[DeltaCredential, DeltaSettings]):
         return cls(credentials=DeltaCredential.from_json(credentials), settings=DeltaSettings.from_json(settings))
 
     def _apply_filter(
-        self, path: DataPath, filter_expression: Expression, columns: list[str]
-    ) -> Union[DataFrame, Iterator[DataFrame]]:
+        self,
+        path: DataPath,
+        filter_expression: Expression,
+        columns: list[str],
+        options: dict[QueryEnabledStoreOptions, any] | None = None,
+        limit: int | None = None,
+    ) -> MetaFrame | Iterator[MetaFrame]:
         return load(
-            auth_client=locate(self.credentials.auth_client_class)(),
+            auth_client=self.credentials.auth_client(credentials=self.credentials.auth_client_credentials()),
             path=path,
             row_filter=filter_expression,
-            columns=columns,
+            columns=columns if columns else None,
+            limit=limit,
+            timeout=options.get(QueryEnabledStoreOptions.TIMEOUT, None),
         )
 
-    def _apply_query(self, query: str) -> Union[DataFrame, Iterator[DataFrame]]:
+    def _apply_query(self, query: str) -> MetaFrame | Iterator[MetaFrame]:
         raise NotImplementedError("Text queries are not supported by Delta QES")

adapta/storage/query_enabled_store/_qes_local.py

@@ -0,0 +1,81 @@
+"""Local Query Enabled Store (QES) for reading local files."""
+import re
+from dataclasses import dataclass
+from typing import final
+from collections.abc import Iterator
+
+from dataclasses_json import DataClassJsonMixin
+from pyarrow.parquet import read_table
+
+from adapta.storage.models import DataPath
+from adapta.storage.models.filter_expression import Expression, compile_expression, ArrowFilterExpression
+from adapta.storage.query_enabled_store._models import (
+    QueryEnabledStore,
+    CONNECTION_STRING_REGEX,
+)
+from adapta.storage.models.enum import QueryEnabledStoreOptions
+from adapta.utils.metaframe import MetaFrame
+
+
+@dataclass
+class LocalCredential(DataClassJsonMixin):
+    """
+    Local credential helper for QES.
+    No authentication is required for local files.
+    """
+
+
+@dataclass
+class LocalSettings(DataClassJsonMixin):
+    """
+    Settings for local QES
+    """
+
+
+@final
+class LocalQueryEnabledStore(QueryEnabledStore[LocalCredential, LocalSettings]):
+    """
+    QES Client for local file reads (e.g., Parquet) using PyArrow.
+    """
+
+    def close(self) -> None:
+        pass
+
+    @classmethod
+    def _from_connection_string(
+        cls, connection_string: str, lazy_init: bool = False
+    ) -> "QueryEnabledStore[LocalCredential, LocalSettings]":
+        """
+        Parses a connection string for local files.
+        """
+        _, credentials, settings = re.findall(re.compile(CONNECTION_STRING_REGEX), connection_string)[0]
+        return cls(credentials=LocalCredential.from_json(credentials), settings=LocalSettings.from_json(settings))
+
+    def _apply_filter(
+        self,
+        path: DataPath,
+        filter_expression: Expression,
+        columns: list[str],
+        options: dict[QueryEnabledStoreOptions, any] | None = None,
+        limit: int = None,
+    ) -> MetaFrame | Iterator[MetaFrame]:
+        """
+        Applies a filter to a local file
+        """
+        row_filter = compile_expression(filter_expression, ArrowFilterExpression) if filter_expression else None
+
+        pyarrow_table = read_table(
+            path.path,
+            columns=columns if columns else None,
+            filters=row_filter,
+        )
+
+        return MetaFrame.from_arrow(
+            data=pyarrow_table,
+        )
+
+    def _apply_query(self, query: str) -> MetaFrame | Iterator[MetaFrame]:
+        """
+        Local QES does not natively support SQL-like queries.
+        """
+        raise NotImplementedError("Text queries are currently not supported by Local QES")

adapta/storage/query_enabled_store/_qes_trino.py

@@ -0,0 +1,133 @@
+"""
+ QES implementations for Trino.
+"""
+import os
+import re
+from dataclasses import dataclass
+from typing import final
+from collections.abc import Iterator
+
+from dataclasses_json import DataClassJsonMixin
+
+from adapta.storage.database.v3.trino_sql import TrinoClient
+from adapta.storage.models import TrinoPath
+from adapta.storage.models.filter_expression import (
+    Expression,
+    compile_expression,
+    TrinoFilterExpression,
+)
+from adapta.storage.query_enabled_store._models import (
+    QueryEnabledStore,
+    CONNECTION_STRING_REGEX,
+)
+from adapta.storage.models.enum import QueryEnabledStoreOptions
+from adapta.utils.metaframe import MetaFrame, concat
+
+
+@dataclass
+class TrinoCredential(DataClassJsonMixin):
+    """
+    Trino credential helper for QES.
+
+    Trino credentials can either be provided via the oauth2_username or via the following environment variables, which
+    is handled inside the TrinoClient:
+    - (ADAPTA__TRINO_USERNAME, ADAPTA__TRINO_PASSWORD)
+
+    Currently, we don't support the credentials_provider option of the TrinoClient.
+    """
+
+    oauth2_username: str | None = None
+
+    def __post_init__(self):
+        self.oauth2_username = self.oauth2_username or os.getenv("ADAPTA__TRINO_OAUTH2_USERNAME")
+
+
+@dataclass
+class TrinoSettings(DataClassJsonMixin):
+    """
+    Trino connection settings for QES.
+    """
+
+    host: str | None = None
+    port: int | None = None
+
+    def __post_init__(self):
+        self.host = self.host or os.getenv("PROTEUS__TRINO_HOST")
+        if not self.host:
+            raise RuntimeError(
+                "Trino host not provided. Please provide it via connection string or via environment variable PROTEUS__TRINO_HOST."
+            )
+        self.port = self.port or int(os.getenv("PROTEUS__TRINO_PORT", "443"))
+
+
+@final
+class TrinoQueryEnabledStore(QueryEnabledStore[TrinoCredential, TrinoSettings]):
+    """
+    QES Client for Trino queries.
+    """
+
+    def close(self) -> None:
+        pass
+
+    def __init__(self, credentials: TrinoCredential, settings: TrinoSettings):
+        super().__init__(credentials, settings)
+        self._trino_client = TrinoClient(
+            host=self.settings.host,
+            port=self.settings.port,
+            oauth2_username=self.credentials.oauth2_username,
+        )
+
+    def _apply_filter(
+        self,
+        path: TrinoPath,
+        filter_expression: Expression,
+        columns: list[str],
+        options: dict[QueryEnabledStoreOptions, any] | None = None,
+        limit: int | None = None,
+    ) -> MetaFrame | Iterator[MetaFrame]:
+        query = self._build_query(query=path.query, filter_expression=filter_expression, columns=columns, limit=limit)
+
+        with self._trino_client as trino_client:
+            if QueryEnabledStoreOptions.BATCH_SIZE in options:
+                data = concat(
+                    trino_client.query(
+                        query=query,
+                        batch_size=options[QueryEnabledStoreOptions.BATCH_SIZE],
+                    )
+                )
+            else:
+                data = concat(trino_client.query(query=query))
+
+        return data
+
+    def _apply_query(self, query: str) -> MetaFrame | Iterator[MetaFrame]:
+        raise NotImplementedError("Text queries are not supported by Trino QES")
+
+    @classmethod
+    def _from_connection_string(
+        cls, connection_string: str, lazy_init: bool = False
+    ) -> "QueryEnabledStore[TrinoCredential, TrinoSettings]":
+        _, credentials, settings = re.findall(re.compile(CONNECTION_STRING_REGEX), connection_string)[0]
+        return cls(
+            credentials=TrinoCredential.from_json(credentials),
+            settings=TrinoSettings.from_json(settings),
+        )
+
+    @staticmethod
+    def _build_query(query: str, filter_expression: Expression, columns: list[str], limit: int | None) -> str:
+        """
+        Build the final query by applying the filter expression, selected columns, and limit to the base query.
+        """
+
+        if filter_expression or columns or limit:
+            columns_to_select = ", ".join(columns) if columns else "*"
+            query = f"SELECT {columns_to_select} FROM ({query})"
+
+            if filter_expression:
+                compiled_expression = compile_expression(expression=filter_expression, target=TrinoFilterExpression)
+                query = f"{query} WHERE {compiled_expression}"
+
+            if limit:
+                query = f"{query} LIMIT {limit}"
+
+        return query

adapta/storage/secrets/__init__.py

@@ -2,7 +2,7 @@
  Import index.
 """
 
-#  Copyright (c) 2023-2024. ECCO Sneaks & Data
+#  Copyright (c) 2023-2026. ECCO Data & AI and other project contributors.
 #
 #  Licensed under the Apache License, Version 2.0 (the "License");
 #  you may not use this file except in compliance with the License.

adapta/storage/secrets/_base.py

@@ -2,7 +2,7 @@
  Abstraction for secret storage operations.
 """
 
-#  Copyright (c) 2023-2024. ECCO Sneaks & Data
+#  Copyright (c) 2023-2026. ECCO Data & AI and other project contributors.
 #
 #  Licensed under the Apache License, Version 2.0 (the "License");
 #  you may not use this file except in compliance with the License.
@@ -18,7 +18,8 @@
 #
 
 from abc import ABC, abstractmethod
-from typing import Union, Dict, Iterable, Any
+from typing import Any
+from collections.abc import Iterable
 
 from adapta.security.clients import AuthenticationClient
 
@@ -32,7 +33,7 @@ class SecretStorageClient(ABC):
         self._base_client = base_client
 
     @abstractmethod
-    def read_secret(self, storage_name: str, secret_name: str) -> Union[bytes, str, Dict[str, str]]:
+    def read_secret(self, storage_name: str, secret_name: str) -> bytes | str | dict[str, str]:
         """
           Reads a secret from the specified storage.
 
@@ -46,7 +47,7 @@ class SecretStorageClient(ABC):
         self,
         storage_name: str,
         secret_name: str,
-        secret_value: Union[str, Dict[str, str]],
+        secret_value: str | dict[str, str],
         b64_encode=False,
     ) -> None:
         """

adapta/storage/secrets/azure_secret_client.py

@@ -1,7 +1,7 @@
 """
  Azure Secret Storage Client (KeyVault).
 """
-#  Copyright (c) 2023-2024. ECCO Sneaks & Data
+#  Copyright (c) 2023-2026. ECCO Data & AI and other project contributors.
 #
 #  Licensed under the Apache License, Version 2.0 (the "License");
 #  you may not use this file except in compliance with the License.
@@ -17,7 +17,6 @@
 #
 
 import base64
-from typing import Union, Dict
 
 from azure.keyvault.secrets import SecretClient
 
@@ -43,14 +42,14 @@ class AzureSecretStorageClient(SecretStorageClient):
         kv_uri = f"https://{keyvault}.vault.azure.net"
         return SecretClient(kv_uri, self._base_client.get_credentials())
 
-    def read_secret(self, storage_name: str, secret_name: str) -> Union[bytes, str, Dict[str, str]]:
+    def read_secret(self, storage_name: str, secret_name: str) -> bytes | str | dict[str, str]:
         return self._get_keyvault(storage_name).get_secret(secret_name).value
 
     def create_secret(
         self,
         storage_name: str,
         secret_name: str,
-        secret_value: Union[str, Dict[str, str]],
+        secret_value: str | dict[str, str],
         b64_encode=False,
     ) -> None:
         if not isinstance(secret_value, str):

adapta/storage/secrets/hashicorp_vault_secret_storage_client.py

@@ -1,7 +1,7 @@
 """
  Hashicorp Vault Secret storage client
 """
-#  Copyright (c) 2023-2024. ECCO Sneaks & Data
+#  Copyright (c) 2023-2026. ECCO Data & AI and other project contributors.
 #
 #  Licensed under the Apache License, Version 2.0 (the "License");
 #  you may not use this file except in compliance with the License.
@@ -16,7 +16,7 @@
 #  limitations under the License.
 #
 
-from typing import Union, Dict, Iterable
+from collections.abc import Iterable
 
 import hvac
 
@@ -41,7 +41,7 @@ class HashicorpSecretStorageClient(SecretStorageClient):
         self.client = hvac.Client(self._base_client.vault_address, self._access_token)
         self._role = role
 
-    def read_secret(self, storage_name: str, secret_name: str) -> Union[bytes, str, Dict[str, str]]:
+    def read_secret(self, storage_name: str, secret_name: str) -> bytes | str | dict[str, str]:
         secret = self.client.secrets.kv.v2.read_secret_version(path=secret_name)
         return secret["data"]["data"]
 
@@ -49,10 +49,10 @@ class HashicorpSecretStorageClient(SecretStorageClient):
         self,
         storage_name: str,
         secret_name: str,
-        secret_value: Union[str, Dict[str, str]],
+        secret_value: str | dict[str, str],
         b64_encode=False,
     ) -> None:
-        if not isinstance(secret_value, Dict):
+        if not isinstance(secret_value, dict):
             raise ValueError(
                 f"Only Dict secret type supported in HashicorpSecretStorageClient but was: {type(secret_value)}"
             )