digitalhub 0.7.0b2__py3-none-any.whl → 0.8.0__py3-none-any.whl

digitalhub/stores/sql/store.py

@@ -0,0 +1,328 @@
+from __future__ import annotations
+
+from pathlib import Path
+
+import pyarrow as pa
+import pyarrow.parquet as pq
+from sqlalchemy import MetaData, Table, create_engine
+from sqlalchemy.engine import Engine
+from sqlalchemy.engine.row import LegacyRow
+from sqlalchemy.exc import SQLAlchemyError
+
+from digitalhub.stores._base.store import Store, StoreConfig
+from digitalhub.utils.exceptions import StoreError
+
+
+class SQLStoreConfig(StoreConfig):
+    """
+    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."""
+
+
+class SqlStore(Store):
+    """
+    SQL store class. It implements the Store interface and provides methods to fetch and persist
+    artifacts on SQL based storage.
+    """
+
+    def __init__(self, name: str, store_type: str, config: SQLStoreConfig) -> None:
+        super().__init__(name, store_type)
+        self.config = config
+
+    ##############################
+    # IO methods
+    ##############################
+
+    def download(
+        self,
+        root: str,
+        dst: Path,
+        src: list[str],
+        overwrite: bool = False,
+    ) -> str:
+        """
+        Download artifacts from storage.
+
+        Parameters
+        ----------
+        root : str
+            The root path of the artifact.
+        dst : str
+            The destination of the artifact on local filesystem.
+        src : list[str]
+            List of sources.
+        overwrite : bool
+            Specify if overwrite existing file(s).
+
+        Returns
+        -------
+        str
+            Destination path of the downloaded artifact.
+        """
+        table_name = self._get_table_name(root) + ".parquet"
+        # Case where dst is not provided
+        if dst is None:
+            dst = Path(self._build_temp("sql")) / table_name
+        else:
+            self._check_local_dst(str(dst))
+            path = Path(dst)
+
+            # Case where dst is a directory
+            if path.suffix == "":
+                dst = path / table_name
+
+            # Case where dst is a file
+            elif path.suffix != ".parquet":
+                raise StoreError("The destination path must be a directory or a parquet file.")
+
+            self._check_overwrite(dst, overwrite)
+            self._build_path(dst)
+
+        schema = self._get_schema(root)
+        table = self._get_table_name(root)
+        return self._download_table(schema, table, str(dst))
+
+    def upload(self, src: str | list[str], dst: str | None = None) -> list[tuple[str, str]]:
+        """
+        Upload an artifact to storage.
+
+        Raises
+        ------
+        StoreError
+            This method is not implemented.
+        """
+        raise StoreError("SQL store does not support upload.")
+
+    def get_file_info(self, paths: list[str]) -> list[dict]:
+        """
+        Get file information from SQL based storage.
+
+        Raises
+        ------
+        NotImplementedError
+            This method is not implemented.
+        """
+        raise NotImplementedError("SQL store does not support upload.")
+
+    ##############################
+    # Private helper methods
+    ##############################
+
+    def _get_connection_string(self) -> str:
+        """
+        Get the connection string.
+
+        Returns
+        -------
+        str
+            The connection string.
+        """
+        return (
+            f"postgresql://{self.config.user}:{self.config.password}@"
+            f"{self.config.host}:{self.config.port}/{self.config.database}"
+        )
+
+    def _get_engine(self, schema: str | None = None) -> Engine:
+        """
+        Create engine from connection string.
+
+        Parameters
+        ----------
+        schema : str
+            The schema.
+
+        Returns
+        -------
+        Engine
+            An SQLAlchemy engine.
+        """
+        connection_string = self._get_connection_string()
+        if not isinstance(connection_string, str):
+            raise StoreError("Connection string must be a string.")
+        try:
+            connect_args = {"connect_timeout": 30}
+            if schema is not None:
+                connect_args["options"] = f"-csearch_path={schema}"
+            return create_engine(connection_string, connect_args=connect_args)
+        except Exception as ex:
+            raise StoreError(f"Something wrong with connection string. Arguments: {str(ex.args)}")
+
+    def _check_factory(self, schema: str | None = None) -> Engine:
+        """
+        Check if the database is accessible and return the engine.
+
+        Parameters
+        ----------
+        schema : str
+            The schema.
+
+        Returns
+        -------
+        Engine
+            The database engine.
+        """
+        engine = self._get_engine(schema)
+        self._check_access_to_storage(engine)
+        return engine
+
+    @staticmethod
+    def _parse_path(path: str) -> dict:
+        """
+        Parse the path and return the components.
+
+        Parameters
+        ----------
+        path : str
+            The path.
+
+        Returns
+        -------
+        dict
+            A dictionary containing the components of the path.
+        """
+        # Parse path
+        err_msg = "Invalid SQL path. Must be sql://<database>/<schema>/<table> or sql://<database>/<table>"
+        protocol, pth = path.split("://")
+        components = pth.split("/")
+        if protocol != "sql" or not (2 <= len(components) <= 3):
+            raise ValueError(err_msg)
+
+        # Get components
+        database = components[0]
+        table = components[-1]
+        schema = components[1] if len(components) == 3 else "public"
+        return {"database": database, "schema": schema, "table": table}
+
+    def _get_schema(self, uri: str) -> str:
+        """
+        Get the name of the SQL schema from the URI.
+
+        Parameters
+        ----------
+        uri : str
+            The URI.
+
+        Returns
+        -------
+        str
+            The name of the SQL schema.
+        """
+        return str(self._parse_path(uri).get("schema"))
+
+    def _get_table_name(self, uri: str) -> str:
+        """
+        Get the name of the table from the URI.
+
+        Parameters
+        ----------
+        uri : str
+            The URI.
+
+        Returns
+        -------
+        str
+            The name of the table
+        """
+        return str(self._parse_path(uri).get("table"))
+
+    @staticmethod
+    def _check_access_to_storage(engine: Engine) -> None:
+        """
+        Check if there is access to the storage.
+
+        Parameters
+        ----------
+        engine : Engine
+            An SQLAlchemy engine.
+
+        Returns
+        -------
+        None
+
+        Raises
+        ------
+        StoreError
+            If there is no access to the storage.
+        """
+        try:
+            engine.connect()
+        except SQLAlchemyError:
+            engine.dispose()
+            raise StoreError("No access to db!")
+
+    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)
+        query = sa_table.select()
+        with engine.begin() as conn:
+            result: list[LegacyRow] = conn.execute(query).fetchall()
+
+        # Parse the result
+        data = self._parse_result(result)
+
+        # 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
+
+    @staticmethod
+    def _parse_result(result: list[LegacyRow]) -> dict:
+        """
+        Convert a list of list of tuples to a dict.
+
+        Parameters
+        ----------
+        result : list[LegacyRow]
+            The data to convert.
+
+        Returns
+        -------
+        dict
+            The converted data.
+        """
+        data_list = [row.items() for row in result]
+        data = {}
+        for row in data_list:
+            for column_name, value in row:
+                if column_name not in data:
+                    data[column_name] = []
+                data[column_name].append(value)
+        return data

digitalhub/utils/data_utils.py

@@ -0,0 +1,127 @@
+from __future__ import annotations
+
+import json
+
+
+def build_data_preview(preview: list[dict] | None = None, rows_count: int | None = None) -> dict:
+    """
+    Build data preview.
+
+    Parameters
+    ----------
+    preview : list[dict] | None
+        Preview.
+    rows_count : int | None
+        Row count.
+
+    Returns
+    -------
+    dict
+        Data preview.
+    """
+    dict_ = {}
+    if preview is not None:
+        dict_["cols"] = preview
+    if rows_count is not None:
+        dict_["rows"] = rows_count
+    return dict_
+
+
+def get_data_preview(columns: list, data: list[list], columnar: bool = False) -> list[dict]:
+    """
+    Prepare preview.
+
+    Parameters
+    ----------
+    columns : list
+        Columns names.
+    data : list[list]
+        Data to preview.
+    columnar : bool
+        If data are arranged in columns. If False, data are arranged in rows.
+
+    Returns
+    -------
+    list[dict]
+        Data preview.
+    """
+    # Reduce data to 10 rows
+    if not columnar:
+        if len(data) > 10:
+            data = data[:10]
+    else:
+        data = [d[:10] for d in data]
+
+    # Transpose data if needed
+    if not columnar:
+        data = list(map(list, list(zip(*data))))
+
+    # Prepare the preview
+    data_dict = prepare_preview(columns, data)
+
+    # Filter memoryview values
+    filtered_memview = filter_memoryview(data_dict)
+
+    # Check the size of the preview data
+    return check_preview_size(filtered_memview)
+
+
+def prepare_preview(column_names: list, data: list[list]) -> list[dict]:
+    """
+    Get preview.
+
+    Parameters
+    ----------
+    data : pd.DataFrame
+        Data.
+
+    Returns
+    -------
+    list[dict]
+        Preview.
+    """
+    if len(column_names) != len(data):
+        raise ValueError("Column names and data must have the same length")
+    return [{"name": column, "value": values} for column, values in zip(column_names, data)]
+
+
+def filter_memoryview(data: list[dict]) -> list[dict]:
+    """
+    Find memoryview values.
+
+    Parameters
+    ----------
+    data : pd.DataFrame
+        Data.
+
+    Returns
+    -------
+    list[str]
+        Column to filter out from preview.
+    """
+    key_to_filter = []
+    for i in data:
+        if any(isinstance(v, memoryview) for v in i["value"]):
+            key_to_filter.append(i["name"])
+    for i in key_to_filter:
+        data = [d for d in data if d["name"] != i]
+    return data
+
+
+def check_preview_size(preview: list[dict]) -> list:
+    """
+    Check preview size. If it's too big, return empty list.
+
+    Parameters
+    ----------
+    preview : list[dict]
+        Preview.
+
+    Returns
+    -------
+    list
+        Preview.
+    """
+    if len(json.dumps(preview).encode("utf-8")) >= 64000:
+        return []
+    return preview

digitalhub/utils/exceptions.py

@@ -0,0 +1,67 @@
+from __future__ import annotations
+
+
+class BuilderError(Exception):
+    """
+    Raised when incontered errors on builders.
+    """
+
+
+class StoreError(Exception):
+    """
+    Raised when incontered errors on stores.
+    """
+
+
+class BackendError(Exception):
+    """
+    Raised when incontered errors from backend.
+    """
+
+
+class EntityNotExistsError(BackendError):
+    """
+    Raised when entity not found.
+    """
+
+
+class EntityAlreadyExistsError(BackendError):
+    """
+    Raised when entity already exists.
+    """
+
+
+class MissingSpecError(BackendError):
+    """
+    Raised when spec is missing in backend.
+    """
+
+
+class UnauthorizedError(BackendError):
+    """
+    Raised when unauthorized.
+    """
+
+
+class ForbiddenError(BackendError):
+    """
+    Raised when forbidden.
+    """
+
+
+class BadRequestError(BackendError):
+    """
+    Raised when bad request.
+    """
+
+
+class EntityError(Exception):
+    """
+    Raised when incontered errors on entities.
+    """
+
+
+class ContextError(Exception):
+    """
+    Raised when context errors.
+    """

digitalhub/utils/file_utils.py

@@ -0,0 +1,204 @@
+from __future__ import annotations
+
+from datetime import datetime
+from hashlib import sha256
+from mimetypes import guess_type
+from pathlib import Path
+
+from pydantic import BaseModel
+
+
+class FileInfo(BaseModel):
+    """
+    File info class.
+    """
+
+    path: str = None
+    name: str = None
+    content_type: str = None
+    size: int = None
+    hash: str = None
+    last_modified: str = None
+
+
+def calculate_blob_hash(data_path: str) -> str:
+    """
+    Calculate the hash of a file.
+
+    Parameters
+    ----------
+    data_path : str
+        Path to the file.
+
+    Returns
+    -------
+    str
+        The hash of the file.
+    """
+    with open(data_path, "rb") as f:
+        data = f.read()
+        return f"sha256:{sha256(data).hexdigest()}"
+
+
+def get_file_size(data_path: str) -> int:
+    """
+    Get the size of a file.
+
+    Parameters
+    ----------
+    data_path : str
+        Path to the file.
+
+    Returns
+    -------
+    int
+        The size of the file.
+    """
+    return Path(data_path).stat().st_size
+
+
+def get_file_mime_type(data_path: str) -> str:
+    """
+    Get the mime type of a file.
+
+    Parameters
+    ----------
+    data_path : str
+        Path to the file.
+
+    Returns
+    -------
+    str
+        The mime type of the file.
+    """
+    return guess_type(data_path)[0]
+
+
+def get_path_name(data_path: str) -> str:
+    """
+    Get the name of a file.
+
+    Parameters
+    ----------
+    data_path : str
+        Path to the file.
+
+    Returns
+    -------
+    str
+        The name of the file.
+    """
+    return Path(data_path).name
+
+
+def get_last_modified(data_path: str) -> str:
+    """
+    Get the last modified date of a file.
+
+    Parameters
+    ----------
+    data_path : str
+        Path to the file.
+
+    Returns
+    -------
+    str
+        The last modified date of the file.
+    """
+    path = Path(data_path)
+    timestamp = path.stat().st_mtime
+    return datetime.fromtimestamp(timestamp).astimezone().isoformat()
+
+
+def get_s3_path(src_path: str) -> str:
+    """
+    Get the S3 path of a file.
+
+    Parameters
+    ----------
+    src_path : str
+        Path to the file.
+
+    Returns
+    -------
+    str
+        The S3 path of the file.
+    """
+    return Path(src_path).as_uri()
+
+
+def get_file_info_from_local(path: str, src_path: str) -> None | dict:
+    """
+    Get file info from path.
+
+    Parameters
+    ----------
+    path : str
+        Target path of the object.
+    src_path : str
+        Local path of some source.
+
+    Returns
+    -------
+    dict
+        File info.
+    """
+    try:
+        name = get_path_name(path)
+        content_type = get_file_mime_type(path)
+        size = get_file_size(path)
+        hash = calculate_blob_hash(path)
+        last_modified = get_last_modified(path)
+
+        return FileInfo(
+            path=src_path,
+            name=name,
+            content_type=content_type,
+            size=size,
+            hash=hash,
+            last_modified=last_modified,
+        ).dict()
+    except Exception:
+        return None
+
+
+def get_file_info_from_s3(path: str, metadata: dict) -> None | dict:
+    """
+    Get file info from path.
+
+    Parameters
+    ----------
+    path : str
+        Object source path.
+    metadata : dict
+        Metadata of the object from S3.
+
+    Returns
+    -------
+    dict
+        File info.
+    """
+    try:
+        size = metadata["ContentLength"]
+        file_hash = metadata["ETag"][1:-1]
+
+        file_size_limit_multipart = 20 * 1024 * 1024
+        if size < file_size_limit_multipart:
+            file_hash = "md5:" + file_hash
+        else:
+            file_hash = "LiteralETag:" + file_hash
+
+        name = get_path_name(path)
+        content_type = metadata["ContentType"]
+        last_modified = metadata["LastModified"].isoformat()
+
+        return FileInfo(
+            path=path,
+            name=name,
+            content_type=content_type,
+            size=size,
+            hash=file_hash,
+            last_modified=last_modified,
+        ).dict()
+    except Exception:
+        return None