Repository: SneaksAndData/adapta (language: python)

Issue:
[FEATURE] Generalize the `load` method
It would be nice if we could generalize the `load` method such that it accepts a socket and infers what driver API to use to load the data.

This could reduce the code duplication in our repos that load data. Especially in [SneaksAndData/omni-channel-crystal-solver](https://github.com/SneaksAndData/omni-channel-crystal-solver) and  [SneaksAndData/return-optimizer-crystal-solver](https://github.com/SneaksAndData/return-optimizer-crystal-solver).

The generalized version might not be a function, but could be OOP-based with a factory class to create a data loader instance.


Discussion / hints:
After a very productive offline discussion with @matt035343, @Jupetus and @lucasgml  we'll implement the following:

1. All storages will be generalized into `StorageConnection` class which accepts a connection string in a format to be developed, similar to data paths in `DataSocket`
2. We'll add the following capability to that class:

```
  con = StorageConnection.from_connection_string('....')
  socket = DataSocket(...) # implemented already
  filters = (FilterField("col_a") == 1) & (FilterField("col_b") == "abc") # implemented already
  
  my_data = con.load(socket, filters) # pd.DataFrame or ps.DataFrame
```

This will allow to use Astra, Azure, S3, SQL and any other storage data source interchangeably without adding boilerplate and reduce the number of environment variables we map to algorithm runner containers + unify all storage loading interfaces.
Implements https://github.com/SneaksAndData/adapta/issues/291
Replaces #201 

Fix patch:
diff --git a/adapta/storage/delta_lake/README.md b/adapta/storage/delta_lake/README.md
index 0d554a42..051c8a21 100644
--- a/adapta/storage/delta_lake/README.md
+++ b/adapta/storage/delta_lake/README.md
@@ -7,47 +7,16 @@ Supported API:
 - read and filter a delta table without loading all rows in memory
 
 ## Example usage for Azure Datalake Gen2
-1. Create table
-```cassandraql
-create table delta_lake_test_filtering using delta as
-select
-    "some-value" as my_column,
-     "123" as my_other_column,
-      456L as another_column
-union all
-select
-    "some-value" as my_column,
-     "another-value" as my_other_column,
-     0L as another_column
-union all
-select
-    "something1" as my_column,
-     "else" as my_other_column,
-     1L as another_column
-union all
-select
-    "something1" as my_column,
-     "nonexistent" as my_other_column,
-     2L as another_column
-union all
-select
-    "something1" as my_column,
-     "nonexistent1" as my_other_column,
-     123L as another_column
-```
-2. Prepare connection and load
+Prepare connection and load
+
 ```python
 import os
-import datetime
 from adapta.security.clients import AzureClient
 from adapta.storage.models.azure import AdlsGen2Path
-from adapta.storage.models.hive import HivePath
-from adapta.storage.delta_lake import load, load_cached
-from adapta.logs import SemanticLogger
-from adapta.storage.cache.redis_cache import RedisCache
-from adapta.storage.models.filter_expression import FilterField
+from adapta.storage.delta_lake import load
 
-azure_client = AzureClient(subscription_id='6c5538ce-b24a-4e2a-877f-979ad71287ff')
+os.environ["PROTEUS__USE_AZURE_CREDENTIAL"] = "1"
+azure_client = AzureClient()
 adls_path = AdlsGen2Path.from_hdfs_path('abfss://container@account.dfs.core.windows.net/path/to/my/table')
 
 # get Iterable[pandas.DataFrame]
@@ -56,6 +25,8 @@ batches = load(azure_client, adls_path, batch_size=1000)
 ## Using the Filtering API.
 1. Create generic filter expressions
 ```python
+from adapta.storage.models.filter_expression import FilterField
+
 simple_filter = FilterField("my_column") == "some-value"
 combined_filter = (FilterField("my_column") == "some-value") & (FilterField("other_column") == "another-value")
 combined_filter_with_collection = (FilterField("my_column") == "something1") & (FilterField("other_column").isin(['else', 'nonexistent']))
diff --git a/adapta/storage/models/astra.py b/adapta/storage/models/astra.py
index 7d9ac750..8b49a559 100644
--- a/adapta/storage/models/astra.py
+++ b/adapta/storage/models/astra.py
@@ -17,10 +17,14 @@
 #
 
 from dataclasses import dataclass
+from pydoc import locate
+from typing import Optional, TypeVar, Type
 from urllib.parse import urlparse
 
 from adapta.storage.models.base import DataPath, DataProtocols
 
+TModel = TypeVar("TModel")  # pylint: disable=C0103
+
 
 @dataclass
 class AstraPath(DataPath):
@@ -40,6 +44,7 @@ def from_uri(cls, url: str) -> "DataPath":
 
     keyspace: str
     table: str
+    model_class_name: Optional[str] = None
     protocol: str = DataProtocols.ASTRA.value
 
     @classmethod
@@ -53,3 +58,12 @@ def to_hdfs_path(self) -> str:
 
     def to_delta_rs_path(self) -> str:
         raise NotImplementedError
+
+    def model_class(self) -> Optional[Type[TModel]]:
+        """
+        Locates and returns model class name.
+        """
+        if self.model_class_name:
+            return locate(self.model_class_name)
+
+        return None
diff --git a/adapta/storage/query_enabled_store/README.md b/adapta/storage/query_enabled_store/README.md
new file mode 100644
index 00000000..830f5ce2
--- /dev/null
+++ b/adapta/storage/query_enabled_store/README.md
@@ -0,0 +1,43 @@
+## Query-Enabled Store Interface
+
+This API greatly simplifies data reads for applications that require transparent swap of a data backend, for example Delta -> Cassandra, Delta (Azure) -> Delta (AWS) or any other combination. Implementing a QES client allows data scientists and engineers to write applications that can be reconfigured to run on any supported store by simply updating the application's configuration.
+
+QES enables abstraction of *both* storage connection and filtering API, which relies on adapta's `FilterExpression` API.
+
+### Usage
+
+You can create a QES instance by providing a QES-formatted connection string. Two formats are supported:
+- `"qes://class=<full class name>;plaintext_credentials=<json object>;settings=<json object>"`
+- `"qes://class=<stable store alias>;plaintext_credentials=<json object>;settings=<json object>"`
+
+Some QES implementations are bundled with `adapta`: `ASTRA`, `DELTA`. In case you have implemented your own QES in this repository, you can add it to `BundledQes` enum, or it can be loaded dynamically from client code, if your application is able to resolve the QES class import.
+
+Example of a connection string, for a delta table stored on Azure:
+```python
+# using dynamic import
+conn = "qes://engine=adapta.storage.query_enabled.DeltaQes;plaintext_credentials={\"auth_client_class\":\"adapta.security.clients.AzureClient\"};settings={}"
+# using bundled QES
+conn_bundled = "qes://engine=DELTA;plaintext_credentials={\"auth_client_class\":\"adapta.security.clients.AzureClient\"};settings={}"
+```
+Now, initialize a QES object from that connection and read some data:
+
+```python
+import os
+from adapta.storage.query_enabled_store import QueryEnabledStore
+from adapta.storage.models.azure import AdlsGen2Path
+from adapta.storage.models.filter_expression import FilterField
+
+# use implicit auth for Azure (for simplicity - check AzureClient documentation for more options)
+os.environ["PROTEUS__USE_AZURE_CREDENTIAL"] = "1"
+adls_path = AdlsGen2Path.from_hdfs_path('abfss://container@account.dfs.core.windows.net/path/to/table')
+
+conn = "qes://engine=DELTA;plaintext_credentials={\"auth_client_class\":\"adapta.security.clients.AzureClient\"};settings={}"
+store = QueryEnabledStore.from_string(conn)
+data = store.open(adls_path).filter(FilterField("date_key") == "20230101").select("date_key", "year").read()
+
+print(data)
+
+#
+#    date_key  year
+# 0  20230101  2023
+```
diff --git a/adapta/storage/query_enabled_store/__init__.py b/adapta/storage/query_enabled_store/__init__.py
new file mode 100644
index 00000000..34cdfa04
--- /dev/null
+++ b/adapta/storage/query_enabled_store/__init__.py
@@ -0,0 +1,21 @@
+"""
+ Import index.
+"""
+#  Copyright (c) 2023. ECCO Sneaks & Data
+#
+#  Licensed under the Apache License, Version 2.0 (the "License");
+#  you may not use this file except in compliance with the License.
+#  You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+#  Unless required by applicable law or agreed to in writing, software
+#  distributed under the License is distributed on an "AS IS" BASIS,
+#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+#  See the License for the specific language governing permissions and
+#  limitations under the License.
+#
+
+from adapta.storage.query_enabled_store._models import *
+from adapta.storage.query_enabled_store._qes_astra import *
+from adapta.storage.query_enabled_store._qes_delta import *
diff --git a/adapta/storage/query_enabled_store/_models.py b/adapta/storage/query_enabled_store/_models.py
new file mode 100644
index 00000000..ea036eaa
--- /dev/null
+++ b/adapta/storage/query_enabled_store/_models.py
@@ -0,0 +1,152 @@
+"""
+ Query Enabled Store Connection interface.
+"""
+
+#  Copyright (c) 2023. ECCO Sneaks & Data
+#
+#  Licensed under the Apache License, Version 2.0 (the "License");
+#  you may not use this file except in compliance with the License.
+#  You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+#  Unless required by applicable law or agreed to in writing, software
+#  distributed under the License is distributed on an "AS IS" BASIS,
+#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+#  See the License for the specific language governing permissions and
+#  limitations under the License.
+#
+
+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
+
+import pandas
+
+from adapta.storage.models.base import DataPath
+from adapta.storage.models.filter_expression import Expression
+
+TCredential = TypeVar("TCredential")  # pylint: disable=C0103
+TSettings = TypeVar("TSettings")  # pylint: disable=C0103
+
+CONNECTION_STRING_REGEX = r"^qes:\/\/engine=(.*?);plaintext_credentials=(.*?);settings=(.*?)$"
+
+
+@final
+class BundledQueryEnabledStores(Enum):
+    """
+    QES Implementations aliases that are bundled with Adapta.
+    """
+
+    DELTA = "adapta.storage.query_enabled_store.DeltaQueryEnabledStore"
+    ASTRA = "adapta.storage.query_enabled_store.AstraQueryEnabledStore"
+
+
+BUNDLED_STORES = {store.name: store.value for store in BundledQueryEnabledStores}
+
+
+class QueryEnabledStore(Generic[TCredential, TSettings], ABC):
+    """
+    QES base class.
+    """
+
+    def __init__(self, credentials: TCredential, settings: TSettings):
+        self._credentials = credentials
+        self._settings = settings
+
+    @property
+    def credentials(self) -> TCredential:
+        """
+        Returns the credentials for this store type.
+        """
+        return self._credentials
+
+    @property
+    def settings(self) -> TSettings:
+        """
+        Returns the address to connect to, if applicable.
+        """
+        return self._settings
+
+    def open(self, path: DataPath) -> "QueryConfigurationBuilder":
+        """
+        Construct a reader object for QES to proxy to the underlying store implementation.
+        """
+        return QueryConfigurationBuilder(self, path)
+
+    @abstractmethod
+    def _apply_filter(
+        self, path: DataPath, filter_expression: Expression, columns: list[str]
+    ) -> Union[pandas.DataFrame, Iterator[pandas.DataFrame]]:
+        """
+        Applies the provided filter expression to this Store and returns the result in a pandas DataFrame
+        """
+
+    @abstractmethod
+    def _apply_query(self, query: str) -> Union[pandas.DataFrame, Iterator[pandas.DataFrame]]:
+        """
+        Applies a plaintext query to this Store and returns the result in a pandas DataFrame
+        """
+
+    @classmethod
+    @abstractmethod
+    def _from_connection_string(cls, connection_string: str) -> "QueryEnabledStore[TCredential, TSettings]":
+        """
+        Constructs the connection from a connection string
+        """
+
+    @staticmethod
+    def from_string(connection_string: str) -> "QueryEnabledStore[TCredential, TSettings]":
+        """
+        Constructs a concrete QES instance from a connection string.
+        """
+
+        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]
+        class_object = get_qes_class(class_name)
+        if class_object is None:
+            raise ModuleNotFoundError(
+                f"Cannot locate QES implementation: {class_name}. Please check the name for spelling errors and make sure your application can resolve the import"
+            )
+        return class_object._from_connection_string(connection_string)
+
+
+@final
+class QueryConfigurationBuilder:
+    """
+    Builder-pattern support for querying via QES.
+    """
+
+    def __init__(self, store: QueryEnabledStore, path: DataPath):
+        self._store = store
+        self._path = path
+        self._filter_expression: Optional[Expression] = None
+        self._columns: list[str] = []
+
+    def filter(self, filter_expression: Expression) -> "QueryConfigurationBuilder":
+        """
+        Use the provided expression when querying the underlying storage.
+        """
+        self._filter_expression = (
+            filter_expression if self._filter_expression is None else self._filter_expression and filter_expression
+        )
+        return self
+
+    def select(self, *columns: str) -> "QueryConfigurationBuilder":
+        """
+        Request the underlying store to project the result onto the provided column set.
+        """
+        self._columns = list(columns)
+        return self
+
+    def read(self) -> Union[pandas.DataFrame, Iterator[pandas.DataFrame]]:
+        """
+        Execute the query on the underlying store.
+        """
+        return self._store._apply_filter(
+            path=self._path, filter_expression=self._filter_expression, columns=self._columns
+        )
diff --git a/adapta/storage/query_enabled_store/_qes_astra.py b/adapta/storage/query_enabled_store/_qes_astra.py
new file mode 100644
index 00000000..23495998
--- /dev/null
+++ b/adapta/storage/query_enabled_store/_qes_astra.py
@@ -0,0 +1,98 @@
+"""
+ QES implementations for DataStax Astra.
+"""
+import os
+import re
+from dataclasses import dataclass
+from typing import final, Optional, Union, Iterator
+
+import pandas
+
+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.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
+
+
+@dataclass
+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
+
+    def __post_init__(self):
+        self.secret_connection_bundle_bytes = self.secret_connection_bundle_bytes or os.getenv(
+            "PROTEUS__ASTRA_BUNDLE_BYTES"
+        )
+        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]):
+            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."
+            )
+
+
+@dataclass
+class AstraSettings(DataClassJsonMixin):
+    """
+    Astra DB connection settings for QES.
+    """
+
+    client_name: Optional[str] = None
+    keyspace: Optional[str] = None
+
+    def __post_init__(self):
+        self.client_name = self.client_name or f"Adapta Client {__version__}"
+        self.keyspace = self.keyspace or os.getenv("PROTEUS__ASTRA_KEYSPACE")
+
+
+@final
+class AstraQueryEnabledStore(QueryEnabledStore[AstraCredential, AstraSettings]):
+    """
+    QES Client for Astra DB (Cassandra).
+    """
+
+    def _apply_filter(
+        self, path: DataPath, filter_expression: Expression, columns: list[str]
+    ) -> Union[pandas.DataFrame, Iterator[pandas.DataFrame]]:
+        assert isinstance(path, AstraPath)
+        astra_path: AstraPath = path
+
+        with AstraClient(
+            client_name=self.settings.client_name,
+            keyspace=astra_path.keyspace,
+            secure_connect_bundle_bytes=self.credentials.secret_connection_bundle_bytes,
+            client_id=self.credentials.client_id,
+            client_secret=self.credentials.client_secret,
+        ) as astra_client:
+            return astra_client.filter_entities(
+                model_class=astra_path.model_class(),
+                key_column_filter_values=filter_expression,
+                table_name=astra_path.table,
+                select_columns=columns,
+            )
+
+    def _apply_query(self, query: str) -> Union[pandas.DataFrame, Iterator[pandas.DataFrame]]:
+        with AstraClient(
+            client_name=self.settings.client_name,
+            keyspace=self.settings.keyspace,
+            secure_connect_bundle_bytes=self.credentials.secret_connection_bundle_bytes,
+            client_id=self.credentials.client_id,
+            client_secret=self.credentials.client_secret,
+        ) as astra_client:
+            return astra_client.get_entities_raw(query)
+
+    @classmethod
+    def _from_connection_string(cls, connection_string: str) -> "QueryEnabledStore[AstraCredential, AstraSettings]":
+        _, credentials, settings = re.findall(re.compile(CONNECTION_STRING_REGEX), connection_string)[0]
+        return cls(credentials=AstraCredential.from_json(credentials), settings=AstraSettings.from_json(settings))
diff --git a/adapta/storage/query_enabled_store/_qes_delta.py b/adapta/storage/query_enabled_store/_qes_delta.py
new file mode 100644
index 00000000..90259297
--- /dev/null
+++ b/adapta/storage/query_enabled_store/_qes_delta.py
@@ -0,0 +1,65 @@
+"""
+ QES implementations for delta-rs.
+"""
+import re
+from dataclasses import dataclass
+from pydoc import locate
+from typing import final, Union, Iterator
+
+import pandas
+from dataclasses_json import DataClassJsonMixin
+
+from adapta.storage.delta_lake 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
+
+
+@dataclass
+class DeltaCredential(DataClassJsonMixin):
+    """
+    Delta-rs credential helper for QES.
+    """
+
+    auth_client_class: str
+
+    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:
+            raise ModuleNotFoundError(
+                "Authentication plugin class name cannot be loaded. Please check the spelling and make sure your application can resolve the import"
+            )
+
+
+@dataclass
+class DeltaSettings(DataClassJsonMixin):
+    """
+    Delta QES has no additional settings.
+    """
+
+
+@final
+class DeltaQueryEnabledStore(QueryEnabledStore[DeltaCredential, DeltaSettings]):
+    """
+    QES Client for Delta Lake reads using delta-rs.
+    """
+
+    @classmethod
+    def _from_connection_string(cls, connection_string: str) -> "QueryEnabledStore[DeltaCredential, DeltaSettings]":
+        _, credentials, settings = re.findall(re.compile(CONNECTION_STRING_REGEX), connection_string)[0]
+        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[pandas.DataFrame, Iterator[pandas.DataFrame]]:
+        return load(
+            auth_client=locate(self.credentials.auth_client_class)(),
+            path=path,
+            row_filter=filter_expression,
+            columns=columns,
+        )
+
+    def _apply_query(self, query: str) -> Union[pandas.DataFrame, Iterator[pandas.DataFrame]]:
+        raise NotImplementedError("Text queries are not supported by Delta QES")


Test patch:
diff --git a/tests/test_query_enabled_store.py b/tests/test_query_enabled_store.py
new file mode 100644
index 00000000..9b943ecf
--- /dev/null
+++ b/tests/test_query_enabled_store.py
@@ -0,0 +1,57 @@
+from typing import Optional, Type, Union
+
+import pytest
+
+from adapta.storage.query_enabled_store import QueryEnabledStore, DeltaQueryEnabledStore, AstraQueryEnabledStore
+
+
+@pytest.mark.parametrize(
+    "connection_string, expected_store_type",
+    [
+        (
+            'qes://engine=adapta.storage.query_enabled_store.DeltaQueryEnabledStore;plaintext_credentials={"auth_client_class":"adapta.security.clients.AzureClient"};settings={}',
+            DeltaQueryEnabledStore,
+        ),
+        (
+            'qes://engine=DELTA;plaintext_credentials={"auth_client_class":"adapta.security.clients.AzureClient"};settings={}',
+            DeltaQueryEnabledStore,
+        ),
+        (
+            'qes://engine=DELTA;plaintext_credentials={"auth_client_class":"adapta.security.clients.AwsClient"};settings={}',
+            DeltaQueryEnabledStore,
+        ),
+        (
+            'qes://engine=DELTA;plaintext_credentials={"auth_client_class":"adapta.security.clients.TestClient"};settings={}',
+            ModuleNotFoundError,
+        ),
+        (
+            'qes://engine=ASTRA;plaintext_credentials={"secret_connection_bundle_bytes":"test", "client_id": "test", "client_secret": "test"};settings={"keyspace": "tmp"}',
+            AstraQueryEnabledStore,
+        ),
+        (
+            'qes://engine=adapta.storage.query_enabled_store.AstraQueryEnabledStore;plaintext_credentials={"secret_connection_bundle_bytes":"test", "client_id": "test", "client_secret": "test"};settings={"keyspace": "tmp"}',
+            AstraQueryEnabledStore,
+        ),
+        (
+            'qes://engine=adapta.storage.query_enabled_store.AstraQueryEnabledStore;plaintext_credentials={"secret_connection_bundle_bytes":"test", "client_id": "test", "client_secret": "test"};settings={"client_name": "test", "keyspace": "tmp"}',
+            AstraQueryEnabledStore,
+        ),
+        (
+            'qes://engine=adapta.storage.query_enabled_store.AstraQueryEnabledStore;plaintext_credentials={};settings={"client_name": "test", "keyspace": "tmp"}',
+            RuntimeError,
+        ),
+        (
+            'qes://engine=DELT;plaintext_credentials={"auth_client_class":"adapta.security.clients.AzureClient"};settings={}',
+            ModuleNotFoundError,
+        ),
+    ],
+)
+def test_query_store_instantiation(
+    connection_string: str, expected_store_type: Union[Type[QueryEnabledStore], Exception]
+):
+    try:
+        store = QueryEnabledStore.from_string(connection_string)
+
+        assert isinstance(store, expected_store_type)
+    except Exception as load_error:
+        assert isinstance(load_error, expected_store_type)