digitalhub 0.9.2__py3-none-any.whl → 0.10.0__py3-none-any.whl

digitalhub/stores/_base/store.py

@@ -1,15 +1,18 @@
 from __future__ import annotations
 
+import typing
 from abc import abstractmethod
 from pathlib import Path
 from tempfile import mkdtemp
 from typing import Any
 
-from pydantic import BaseModel, ConfigDict
-
-from digitalhub.readers.api import get_reader_by_engine
+from digitalhub.readers.data.api import get_reader_by_engine
 from digitalhub.utils.exceptions import StoreError
-from digitalhub.utils.uri_utils import SchemeCategory, has_local_scheme
+from digitalhub.utils.types import SourcesOrListOfSources
+from digitalhub.utils.uri_utils import has_local_scheme
+
+if typing.TYPE_CHECKING:
+    from digitalhub.readers.data._base.reader import DataframeReader
 
 
 class Store:
@@ -17,24 +20,6 @@ class Store:
     Store abstract class.
     """
 
-    def __init__(self, name: str, store_type: str) -> None:
-        """
-        Constructor.
-
-        Parameters
-        ----------
-        name : str
-            Store name.
-        store_type : str
-            Store type. Used to choose the right store implementation.
-
-        Returns
-        -------
-        None
-        """
-        self.name = name
-        self.type = store_type
-
     ##############################
     # I/O methods
     ##############################
@@ -52,7 +37,7 @@ class Store:
         """
 
     @abstractmethod
-    def upload(self, src: str | list[str], dst: str) -> list[tuple[str, str]]:
+    def upload(self, src: SourcesOrListOfSources, dst: str) -> list[tuple[str, str]]:
         """
         Method to upload artifact to storage.
         """
@@ -72,39 +57,32 @@ class Store:
     ##############################
 
     @abstractmethod
-    def write_df(self, df: Any, dst: str, extension: str | None = None, **kwargs) -> str:
-        """
-        Write DataFrame as parquet or csv.
-        """
-
     def read_df(
         self,
-        path: str | list[str],
-        extension: str,
+        path: SourcesOrListOfSources,
+        file_format: str | None = None,
         engine: str | None = None,
         **kwargs,
     ) -> Any:
         """
         Read DataFrame from path.
+        """
 
-        Parameters
-        ----------
-        path : str | list[str]
-            Path(s) to read DataFrame from.
-        extension : str
-            Extension of the file.
-        engine : str
-            Dataframe engine (pandas, polars, etc.).
-        **kwargs : dict
-            Keyword arguments.
+    @abstractmethod
+    def query(
+        self,
+        query: str,
+        engine: str | None = None,
+    ) -> Any:
+        """
+        Query data from database.
+        """
 
-        Returns
-        -------
-        Any
-            DataFrame.
+    @abstractmethod
+    def write_df(self, df: Any, dst: str, extension: str | None = None, **kwargs) -> str:
+        """
+        Write DataFrame as parquet or csv.
         """
-        reader = get_reader_by_engine(engine)
-        return reader.read_df(path, extension, **kwargs)
 
     ##############################
     # Helpers methods
@@ -182,7 +160,7 @@ class Store:
 
         Parameters
         ----------
-        path : str
+        path : str | Path
             The path to build.
 
         Returns
@@ -208,30 +186,42 @@ class Store:
         tmpdir = mkdtemp()
         return Path(tmpdir)
 
+    @staticmethod
+    def _get_reader(engine: str | None = None) -> DataframeReader:
+        """
+        Get Dataframe reader.
 
-class StoreConfig(BaseModel):
-    """
-    Store configuration base class.
-    """
-
-    model_config = ConfigDict(use_enum_values=True)
-
-
-class StoreParameters(BaseModel):
-    """
-    Store configuration class.
-    """
-
-    model_config = ConfigDict(use_enum_values=True)
+        Parameters
+        ----------
+        engine : str
+            Dataframe engine (pandas, polars, etc.).
 
-    name: str
-    """Store id."""
+        Returns
+        -------
+        Any
+            Reader object.
+        """
+        return get_reader_by_engine(engine)
 
-    type: SchemeCategory
-    """Store type to instantiate."""
+    @staticmethod
+    def _get_extension(extension: str | None = None, path: str | None = None) -> str:
+        """
+        Get extension from path.
 
-    config: StoreConfig = None
-    """Configuration for the store."""
+        Parameters
+        ----------
+        extension : str
+            The extension to get.
+        path : str
+            The path to get the extension from.
 
-    is_default: bool = False
-    """Flag to determine if the store is the default one."""
+        Returns
+        -------
+        str
+            The extension.
+        """
+        if extension is not None:
+            return extension
+        if path is not None:
+            return Path(path).suffix.removeprefix(".")
+        raise ValueError("Extension or path must be provided.")

digitalhub/stores/api.py

@@ -5,50 +5,25 @@ import typing
 from digitalhub.stores.builder import store_builder
 
 if typing.TYPE_CHECKING:
-    from digitalhub.stores._base.store import Store, StoreParameters
+    from digitalhub.stores._base.store import Store
 
 
-def set_store(store_cfg: StoreParameters) -> None:
+def get_store(project: str, uri: str, config: dict | None = None) -> Store:
     """
-    Set a new store instance with the given configuration.
-
-    Parameters
-    ----------
-    store_cfg : StoreParameters
-        Store configuration.
-
-    Returns
-    -------
-    None
-    """
-    store_builder.build(store_cfg)
-
-
-def get_store(uri: str) -> Store:
-    """
-    Get store instance by uri.
+    Get store instance by URI.
 
     Parameters
     ---------
+    project : str
+        Project name.
     uri : str
         URI to parse.
+    config : dict
+        Store configuration.
 
     Returns
     -------
     Store
         Store instance.
     """
-    return store_builder.get(uri)
-
-
-def get_default_store() -> Store:
-    """
-    Get the default store instance. The default store is the one that
-    can persist artifacts and dataitems.
-
-    Returns
-    -------
-    Store
-        Default store instance.
-    """
-    return store_builder.default()
+    return store_builder.get(uri, config)

digitalhub/stores/builder.py

@@ -1,28 +1,41 @@
 from __future__ import annotations
 
-import os
 import typing
 
-from pydantic import ValidationError
-
-from digitalhub.stores._base.store import StoreParameters
-from digitalhub.stores.local.store import LocalStore, LocalStoreConfig
-from digitalhub.stores.remote.store import RemoteStore, RemoteStoreConfig
-from digitalhub.stores.s3.store import S3Store, S3StoreConfig
-from digitalhub.stores.sql.store import SqlStore, SQLStoreConfig
-from digitalhub.utils.exceptions import StoreError
+from digitalhub.configurator.api import get_current_env
+from digitalhub.stores.local.store import LocalStore
+from digitalhub.stores.remote.store import RemoteStore
+from digitalhub.stores.s3.store import S3Store
+from digitalhub.stores.sql.store import SqlStore
 from digitalhub.utils.uri_utils import SchemeCategory, map_uri_scheme
 
 if typing.TYPE_CHECKING:
     from digitalhub.stores._base.store import Store
 
 
-REGISTRY_STORES = {
-    "local": LocalStore,
-    "s3": S3Store,
-    "remote": RemoteStore,
-    "sql": SqlStore,
-}
+def _get_class_from_type(type: str) -> Store:
+    """
+    Get a store class from its type.
+
+    Parameters
+    ----------
+    type : str
+        Store type.
+
+    Returns
+    -------
+    Store
+        The store class.
+    """
+    if type == SchemeCategory.LOCAL.value:
+        return LocalStore
+    if type == SchemeCategory.S3.value:
+        return S3Store
+    if type == SchemeCategory.REMOTE.value:
+        return RemoteStore
+    if type == SchemeCategory.SQL.value:
+        return SqlStore
+    raise ValueError(f"Unknown store type: {type}")
 
 
 class StoreBuilder:
@@ -32,27 +45,28 @@ class StoreBuilder:
 
     def __init__(self) -> None:
         self._instances: dict[str, Store] = {}
-        self._default: Store | None = None
-        self._def_scheme = "s3"
 
-    def build(self, store_cfg: StoreParameters) -> None:
+    def build(self, store_type: str, config: dict | None = None) -> None:
         """
         Build a store instance and register it.
-        It overrides any existing instance.
 
         Parameters
         ----------
-        store_cfg : StoreParameters
+        store_type : str
+            Store type.
+        config : dict
             Store configuration.
 
         Returns
         -------
         None
         """
-        scheme = map_uri_scheme(store_cfg.type)
-        self._instances[scheme] = self.build_store(store_cfg)
+        env = get_current_env()
+        if env not in self._instances:
+            self._instances[env] = {}
+        self._instances[env][store_type] = _get_class_from_type(store_type)(config)
 
-    def get(self, uri: str) -> Store:
+    def get(self, uri: str, config: dict | None = None) -> Store:
         """
         Get a store instance by URI.
 
@@ -60,152 +74,21 @@ class StoreBuilder:
         ----------
         uri : str
             URI to parse.
-
-        Returns
-        -------
-        Store
-            The store instance.
-        """
-        scheme = map_uri_scheme(uri)
-        if scheme not in self._instances:
-            store_cfg = get_env_store_config(scheme)
-            self._instances[scheme] = self.build_store(store_cfg)
-        return self._instances[scheme]
-
-    def default(self) -> Store:
-        """
-        Get the default store instance.
-
-        Returns
-        -------
-        Store
-            The default store instance.
-
-        Raises
-        ------
-        StoreError
-            If no default store is set.
-        """
-        if self._default is None:
-            store_cfg = get_env_store_config(self._def_scheme)
-            self._default = self.build_store(store_cfg)
-        return self._default
-
-    def build_store(self, cfg: StoreParameters) -> Store:
-        """
-        Build a store instance.
-
-        Parameters
-        ----------
-        cfg : StoreParameters
+        config : dict
             Store configuration.
 
         Returns
         -------
         Store
             The store instance.
-
-        Raises
-        ------
-        NotImplementedError
-            If the store type is not implemented.
         """
+        env = get_current_env()
+        store_type = map_uri_scheme(uri)
         try:
-            obj = REGISTRY_STORES[cfg.type](cfg.name, cfg.type, cfg.config)
-            if cfg.is_default and self._default is not None:
-                raise StoreError("Only one default store!")
-            return obj
-        except KeyError as e:
-            raise NotImplementedError from e
-
-    @staticmethod
-    def _check_config(config: StoreParameters | dict) -> StoreParameters:
-        """
-        Check the store configuration validity.
-
-        Parameters
-        ----------
-        config : StoreParameters | dict
-            The store configuration.
-
-        Returns
-        -------
-        StoreParameters
-            The store configuration.
-
-        Raises
-        ------
-        TypeError
-            If the config parameter is not a StoreParameters instance or a well-formed dictionary.
-        """
-        if not isinstance(config, StoreParameters):
-            try:
-                return StoreParameters(**config)
-            except TypeError as e:
-                raise StoreError("Invalid store configuration type.") from e
-            except ValidationError as e:
-                raise StoreError("Malformed store configuration parameters.") from e
-        return config
-
-
-def get_env_store_config(scheme: str) -> StoreParameters:
-    """
-    Get a store configuration from the environment.
-
-    Parameters
-    ----------
-    scheme : str
-        URI scheme.
-
-    Returns
-    -------
-    StoreParameters
-        The store configuration based on the scheme.
-
-    Raises
-    ------
-    ValueError
-        If the scheme is not supported.
-    """
-    if scheme == SchemeCategory.S3.value:
-        return StoreParameters(
-            name="s3",
-            type="s3",
-            config=S3StoreConfig(
-                endpoint_url=os.getenv("S3_ENDPOINT_URL"),  # type: ignore
-                aws_access_key_id=os.getenv("AWS_ACCESS_KEY_ID"),  # type: ignore
-                aws_secret_access_key=os.getenv("AWS_SECRET_ACCESS_KEY"),  # type: ignore
-                bucket_name=os.getenv("S3_BUCKET_NAME"),  # type: ignore
-            ),
-        )
-    if scheme == SchemeCategory.SQL.value:
-        return StoreParameters(
-            name="sql",
-            type="sql",
-            config=SQLStoreConfig(
-                host=os.getenv("POSTGRES_HOST"),  # type: ignore
-                port=os.getenv("POSTGRES_PORT"),  # type: ignore
-                user=os.getenv("POSTGRES_USER"),  # type: ignore
-                password=os.getenv("POSTGRES_PASSWORD"),  # type: ignore
-                database=os.getenv("POSTGRES_DATABASE"),  # type: ignore
-                pg_schema=os.getenv("POSTGRES_SCHEMA"),  # type: ignore
-            ),
-        )
-    if scheme == SchemeCategory.REMOTE.value:
-        return StoreParameters(
-            name="remote",
-            type="remote",
-            config=RemoteStoreConfig(),
-        )
-    if scheme == SchemeCategory.LOCAL.value:
-        return StoreParameters(
-            name="local",
-            type="local",
-            config=LocalStoreConfig(
-                path="tempsdk",
-            ),
-        )
-    raise ValueError(f"Unsupported scheme {scheme}")
+            return self._instances[env][store_type]
+        except KeyError:
+            self.build(store_type, config)
+            return self._instances[env][store_type]
 
 
 store_builder = StoreBuilder()