digitalhub 0.8.0b0__py3-none-any.whl → 0.8.0b2__py3-none-any.whl

digitalhub/stores/objects/sql.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.objects.base 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/env_utils.py

@@ -0,0 +1,123 @@
+from __future__ import annotations
+
+import os
+import typing
+
+from digitalhub.client.builder import check_client_exists, get_client
+
+if typing.TYPE_CHECKING:
+    from digitalhub.client.objects.dhcore import ClientDHCore
+
+
+def set_dhcore_env(
+    endpoint: str | None = None,
+    user: str | None = None,
+    password: str | None = None,
+    access_token: str | None = None,
+    refresh_token: str | None = None,
+    client_id: str | None = None,
+) -> None:
+    """
+    Function to set environment variables for DHCore config.
+    Note that if the environment variable is already set, it
+    will be overwritten. It also ovverides the remote client
+    configuration.
+
+    Parameters
+    ----------
+    endpoint : str
+        The endpoint of DHCore.
+    user : str
+        The user of DHCore.
+    password : str
+        The password of DHCore.
+    access_token : str
+        The access token of DHCore.
+    refresh_token : str
+        The refresh token of DHCore.
+    client_id : str
+        The client id of DHCore.
+
+    Returns
+    -------
+    None
+    """
+    if endpoint is not None:
+        os.environ["DHCORE_ENDPOINT"] = endpoint
+    if user is not None:
+        os.environ["DHCORE_USER"] = user
+    if password is not None:
+        os.environ["DHCORE_PASSWORD"] = password
+    if access_token is not None:
+        os.environ["DHCORE_ACCESS_TOKEN"] = access_token
+    if refresh_token is not None:
+        os.environ["DHCORE_REFRESH_TOKEN"] = refresh_token
+    if client_id is not None:
+        os.environ["DHCORE_CLIENT_ID"] = client_id
+
+    if check_client_exists(local=False):
+        update_client_from_env()
+
+
+def update_client_from_env() -> None:
+    """
+    Function to update client from environment variables.
+
+    Returns
+    -------
+    None
+    """
+    client: ClientDHCore = get_client(local=False)
+
+    # Update endpoint
+    endpoint = os.getenv("DHCORE_ENDPOINT")
+    if endpoint is not None:
+        client._endpoint_core = endpoint
+
+    # Update auth
+
+    # If token is set, it will override the other auth options
+    access_token = os.getenv("DHCORE_ACCESS_TOKEN")
+    refresh_token = os.getenv("DHCORE_REFRESH_TOKEN")
+    client_id = os.getenv("DHCORE_CLIENT_ID")
+
+    if access_token is not None:
+        if refresh_token is not None:
+            client._refresh_token = refresh_token
+        if client_id is not None:
+            client._client_id = client_id
+        client._access_token = access_token
+        client._auth_type = "oauth2"
+        return
+
+    # Otherwise, if user and password are set, basic auth will be used
+    username = os.getenv("DHCORE_USER")
+    password = os.getenv("DHCORE_PASSWORD")
+    if username is not None and password is not None:
+        client._user = username
+        client._password = password
+        client._auth_type = "basic"
+
+
+def refresh_token() -> None:
+    """
+    Function to refresh token.
+
+    Returns
+    -------
+    None
+    """
+    client: ClientDHCore = get_client(local=False)
+    client._get_new_access_token()
+
+
+def get_s3_bucket() -> str | None:
+    """
+    Function to get S3 bucket name.
+
+    Returns
+    -------
+    str
+        The S3 bucket name.
+    """
+    return os.getenv("S3_BUCKET_NAME", "datalake")

digitalhub/utils/exceptions.py

@@ -0,0 +1,55 @@
+from __future__ import annotations
+
+
+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.
+    """