plexus-python-pulse-common 1.1.76__tar.gz → 1.1.77__tar.gz

{plexus_python_pulse_common-1.1.76 → plexus_python_pulse_common-1.1.77}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: plexus-python-pulse-common
-Version: 1.1.76
+Version: 1.1.77
 Classifier: Programming Language :: Python :: 3
 Classifier: Programming Language :: Python :: 3.12
 Classifier: Programming Language :: Python :: 3.13

plexus_python_pulse_common-1.1.77/src/plexus/pulse/common/utils/flowutils.py

@@ -0,0 +1,167 @@
+import contextlib
+import enum
+import os
+import pathlib
+import threading
+import typing
+import uuid as py_uuid
+from collections.abc import Generator
+from typing import Self
+
+import pydantic as pdt
+import sqlalchemy as sa
+import sqlalchemy.dialects.sqlite as sa_sqlite
+import sqlalchemy.orm as sa_orm
+from plexus.common.utils.dbutils import ConnectionMaker
+from plexus.common.utils.funcutils import Chainable
+from plexus.common.utils.funcutils import chainable, singleton
+from sqlmodel import Field
+
+from plexus.pulse.common.utils.ormutils import SQLiteUUID, SequenceModelMixinProtocol
+from plexus.pulse.common.utils.ormutils import clone_sequence_model_instance, make_base_model, make_sequence_model_mixin
+
+__all__ = [
+    "FlowStageEnum",
+    "current_flow_stage",
+    "current_flow_run_id",
+    "DatasetRef",
+    "DatasetRefTable",
+    "flow_context_file_path",
+    "FlowContext",
+    "context",
+]
+
+
+class FlowStageEnum(enum.Enum):
+    Preprocessing = "preprocessing"
+    Annotation = "annotation"
+    Curation = "curation"
+    Training = "training"
+
+
+def current_flow_stage() -> FlowStageEnum:
+    stage_str = os.environ.get("PULSE_FLOW_STAGE")
+    if stage_str is None:
+        raise ValueError("Missing required environment variable 'PULSE_FLOW_STAGE'")
+    for stage in list(FlowStageEnum):
+        if stage_str == stage.value:
+            return stage
+    raise ValueError(f"Illegal 'PULSE_FLOW_STAGE' value '{stage_str}'")
+
+
+def current_flow_run_id() -> str:
+    run_id_str = os.environ.get("PULSE_FLOW_RUN_ID")
+    if run_id_str is None:
+        raise ValueError("Missing required environment variable 'PULSE_FLOW_RUN_ID'")
+    return run_id_str
+
+
+def make_dataset_ref_model_mixin() -> type[pdt.BaseModel]:
+    class ModelMixin(pdt.BaseModel):
+        dataset_uuid: py_uuid.UUID = Field(
+            sa_column=sa.Column(SQLiteUUID(), nullable=False),
+            description="UUID of the dataset",
+        )
+        dataset_name: str | None = Field(
+            sa_column=sa.Column(sa_sqlite.VARCHAR(128), nullable=True),
+            description="Name of the dataset",
+        )
+
+    return ModelMixin
+
+
+BaseModel = make_base_model()
+
+
+class DatasetRef(BaseModel, make_dataset_ref_model_mixin()):
+    pass
+
+
+class DatasetRefTable(DatasetRef, make_sequence_model_mixin("sqlite"), table=True):
+    __tablename__ = "dataset_ref"
+
+
+if typing.TYPE_CHECKING:
+    class DatasetRef(BaseModel):
+        dataset_uuid: sa_orm.Mapped[py_uuid.UUID] = ...
+        dataset_name: sa_orm.Mapped[str | None] = ...
+
+
+    class DatasetRefTable(DatasetRef, SequenceModelMixinProtocol):
+        pass
+
+
+def is_resolve_dataset_ref_stage(stage: FlowStageEnum | None) -> bool:
+    if stage is None:
+        stage = current_flow_stage()
+    return stage in {FlowStageEnum.Curation}
+
+
+def is_construct_dataset_ref_stage(stage: FlowStageEnum | None) -> bool:
+    if stage is None:
+        stage = current_flow_stage()
+    return stage in {FlowStageEnum.Training}
+
+
+def flow_context_file_path() -> pathlib.Path:
+    root_dir = os.environ.get("PULSE_FLOW_CONTEXT_ROOT_DIR")
+    if root_dir is None:
+        raise ValueError("Missing required environment variable 'PULSE_FLOW_CONTEXT_ROOT_DIR'")
+    return pathlib.Path(root_dir) / current_flow_run_id() / "flow_context.db"
+
+
+class FlowContext(object):
+    def __init__(
+        self,
+        *,
+        file_path: str | os.PathLike[str] | None = None,
+        fail_if_exists: bool = False,
+    ):
+        self.file_path = pathlib.Path(file_path or flow_context_file_path())
+        if fail_if_exists and self.file_path.exists():
+            raise FileExistsError(f"flow context file '{self.file_path}' already exists")
+        self.file_path.parent.mkdir(parents=True, exist_ok=True)
+
+        self.conn_maker = ConnectionMaker.from_url(f"sqlite:///{self.file_path.absolute()}",
+                                                   engine_opts=dict(connect_args={"check_same_thread": False}))
+
+        self.thread_lock = threading.Lock()
+
+        BaseModel.metadata.create_all(self.conn_maker.engine)
+
+    @contextlib.contextmanager
+    def make_session(self) -> Generator[sa_orm.Session, None, None]:
+        with self.thread_lock:
+            with self.conn_maker.make_session(expire_on_commit=False) as session:
+                try:
+                    yield session
+                    session.commit()
+                except Exception:
+                    session.rollback()
+                    raise
+
+    def count_dataset_ref(self) -> int:
+        with self.make_session() as session:
+            return session.query(DatasetRefTable).count()
+
+    def list_dataset_refs(self) -> list[DatasetRefTable]:
+        with self.make_session() as session:
+            return session.query(DatasetRefTable).all()
+
+    def add_dataset_ref(
+        self,
+        dataset_uuid: py_uuid.UUID,
+        dataset_name: str | None = None,
+    ) -> Chainable[Self, DatasetRefTable]:
+        with self.make_session() as session:
+            dataset_ref = DatasetRef(dataset_uuid=dataset_uuid, dataset_name=dataset_name)
+            db_dataset_ref = clone_sequence_model_instance(DatasetRefTable, dataset_ref)
+            session.add(db_dataset_ref)
+            session.commit()
+
+            return chainable(self, db_dataset_ref)
+
+
+@singleton
+def context() -> FlowContext:
+    return FlowContext()

{plexus_python_pulse_common-1.1.76 → plexus_python_pulse_common-1.1.77}/src/plexus/pulse/common/utils/ormutils.py

@@ -1,4 +1,5 @@
 import datetime
+import uuid as py_uuid
 from collections.abc import Iterable, Sequence
 from typing import Protocol, Self
 
@@ -28,8 +29,10 @@ __all__ = [
     "ChangingModelMixinProtocol",
     "SnapshotModelMixinProtocol",
     "RevisionModelMixinProtocol",
+    "SQLiteUUID",
     "SQLiteDateTime",
     "model_sqn_type",
+    "model_uuid_type",
     "model_datetime_tz_type",
     "model_revision_type",
     "SequenceModelMixin",
@@ -395,11 +398,29 @@ class RevisionModelMixinProtocol(SequenceModelMixinProtocol):
         ...
 
 
+class SQLiteUUID(sa.TypeDecorator):
+    """
+    Custom SQLAlchemy type decorator for handling UUIDs in SQLite.
+    """
+    impl = sa.CHAR(36)
+    cache_ok = True
+
+    def process_bind_param(self, value: py_uuid.UUID | None, dialect) -> str | None:
+        if value is None:
+            return None
+        return str(value)
+
+    def process_result_value(self, value: str | None, dialect) -> py_uuid.UUID | None:
+        if value is None:
+            return None
+        return py_uuid.UUID(value)
+
+
 class SQLiteDateTime(sa.TypeDecorator):
     """
     Custom SQLAlchemy type decorator for handling timezone-aware datetimes in SQLite.
     """
-    impl = sa.DateTime
+    impl = sa.DATETIME
     cache_ok = True
 
     def process_bind_param(self, value: datetime.datetime | None, dialect) -> datetime.datetime | None:
@@ -433,6 +454,22 @@ def model_sqn_type(dialect: str | None = None) -> sa.types.TypeEngine[int]:
     raise ValueError(f"unsupported database dialect '{dialect}'")
 
 
+def model_uuid_type(dialect: str | None = None) -> sa.types.TypeEngine[py_uuid.UUID]:
+    """
+    Returns the appropriate SQLAlchemy column type for a UUID based on the specified database dialect.
+
+    :param dialect: The database dialect to determine the column type for. If None, defaults to PostgreSQL.
+    :return: The SQLAlchemy column type to use for the UUID field.
+    """
+    if dialect is None:
+        dialect = dialects.postgresql
+    if dialect == dialects.postgresql:
+        return sa_pg.UUID(as_uuid=True)
+    if dialect == dialects.sqlite:
+        return SQLiteUUID
+    raise ValueError(f"unsupported database dialect '{dialect}'")
+
+
 def model_datetime_tz_type(dialect: str | None = None) -> sa.types.TypeEngine[datetime.datetime]:
     """
     Returns the appropriate SQLAlchemy column type for a timezone-aware datetime field based on the specified database

{plexus_python_pulse_common-1.1.76 → plexus_python_pulse_common-1.1.77}/src/plexus_python_pulse_common.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: plexus-python-pulse-common
-Version: 1.1.76
+Version: 1.1.77
 Classifier: Programming Language :: Python :: 3
 Classifier: Programming Language :: Python :: 3.12
 Classifier: Programming Language :: Python :: 3.13

{plexus_python_pulse_common-1.1.76 → plexus_python_pulse_common-1.1.77}/src/plexus_python_pulse_common.egg-info/SOURCES.txt

@@ -54,6 +54,7 @@ src/plexus/pulse/common/utils/bagutils.py
 src/plexus/pulse/common/utils/config.py
 src/plexus/pulse/common/utils/datautils.py
 src/plexus/pulse/common/utils/dockerutils.py
+src/plexus/pulse/common/utils/flowutils.py
 src/plexus/pulse/common/utils/gisutils.py
 src/plexus/pulse/common/utils/jsonutils.py
 src/plexus/pulse/common/utils/ormutils.py
@@ -80,6 +81,7 @@ test/plexus_tests/pulse/common/utils/__init__.py
 test/plexus_tests/pulse/common/utils/bagutils_test.py
 test/plexus_tests/pulse/common/utils/datautils_test.py
 test/plexus_tests/pulse/common/utils/dockerutils_test.py
+test/plexus_tests/pulse/common/utils/flowutils_test.py
 test/plexus_tests/pulse/common/utils/gisutils_test.py
 test/plexus_tests/pulse/common/utils/jsonutils_test.py
 test/plexus_tests/pulse/common/utils/ormutils_test.py

plexus_python_pulse_common-1.1.77/test/plexus_tests/pulse/common/utils/flowutils_test.py

@@ -0,0 +1,36 @@
+import os
+import tempfile
+import unittest
+import uuid as pyuuid
+from unittest.mock import patch
+
+import ddt
+
+from plexus.pulse.common.utils.flowutils import FlowContext
+
+
+@ddt.ddt
+class FlowUtilsTest(unittest.TestCase):
+
+    def test_flow_context(self):
+        with tempfile.TemporaryDirectory(delete=True) as temp_dir:
+            with patch.dict(os.environ,
+                            {
+                                "PULSE_FLOW_CONTEXT_ROOT_DIR": temp_dir,
+                                "PULSE_FLOW_RUN_ID": "01234567-89ab-cdef-0123-456789abcdef",
+                            },
+                            clear=True):
+                context = (
+                    FlowContext()
+                    .add_dataset_ref(pyuuid.UUID("01234567-89ab-cdef-0123-456789abcdef"), "test_dataset")
+                    .add_dataset_ref(pyuuid.UUID("89abcdef-0123-4567-89ab-cdef01234567"), "another_test_dataset")
+                    .add_dataset_ref(pyuuid.UUID("00000000-0000-0000-0000-000000000000"))
+                )
+
+                self.assertEqual(context.count_dataset_ref(), 3)
+
+                dataset, another_dataset, anonymous_dataset = context.list_dataset_refs()
+
+                self.assertEqual(dataset.dataset_uuid, pyuuid.UUID("01234567-89ab-cdef-0123-456789abcdef"))
+                self.assertEqual(another_dataset.dataset_uuid, pyuuid.UUID("89abcdef-0123-4567-89ab-cdef01234567"))
+                self.assertEqual(anonymous_dataset.dataset_uuid, pyuuid.UUID("00000000-0000-0000-0000-000000000000"))