plexus-python-pulse-common 1.1.76__py3-none-any.whl → 1.1.77__py3-none-any.whl

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/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.dist-info → plexus_python_pulse_common-1.1.77.dist-info}/METADATA

@@ -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.dist-info → plexus_python_pulse_common-1.1.77.dist-info}/RECORD

@@ -14,16 +14,17 @@ plexus/pulse/common/utils/bagutils.py,sha256=5JGnMgR8Q9OciW5dE7_2U7n1FGVEXMZXv1S
 plexus/pulse/common/utils/config.py,sha256=YPtFK5QHTVsXwcFBGV0rlZad0GFtsgwysQNyKYjdzvk,2409
 plexus/pulse/common/utils/datautils.py,sha256=OWlnyFSa6UaPbSSG3HzCJHxdiHyfHAE-FUVHaLfNZoM,9335
 plexus/pulse/common/utils/dockerutils.py,sha256=7DF9yFGz8uHwpYk115xBWqOjs5etz0UE7yh66eDvBzs,5793
+plexus/pulse/common/utils/flowutils.py,sha256=VmbVcCVztH8DH47s_5JbdOZchsEc7hZbq1am95amJzs,5307
 plexus/pulse/common/utils/gisutils.py,sha256=UX_IHPeAx9mbA9p_SsbPFwdsdp7bGeaQhWaCTCzvrBM,11466
 plexus/pulse/common/utils/jsonutils.py,sha256=zSbW94klN2jrkMdBECu4RNJX6zzNaSxauQ3FuihuB6E,3342
-plexus/pulse/common/utils/ormutils.py,sha256=MxSUIyxzGFY9wc12Y6QwaWZkbsYivSF5hNTyGmXFQYU,60683
+plexus/pulse/common/utils/ormutils.py,sha256=JD-hKLEGkFWy7ojc3dC_90qtexe2UnLEvD90tEOHtCg,61897
 plexus/pulse/common/utils/pathutils.py,sha256=S5xk2WEmxP8MiuSKIU3cAEtk4Nk5Lz7IqC5FCd8BEFI,9585
 plexus/pulse/common/utils/s3utils.py,sha256=RArl95rWoyBEFyWU9X0c8WgoQBjP_jOK3OkPoKxIMnA,36025
 plexus/pulse/common/utils/sqlutils.py,sha256=D6kTBjhO5YlNRt3uFlPt6z3uH61m9ajEzPYmsI6NoFc,231
 plexus/pulse/common/utils/strutils.py,sha256=PkmMHCcHN9vALCy_8fJTf00Dg7I2N4QTCXdBmB3t7rk,16440
 plexus/pulse/common/utils/tagutils.py,sha256=smM79dyRCiswKTTcI-At7FLxDH54VN4Id3aFHbDxzYo,61696
 plexus/pulse/common/utils/testutils.py,sha256=hKPRvoN4mg8QGQJH121l59a0uo8jN5rntGQAVYplGBw,6166
-plexus_python_pulse_common-1.1.76.dist-info/METADATA,sha256=8BRlnX_wIRMJ1zMz19hv6c-iZPe8koMARlaIet51bL8,1495
-plexus_python_pulse_common-1.1.76.dist-info/WHEEL,sha256=aeYiig01lYGDzBgS8HxWXOg3uV61G9ijOsup-k9o1sk,91
-plexus_python_pulse_common-1.1.76.dist-info/top_level.txt,sha256=ug_g7CVwaMQuas5UzAXbHUrQvKGCn8ezc6ZNvvRlJOE,7
-plexus_python_pulse_common-1.1.76.dist-info/RECORD,,
+plexus_python_pulse_common-1.1.77.dist-info/METADATA,sha256=Aei_-DjiA9AuGRzaDCZ5Vp3R8dJxcjXJSKFZ7SZwRHQ,1495
+plexus_python_pulse_common-1.1.77.dist-info/WHEEL,sha256=aeYiig01lYGDzBgS8HxWXOg3uV61G9ijOsup-k9o1sk,91
+plexus_python_pulse_common-1.1.77.dist-info/top_level.txt,sha256=ug_g7CVwaMQuas5UzAXbHUrQvKGCn8ezc6ZNvvRlJOE,7
+plexus_python_pulse_common-1.1.77.dist-info/RECORD,,