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

digitalhub/stores/local/store.py

@@ -4,19 +4,11 @@ import shutil
 from pathlib import Path
 from typing import Any
 
-from digitalhub.readers.api import get_reader_by_object
-from digitalhub.stores._base.store import Store, StoreConfig
+from digitalhub.readers.data.api import get_reader_by_object
+from digitalhub.stores._base.store import Store
 from digitalhub.utils.exceptions import StoreError
 from digitalhub.utils.file_utils import get_file_info_from_local
-
-
-class LocalStoreConfig(StoreConfig):
-    """
-    Local store configuration class.
-    """
-
-    path: str
-    """Local path."""
+from digitalhub.utils.types import SourcesOrListOfSources
 
 
 class LocalStore(Store):
@@ -25,9 +17,8 @@ class LocalStore(Store):
     artifacts on local filesystem based storage.
     """
 
-    def __init__(self, name: str, store_type: str, config: LocalStoreConfig) -> None:
-        super().__init__(name, store_type)
-        self.config = config
+    def __init__(self, config: dict | None = None) -> None:
+        super().__init__()
 
     ##############################
     # I/O methods
@@ -61,7 +52,7 @@ class LocalStore(Store):
         """
         raise StoreError("Local store does not support download.")
 
-    def upload(self, src: str | list[str], dst: str) -> list[tuple[str, str]]:
+    def upload(self, src: SourcesOrListOfSources, dst: str) -> list[tuple[str, str]]:
         """
         Upload an artifact to storage.
 
@@ -92,6 +83,106 @@ class LocalStore(Store):
         """
         return [get_file_info_from_local(p) for p in paths]
 
+    ##############################
+    # Datastore methods
+    ##############################
+
+    def read_df(
+        self,
+        path: SourcesOrListOfSources,
+        file_format: str | None = None,
+        engine: str | None = None,
+        **kwargs,
+    ) -> Any:
+        """
+        Read DataFrame from path.
+
+        Parameters
+        ----------
+        path : SourcesOrListOfSources
+            Path(s) to read DataFrame from.
+        file_format : str
+            Extension of the file.
+        engine : str
+            Dataframe engine (pandas, polars, etc.).
+        **kwargs : dict
+            Keyword arguments.
+
+        Returns
+        -------
+        Any
+            DataFrame.
+        """
+        reader = self._get_reader(engine)
+
+        dfs = []
+        if isinstance(path, list):
+            for p in path:
+                file_format = self._get_extension(file_format, p)
+                dfs.append(reader.read_df(p, file_format, **kwargs))
+        elif Path(path).is_dir():
+            import glob
+
+            paths = glob.glob(f"{path}/*")
+            for p in paths:
+                file_format = self._get_extension(file_format, p)
+                dfs.append(reader.read_df(p, file_format, **kwargs))
+        else:
+            file_format = self._get_extension(file_format, path)
+            dfs.append(reader.read_df(path, file_format, **kwargs))
+
+        if len(dfs) == 1:
+            return dfs[0]
+
+        return reader.concat_dfs(dfs)
+
+    def query(
+        self,
+        query: str,
+        engine: str | None = None,
+    ) -> Any:
+        """
+        Query data from database.
+
+        Parameters
+        ----------
+        query : str
+            The query to execute.
+        engine : str
+            Dataframe engine (pandas, polars, etc.).
+
+        Returns
+        -------
+        Any
+            DataFrame.
+        """
+        raise StoreError("Local store does not support query.")
+
+    def write_df(self, df: Any, dst: str, extension: str | None = None, **kwargs) -> str:
+        """
+        Method to write a dataframe to a file. Kwargs are passed to df.to_parquet().
+        If destination is not provided, the dataframe is written to the default
+        store path with generated name.
+
+        Parameters
+        ----------
+        df : Any
+            The dataframe to write.
+        dst : str
+            The destination of the dataframe.
+        **kwargs : dict
+            Keyword arguments.
+
+        Returns
+        -------
+        str
+            Path of written dataframe.
+        """
+        self._check_local_dst(dst)
+        reader = get_reader_by_object(df)
+        reader.write_df(df, dst, extension=extension, **kwargs)
+        return dst
+
     ##############################
     # Private I/O methods
     ##############################
@@ -193,77 +284,3 @@ class LocalStore(Store):
             dst = dst / src
         self._build_path(dst)
         return dst
-
-    ##############################
-    # Datastore methods
-    ##############################
-
-    def write_df(self, df: Any, dst: str, extension: str | None = None, **kwargs) -> str:
-        """
-        Method to write a dataframe to a file. Kwargs are passed to df.to_parquet().
-        If destination is not provided, the dataframe is written to the default
-        store path with generated name.
-
-        Parameters
-        ----------
-        df : Any
-            The dataframe to write.
-        dst : str
-            The destination of the dataframe.
-        **kwargs : dict
-            Keyword arguments.
-
-        Returns
-        -------
-        str
-            Path of written dataframe.
-        """
-        self.store._check_local_dst(dst)
-        self._validate_extension(Path(dst).suffix.removeprefix("."))
-
-        # Write dataframe
-        reader = get_reader_by_object(df)
-        reader.write_df(df, dst, extension=extension, **kwargs)
-
-        return dst
-
-    ##############################
-    # Helper methods
-    ##############################
-
-    @staticmethod
-    def is_partition_or_dir(path: str) -> bool:
-        """
-        Check if path is a directory or a partition.
-
-        Parameters
-        ----------
-        path : str
-            The path to check.
-
-        Returns
-        -------
-        bool
-        """
-        return Path(path).is_dir()
-
-    @staticmethod
-    def build_object_path(root: str, paths: str | list[str]) -> list[str]:
-        """
-        Method to build object path.
-
-        Parameters
-        ----------
-        root : str
-            The root of the object path.
-        paths : str | list[str]
-            The path to build.
-
-        Returns
-        -------
-        list[str]
-            Returns the path of the object.
-        """
-        if isinstance(paths, str):
-            paths = [paths]
-        return [str(Path(root) / path) for path in paths]

digitalhub/stores/remote/store.py

@@ -5,14 +5,9 @@ from typing import Any
 
 import requests
 
-from digitalhub.stores._base.store import Store, StoreConfig
+from digitalhub.stores._base.store import Store
 from digitalhub.utils.exceptions import StoreError
-
-
-class RemoteStoreConfig(StoreConfig):
-    """
-    Remote store configuration class.
-    """
+from digitalhub.utils.types import SourcesOrListOfSources
 
 
 class RemoteStore(Store):
@@ -21,9 +16,8 @@ class RemoteStore(Store):
     artifacts from remote HTTP based storage.
     """
 
-    def __init__(self, name: str, store_type: str, config: RemoteStoreConfig) -> None:
-        super().__init__(name, store_type)
-        self.config = config
+    def __init__(self, config: dict | None = None) -> None:
+        super().__init__()
 
     ##############################
     # I/O methods
@@ -69,7 +63,7 @@ class RemoteStore(Store):
 
         return self._download_file(root, dst, overwrite)
 
-    def upload(self, src: str | list[str], dst: str) -> list[tuple[str, str]]:
+    def upload(self, src: SourcesOrListOfSources, dst: str) -> list[tuple[str, str]]:
         """
         Upload an artifact to storage.
 
@@ -104,6 +98,58 @@ class RemoteStore(Store):
     # Datastore methods
     ##############################
 
+    def read_df(
+        self,
+        path: SourcesOrListOfSources,
+        file_format: str | None = None,
+        engine: str | None = None,
+        **kwargs,
+    ) -> Any:
+        """
+        Read DataFrame from path.
+
+        Parameters
+        ----------
+        path : SourcesOrListOfSources
+            Path(s) to read DataFrame from.
+        file_format : str
+            Extension of the file.
+        engine : str
+            Dataframe engine (pandas, polars, etc.).
+        **kwargs : dict
+            Keyword arguments.
+
+        Returns
+        -------
+        Any
+            DataFrame.
+        """
+        reader = self._get_reader(engine)
+        extension = self._head_extension(path, file_format)
+        return reader.read_df(path, extension, **kwargs)
+
+    def query(
+        self,
+        query: str,
+        engine: str | None = None,
+    ) -> Any:
+        """
+        Query data from database.
+
+        Parameters
+        ----------
+        query : str
+            The query to execute.
+        engine : str
+            Dataframe engine (pandas, polars, etc.).
+
+        Returns
+        -------
+        Any
+            DataFrame.
+        """
+        raise StoreError("Remote store does not support query.")
+
     def write_df(self, df: Any, dst: str, extension: str | None = None, **kwargs) -> str:
         """
         Method to write a dataframe to a file. Note that this method is not implemented
@@ -167,3 +213,32 @@ class RemoteStore(Store):
                 for chunk in r.iter_content(chunk_size=8192):
                     f.write(chunk)
         return str(dst)
+
+    def _head_extension(self, url: str, file_format: str | None = None) -> str:
+        """
+        Method to get the extension of a file from a given url.
+
+        Parameters
+        ----------
+        url : str
+            The url of the file to get the extension.
+        file_format : str
+            The file format to check.
+
+        Returns
+        -------
+        str
+            File extension.
+        """
+        if file_format is not None:
+            return file_format
+        try:
+            r = requests.head(url, timeout=60)
+            r.raise_for_status()
+            content_type = r.headers["content-type"]
+            if "text" in content_type:
+                return "csv"
+            else:
+                raise ValueError("Content type not supported.")
+        except Exception as e:
+            raise e

digitalhub/stores/s3/configurator.py

@@ -0,0 +1,108 @@
+from __future__ import annotations
+
+from botocore.config import Config
+
+from digitalhub.configurator.configurator import configurator
+from digitalhub.stores.s3.enums import S3StoreEnv
+from digitalhub.stores.s3.models import S3StoreConfig
+from digitalhub.utils.exceptions import StoreError
+
+
+class S3StoreConfigurator:
+    """
+    Configure the store by getting the credentials from user
+    provided config or from environment.
+    """
+
+    def __init__(self, config: dict | None = None) -> None:
+        self.configure(config)
+
+    ##############################
+    # Configuration methods
+    ##############################
+
+    def configure(self, config: dict | None = None) -> None:
+        """
+        Configure the store by getting the credentials from user
+        provided config or from environment.
+
+        Parameters
+        ----------
+        config : dict
+            Configuration dictionary.
+
+        Returns
+        -------
+        None
+        """
+        # Validate config
+        if config is None:
+            self._get_config()
+        else:
+            config = S3StoreConfig(**config)
+            for pair in [
+                (S3StoreEnv.ENDPOINT_URL.value, config.endpoint),
+                (S3StoreEnv.ACCESS_KEY_ID.value, config.access_key),
+                (S3StoreEnv.SECRET_ACCESS_KEY.value, config.secret_key),
+                (S3StoreEnv.SESSION_TOKEN.value, config.session_token),
+                (S3StoreEnv.BUCKET_NAME.value, config.bucket_name),
+                (S3StoreEnv.REGION.value, config.region),
+                (S3StoreEnv.SIGNATURE_VERSION.value, config.signature_version),
+            ]:
+                configurator.set_credential(*pair)
+
+    def get_boto3_client_config(self) -> dict:
+        """
+        Get S3 credentials (access key, secret key,
+        session token and other config).
+
+        Returns
+        -------
+        dict
+            The credentials.
+        """
+        creds = configurator.get_all_credentials()
+        try:
+            return {
+                "endpoint_url": creds[S3StoreEnv.ENDPOINT_URL.value],
+                "aws_access_key_id": creds[S3StoreEnv.ACCESS_KEY_ID.value],
+                "aws_secret_access_key": creds[S3StoreEnv.SECRET_ACCESS_KEY.value],
+                "aws_session_token": creds[S3StoreEnv.SESSION_TOKEN.value],
+                "config": Config(
+                    region_name=creds[S3StoreEnv.REGION.value],
+                    signature_version=creds[S3StoreEnv.SIGNATURE_VERSION.value],
+                ),
+            }
+        except KeyError as e:
+            raise StoreError(f"Missing credentials for S3 store. {str(e)}")
+
+    def _get_config(self) -> None:
+        """
+        Get the store configuration.
+
+        Returns
+        -------
+        None
+        """
+        required_vars = [
+            S3StoreEnv.ENDPOINT_URL,
+            S3StoreEnv.ACCESS_KEY_ID,
+            S3StoreEnv.SECRET_ACCESS_KEY,
+            S3StoreEnv.BUCKET_NAME,
+        ]
+        optional_vars = [S3StoreEnv.REGION, S3StoreEnv.SIGNATURE_VERSION, S3StoreEnv.SESSION_TOKEN]
+
+        # Load required environment variables
+        credentials = {var.value: configurator.load_var(var.value) for var in required_vars}
+
+        # Check for missing required credentials
+        missing_vars = [key for key, value in credentials.items() if value is None]
+        if missing_vars:
+            raise StoreError(f"Missing credentials for S3 store: {', '.join(missing_vars)}")
+
+        # Load optional environment variables
+        credentials.update({var.value: configurator.load_var(var.value) for var in optional_vars})
+
+        # Set credentials
+        for key, value in credentials.items():
+            configurator.set_credential(key, value)

digitalhub/stores/s3/enums.py

@@ -0,0 +1,17 @@
+from __future__ import annotations
+
+from enum import Enum
+
+
+class S3StoreEnv(Enum):
+    """
+    S3Store environment
+    """
+
+    ENDPOINT_URL = "S3_ENDPOINT_URL"
+    ACCESS_KEY_ID = "AWS_ACCESS_KEY_ID"
+    SECRET_ACCESS_KEY = "AWS_SECRET_ACCESS_KEY"
+    SESSION_TOKEN = "AWS_SESSION_TOKEN"
+    BUCKET_NAME = "S3_BUCKET"
+    REGION = "S3_REGION"
+    SIGNATURE_VERSION = "S3_SIGNATURE_VERSION"

digitalhub/stores/s3/models.py

@@ -0,0 +1,21 @@
+from __future__ import annotations
+
+from pydantic import BaseModel
+
+
+class S3StoreConfig(BaseModel):
+    """
+    S3 store configuration class.
+    """
+
+    endpoint_url: str
+    """S3 endpoint URL."""
+
+    aws_access_key_id: str
+    """AWS access key ID."""
+
+    aws_secret_access_key: str
+    """AWS secret access key."""
+
+    bucket_name: str
+    """S3 bucket name."""