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

digitalhub/stores/s3/store.py

@@ -1,5 +1,6 @@
 from __future__ import annotations
 
+import typing
 from io import BytesIO
 from pathlib import Path
 from typing import Any, Type
@@ -9,43 +10,31 @@ import boto3
 import botocore.client  # pylint: disable=unused-import
 from botocore.exceptions import ClientError
 
-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.stores.s3.configurator import S3StoreConfigurator
+from digitalhub.stores.s3.utils import get_bucket_name
 from digitalhub.utils.exceptions import StoreError
 from digitalhub.utils.file_utils import get_file_info_from_s3, get_file_mime_type
-from digitalhub.utils.s3_utils import get_bucket_name
+from digitalhub.utils.types import SourcesOrListOfSources
+
+if typing.TYPE_CHECKING:
+    pass
 
 # Type aliases
 S3Client = Type["botocore.client.S3"]
 
 
-class S3StoreConfig(StoreConfig):
-    """
-    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."""
-
-
 class S3Store(Store):
     """
     S3 store class. It implements the Store interface and provides methods to fetch and persist
     artifacts on S3 based storage.
     """
 
-    def __init__(self, name: str, store_type: str, config: S3StoreConfig) -> None:
-        super().__init__(name, store_type)
-        self.config = config
+    def __init__(self, config: dict | None = None) -> None:
+        super().__init__()
+        self._configurator = S3StoreConfigurator()
+        self._configurator.configure(config)
 
     ##############################
     # I/O methods
@@ -129,7 +118,7 @@ class S3Store(Store):
 
     def upload(
         self,
-        src: str | list[str],
+        src: SourcesOrListOfSources,
         dst: str,
     ) -> list[tuple[str, str]]:
         """
@@ -137,7 +126,7 @@ class S3Store(Store):
 
         Parameters
         ----------
-        src : str | list[str]
+        src : SourcesOrListOfSources
             Source(s).
         dst : str
             The destination of the artifact on storage.
@@ -244,12 +233,121 @@ class S3Store(Store):
 
         return infos
 
+    ##############################
+    # 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)
+
+        # Verify if partition or single file
+        if self.is_partition(path):
+            client, bucket = self._check_factory(path)
+            objects = self._list_objects(client, bucket, path)
+            keys = [self._get_key(o) for o in objects]
+
+        else:
+            if isinstance(path, list):
+                client, bucket = self._check_factory(path[0])
+                keys = [self._get_key(p) for p in path]
+            else:
+                client, bucket = self._check_factory(path)
+                keys = [self._get_key(path)]
+
+        dfs = []
+        for key in keys:
+            file_format = self._get_extension(file_format, key)
+            obj = self._download_fileobject(key, client, bucket)
+            dfs.append(reader.read_df(obj, extension=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("S3 store does not support query.")
+
+    def write_df(
+        self,
+        df: Any,
+        dst: str,
+        extension: str | None = None,
+        **kwargs,
+    ) -> str:
+        """
+        Write a dataframe to S3 based storage. Kwargs are passed to df.to_parquet().
+
+        Parameters
+        ----------
+        df : Any
+            The dataframe.
+        dst : str
+            The destination path on S3 based storage.
+        extension : str
+            The extension of the file.
+        **kwargs : dict
+            Keyword arguments.
+
+        Returns
+        -------
+        str
+            The S3 path where the dataframe was saved.
+        """
+        fileobj = BytesIO()
+        reader = get_reader_by_object(df)
+        reader.write_df(df, fileobj, extension=extension, **kwargs)
+        return self.upload_fileobject(fileobj, dst)
+
     ##############################
     # Private I/O methods
     ##############################
 
+    @staticmethod
     def _download_file(
-        self,
         key: str,
         dst_pth: Path,
         client: S3Client,
@@ -260,8 +358,8 @@ class S3Store(Store):
 
         Parameters
         ----------
-        keys : str
-            The list of keys to be downloaded.
+        key : str
+            The key to be downloaded.
         dst_pth : str
             The destination of the files on local filesystem.
         client : S3Client
@@ -274,9 +372,34 @@ class S3Store(Store):
         list[str]
             The list of paths of the downloaded files.
         """
-        # Download file
         client.download_file(bucket, key, dst_pth)
 
+    @staticmethod
+    def _download_fileobject(
+        key: str,
+        client: S3Client,
+        bucket: str,
+    ) -> BytesIO:
+        """
+        Download fileobject from S3 partition.
+
+        Parameters
+        ----------
+        key : str
+            The key of the file.
+        client : S3Client
+            The S3 client object.
+        bucket : str
+            The name of the S3 bucket.
+
+        Returns
+        -------
+        BytesIO
+            The fileobject of the downloaded file.
+        """
+        obj = client.get_object(Bucket=bucket, Key=key)
+        return BytesIO(obj["Body"].read())
+
     def _upload_dir(
         self,
         src: str,
@@ -453,41 +576,6 @@ class S3Store(Store):
         """
         client.put_object(Bucket=bucket, Key=key, Body=fileobj.getvalue())
 
-    ##############################
-    # Datastore methods
-    ##############################
-
-    def write_df(
-        self,
-        df: Any,
-        dst: str,
-        extension: str | None = None,
-        **kwargs,
-    ) -> str:
-        """
-        Write a dataframe to S3 based storage. Kwargs are passed to df.to_parquet().
-
-        Parameters
-        ----------
-        df : Any
-            The dataframe.
-        dst : str
-            The destination path on S3 based storage.
-        extension : str
-            The extension of the file.
-        **kwargs : dict
-            Keyword arguments.
-
-        Returns
-        -------
-        str
-            The S3 path where the dataframe was saved.
-        """
-        fileobj = BytesIO()
-        reader = get_reader_by_object(df)
-        reader.write_df(df, fileobj, extension=extension, **kwargs)
-        return self.upload_fileobject(fileobj, dst)
-
     ##############################
     # Helper methods
     ##############################
@@ -512,11 +600,7 @@ class S3Store(Store):
         S3Client
             Returns a client object that interacts with the S3 storage service.
         """
-        cfg = {
-            "endpoint_url": self.config.endpoint_url,
-            "aws_access_key_id": self.config.aws_access_key_id,
-            "aws_secret_access_key": self.config.aws_secret_access_key,
-        }
+        cfg = self._configurator.get_boto3_client_config()
         return boto3.client("s3", **cfg)
 
     def _check_factory(self, root: str) -> tuple[S3Client, str]:

digitalhub/{utils/s3_utils.py → stores/s3/utils.py}

@@ -6,6 +6,10 @@ from urllib.parse import urlparse
 
 from boto3 import client as boto3_client
 
+from digitalhub.stores.s3.enums import S3StoreEnv
+
+DEFAULT_BUCKET = "datalake"
+
 
 def get_bucket_name(path: str) -> str:
     """
@@ -59,11 +63,11 @@ def get_s3_source(bucket: str, key: str, filename: Path) -> None:
     -------
     None
     """
-    s3 = boto3_client("s3", endpoint_url=os.getenv("S3_ENDPOINT_URL"))
+    s3 = boto3_client("s3", endpoint_url=os.getenv(S3StoreEnv.ENDPOINT_URL.value))
     s3.download_file(bucket, key, filename)
 
 
-def get_s3_bucket() -> str | None:
+def get_s3_bucket_from_env() -> str | None:
     """
     Function to get S3 bucket name.
 
@@ -72,4 +76,4 @@ def get_s3_bucket() -> str | None:
     str
         The S3 bucket name.
     """
-    return os.getenv("S3_BUCKET_NAME", "datalake")
+    return os.getenv(S3StoreEnv.BUCKET_NAME.value, DEFAULT_BUCKET)

digitalhub/stores/sql/configurator.py

@@ -0,0 +1,88 @@
+from __future__ import annotations
+
+from digitalhub.configurator.configurator import configurator
+from digitalhub.stores.sql.enums import SqlStoreEnv
+from digitalhub.stores.sql.models import SqlStoreConfig
+from digitalhub.utils.exceptions import StoreError
+
+
+class SqlStoreConfigurator:
+    """
+    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
+        """
+        if config is None:
+            self._get_config()
+        else:
+            config: SqlStoreConfig = SqlStoreConfig(**config)
+            for pair in [
+                (SqlStoreEnv.USERNAME.value, config.user),
+                (SqlStoreEnv.PASSWORD.value, config.password),
+                (SqlStoreEnv.HOST.value, config.host),
+                (SqlStoreEnv.PORT.value, config.port),
+                (SqlStoreEnv.DATABASE.value, config.database),
+            ]:
+                configurator.set_credential(*pair)
+
+    def get_sql_conn_string(self) -> str:
+        """
+        Get the connection string from environment variables.
+
+        Returns
+        -------
+        str
+            The connection string.
+        """
+        creds = configurator.get_all_credentials()
+        try:
+            user = creds[SqlStoreEnv.USERNAME.value]
+            password = creds[SqlStoreEnv.PASSWORD.value]
+            host = creds[SqlStoreEnv.HOST.value]
+            port = creds[SqlStoreEnv.PORT.value]
+            database = creds[SqlStoreEnv.DATABASE.value]
+            return f"postgresql://{user}:{password}@{host}:{port}/{database}"
+        except KeyError as e:
+            raise StoreError(f"Missing credentials for SQL store. {str(e)}")
+
+    def _get_config(self) -> None:
+        """
+        Get the credentials from environment variables.
+
+        Returns
+        -------
+        None
+        """
+        user = configurator.load_var(SqlStoreEnv.USERNAME.value)
+        password = configurator.load_var(SqlStoreEnv.PASSWORD.value)
+        host = configurator.load_var(SqlStoreEnv.HOST.value)
+        port = configurator.load_var(SqlStoreEnv.PORT.value)
+        database = configurator.load_var(SqlStoreEnv.DATABASE.value)
+        if user is None or password is None or host is None or port is None or database is None:
+            raise StoreError("Missing credentials for SQL store.")
+        configurator.set_credential(SqlStoreEnv.USERNAME.value, user)
+        configurator.set_credential(SqlStoreEnv.PASSWORD.value, password)
+        configurator.set_credential(SqlStoreEnv.HOST.value, host)
+        configurator.set_credential(SqlStoreEnv.PORT.value, port)
+        configurator.set_credential(SqlStoreEnv.DATABASE.value, database)

digitalhub/stores/sql/enums.py

@@ -0,0 +1,16 @@
+from __future__ import annotations
+
+from enum import Enum
+
+
+class SqlStoreEnv(Enum):
+    """
+    SqlStore environment
+    """
+
+    HOST = "DB_HOST"
+    PORT = "DB_PORT"
+    USERNAME = "DB_USERNAME"
+    PASSWORD = "DB_PASSWORD"
+    DATABASE = "DB_DATABASE"
+    PG_SCHEMA = "DB_SCHEMA"

digitalhub/stores/sql/models.py

@@ -0,0 +1,24 @@
+from __future__ import annotations
+
+from pydantic import BaseModel
+
+
+class SqlStoreConfig(BaseModel):
+    """
+    SQL store configuration class.
+    """
+
+    host: str
+    """SQL host."""
+
+    port: int
+    """SQL port."""
+
+    user: str
+    """SQL user."""
+
+    password: str
+    """SQL password."""
+
+    database: str
+    """SQL database name."""