digitalhub 0.10.0b1__py3-none-any.whl → 0.10.0b2__py3-none-any.whl

digitalhub/stores/local/store.py

@@ -4,7 +4,7 @@ import shutil
 from pathlib import Path
 from typing import Any
 
-from digitalhub.readers.api import get_reader_by_object
+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
@@ -82,6 +82,106 @@ class LocalStore(Store):
         """
         return [get_file_info_from_local(p) for p in paths]
 
+    ##############################
+    # Datastore methods
+    ##############################
+
+    def read_df(
+        self,
+        path: str | list[str],
+        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.
+        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
     ##############################
@@ -183,73 +283,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._check_local_dst(dst)
-        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

@@ -97,6 +97,58 @@ class RemoteStore(Store):
     # Datastore methods
     ##############################
 
+    def read_df(
+        self,
+        path: str | list[str],
+        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.
+        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
@@ -160,3 +212,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

@@ -51,9 +51,10 @@ class S3StoreConfigurator:
             ]:
                 configurator.set_credential(*pair)
 
-    def get_s3_creds(self) -> dict:
+    def get_boto3_client_config(self) -> dict:
         """
-        Get endpoint, access key and secret key.
+        Get S3 credentials (access key, secret key,
+        session token and other config).
 
         Returns
         -------

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,13 +10,16 @@ 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.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
 
+if typing.TYPE_CHECKING:
+    pass
+
 # Type aliases
 S3Client = Type["botocore.client.S3"]
 
@@ -228,12 +232,121 @@ class S3Store(Store):
 
         return infos
 
+    ##############################
+    # Datastore methods
+    ##############################
+
+    def read_df(
+        self,
+        path: str | list[str],
+        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.
+        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,
@@ -244,8 +357,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
@@ -258,9 +371,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,
@@ -437,41 +575,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
     ##############################
@@ -496,7 +599,7 @@ class S3Store(Store):
         S3Client
             Returns a client object that interacts with the S3 storage service.
         """
-        cfg = self._configurator.get_s3_creds()
+        cfg = self._configurator.get_boto3_client_config()
         return boto3.client("s3", **cfg)
 
     def _check_factory(self, root: str) -> tuple[S3Client, str]:

digitalhub/stores/sql/store.py

@@ -10,7 +10,7 @@ from sqlalchemy import MetaData, Table, create_engine, select
 from sqlalchemy.engine import Engine
 from sqlalchemy.exc import SQLAlchemyError
 
-from digitalhub.readers.api import get_reader_by_object
+from digitalhub.readers.data.api import get_reader_by_object
 from digitalhub.stores._base.store import Store
 from digitalhub.stores.sql.configurator import SqlStoreConfigurator
 from digitalhub.utils.exceptions import StoreError
@@ -119,50 +119,72 @@ class SqlStore(Store):
         return []
 
     ##############################
-    # Private I/O methods
+    # Datastore methods
     ##############################
 
-    def _download_table(self, schema: str, table: str, dst: str) -> str:
+    def read_df(
+        self,
+        path: str | list[str],
+        file_format: str | None = None,
+        engine: str | None = None,
+        **kwargs,
+    ) -> Any:
         """
-        Download a table from SQL based storage.
+        Read DataFrame from path.
 
         Parameters
         ----------
-        schema : str
-            The origin schema.
-        table : str
-            The origin table.
-        dst : str
-            The destination path.
+        path : str | list[str]
+            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
         -------
-        str
-            The destination path.
+        Any
+            DataFrame.
         """
-        engine = self._check_factory(schema=schema)
+        reader = self._get_reader(engine)
+        schema = self._get_schema(path)
+        table = self._get_table_name(path)
+        sql_engine = self._check_factory(schema=schema)
 
-        # Read the table from the database
-        sa_table = Table(table, MetaData(), autoload_with=engine)
+        sa_table = Table(table, MetaData(), autoload_with=sql_engine)
         stm = select(sa_table)
-        with engine.begin() as conn:
-            result: list[Row] = conn.execute(stm).fetchall()
 
-        # Parse the result
-        data = {col: [row[idx] for row in result] for idx, col in enumerate(sa_table.columns.keys())}
-
-        # Convert the result to a pyarrow table and
-        # write the pyarrow table to a Parquet file
-        arrow_table = pa.Table.from_pydict(data)
-        pq.write_table(arrow_table, dst)
+        return reader.read_table(stm, sql_engine, **kwargs)
 
-        engine.dispose()
+    def query(
+        self,
+        query: str,
+        path: str,
+        engine: str | None = None,
+    ) -> Any:
+        """
+        Query data from database.
 
-        return dst
+        Parameters
+        ----------
+        query : str
+            The query to execute.
+        path : str
+            Path to the database.
+        engine : str
+            Dataframe engine (pandas, polars, etc.).
 
-    ##############################
-    # Datastore methods
-    ##############################
+        Returns
+        -------
+        Any
+            DataFrame.
+        """
+        reader = self._get_reader(engine)
+        schema = self._get_schema(path)
+        sql_engine = self._check_factory(schema=schema)
+        return reader.read_table(query, sql_engine)
 
     def write_df(self, df: Any, dst: str, extension: str | None = None, **kwargs) -> str:
         """
@@ -187,9 +209,47 @@ class SqlStore(Store):
         return self._upload_table(df, schema, table, **kwargs)
 
     ##############################
-    # Private Datastore methods
+    # Private I/O methods
     ##############################
 
+    def _download_table(self, schema: str, table: str, dst: str) -> str:
+        """
+        Download a table from SQL based storage.
+
+        Parameters
+        ----------
+        schema : str
+            The origin schema.
+        table : str
+            The origin table.
+        dst : str
+            The destination path.
+
+        Returns
+        -------
+        str
+            The destination path.
+        """
+        engine = self._check_factory(schema=schema)
+
+        # Read the table from the database
+        sa_table = Table(table, MetaData(), autoload_with=engine)
+        stm = select(sa_table)
+        with engine.begin() as conn:
+            result: list[Row] = conn.execute(stm).fetchall()
+
+        # Parse the result
+        data = {col: [row[idx] for row in result] for idx, col in enumerate(sa_table.columns.keys())}
+
+        # Convert the result to a pyarrow table and
+        # write the pyarrow table to a Parquet file
+        arrow_table = pa.Table.from_pydict(data)
+        pq.write_table(arrow_table, dst)
+
+        engine.dispose()
+
+        return dst
+
     def _upload_table(self, df: Any, schema: str, table: str, **kwargs) -> str:
         """
         Upload a table to SQL based storage.