Use ogsi model #32

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open

AlexAxthelm wants to merge 56 commits into main from use-ogsi-model

deployments/api/pyproject.toml

-Original file line number
+Diff line change
@@ Expand Up / @@ -12,6 +12,8 @@ dependencies = [ @@
       "pydantic-settings>=2.12.0",
       "sqlalchemy>=2.0.44",
       "stitch-auth",
+      "stitch-models",
+      "stitch-ogsi",
     ]
     [project.scripts]
@@ Expand Down Expand Up / @@ -41,3 +43,5 @@ addopts = ["-v", "--strict-markers", "--tb=short"] @@
     [tool.uv.sources]
     stitch-auth = { workspace = true }
+    stitch-models = { workspace = true }
+    stitch-ogsi = { workspace = true }

deployments/api/src/stitch/api/db/errors.py

-Original file line number
+Diff line change
@@ -0,0 +1,13 @@
+    from stitch.api.errors import StitchAPIError
+    class ResourceNotFoundError(StitchAPIError): ...
+    class ResourceIntegrityError(StitchAPIError): ...
+    class SourceNotFoundError(StitchAPIError): ...
+    class SourceIntegrityError(StitchAPIError): ...

deployments/api/src/stitch/api/db/init_job.py

-Original file line number
+Diff line change
@@ Expand Up / @@ -5,29 +5,26 @@ @@
     import time
     from enum import Enum
     from dataclasses import dataclass
-    from typing import Iterable
+    from typing import Any
     from sqlalchemy import create_engine, inspect, text
     from sqlalchemy.exc import OperationalError
     from sqlalchemy.orm import Session
     from stitch.api.db.model import (
-        CCReservoirsSourceModel,
-        GemSourceModel,
-        MembershipModel,
-        RMIManualSourceModel,
         ResourceModel,
+        MembershipModel,
         StitchBase,
         UserModel,
-        WMSourceModel,
+        OilGasFieldSourceModel,
     )
     from stitch.api.entities import (
-        GemData,
-        RMIManualData,
         User as UserEntity,
-        WMData,
     )
+    # Domain model from stitch-ogsi package
+    from stitch.ogsi.model.og_field import OilGasFieldBase
     """
     DB init/seed job.
@@ Expand Down Expand Up @@
     def create_seed_user() -> UserModel:
         return UserModel(
-            id=1,
             sub="seed|system",
             name="Seed User",
             email="seed@example.com",
@@ Expand All / @@ -266,97 +262,79 @@ def create_seed_user() -> UserModel: @@
     def create_dev_user() -> UserModel:
         return UserModel(
-            id=2,
             sub="dev|local-placeholder",
             name="Dev Deverson",
             email="dev@example.com",
         )
-    def create_seed_sources():
-        gem_sources = [
-            GemSourceModel.from_entity(
-                GemData(name="Permian Basin Field", country="USA", lat=31.8, lon=-102.3)
-            ),
-            GemSourceModel.from_entity(
-                GemData(name="North Sea Platform", country="GBR", lat=57.5, lon=1.5)
-            ),
-        ]
-        for i, src in enumerate(gem_sources, start=1):
-            src.id = i
-        wm_sources = [
-            WMSourceModel.from_entity(
-                WMData(
-                    field_name="Eagle Ford Shale", field_country="USA", production=125000.5
-                )
-            ),
-            WMSourceModel.from_entity(
-                WMData(field_name="Ghawar Field", field_country="SAU", production=500000.0)
-            ),
-        ]
-        for i, src in enumerate(wm_sources, start=1):
-            src.id = i
-        rmi_sources = [
-            RMIManualSourceModel.from_entity(
-                RMIManualData(
-                    name_override="Custom Override Name",
-                    gwp=25.5,
-                    gor=0.45,
-                    country="CAN",
-                    latitude=56.7,
-                    longitude=-111.4,
-                )
-            ),
-        ]
-        for i, src in enumerate(rmi_sources, start=1):
-            src.id = i
-        # CC Reservoir sources are intentionally omitted from the dev seed profile;
-        # the CCReservoirsSourceModel table is still created from SQLAlchemy metadata.
-        cc_sources: list[CCReservoirsSourceModel] = []
-        return gem_sources, wm_sources, rmi_sources, cc_sources
     def create_seed_resources(user: UserEntity) -> list[ResourceModel]:
         resources = [
-            ResourceModel.create(user, name="Multi-Source Asset", country="USA"),
-            ResourceModel.create(user, name="Single Source Asset", country="GBR"),
+            ResourceModel.create(user, name="Resource Foo01"),
+            ResourceModel.create(user, name="Resource Bar01"),
         ]
-        for i, res in enumerate(resources, start=1):
-            res.id = i
         return resources
     def create_seed_memberships(
         user: UserEntity,
         resources: list[ResourceModel],
-        gem_sources: list[GemSourceModel],
-        wm_sources: list[WMSourceModel],
-        rmi_sources: list[RMIManualSourceModel],
+        sources: list[OilGasFieldSourceModel],
     ) -> list[MembershipModel]:
         memberships = [
-            MembershipModel.create(user, resources[0], "gem", gem_sources[0].id),
-            MembershipModel.create(user, resources[0], "wm", wm_sources[0].id),
-            MembershipModel.create(user, resources[0], "rmi", rmi_sources[0].id),
-            MembershipModel.create(user, resources[1], "gem", gem_sources[1].id),
+            MembershipModel.create(user, resources[0], "gem", 1),
+            MembershipModel.create(user, resources[1], "wm", 2),
         ]
         for i, mem in enumerate(memberships, start=1):
             mem.id = i
         return memberships
-    def reset_sequences(engine, tables: Iterable[str]) -> None:
-        with engine.begin() as conn:
-            for t in tables:
-                conn.execute(
-                    text(
-                        f"SELECT setval('{t}_id_seq', "
-                        f"(SELECT COALESCE(MAX(id), 0) + 1 FROM {t}), false);"
-                    )
-                )
+    def create_seed_oil_gas_source_fields(
+        user: UserEntity,
+        resources: list[ResourceModel],
+    ) -> list[OilGasFieldSourceModel]:
+        """Create example OilGasField rows linked 1:1 with seeded resources."""
+        raw_payloads: list[dict[str, Any]] = [
+            # pretend this came from some upstream system (GEM/WM/etc)
+            {
+                "name": "Permian Alpha",
+                "country": "USA",
+                "basin": "Permian",
+                # extra keys demonstrate why we keep original_payload
+                "upstream_id": "seed-gem-0001",
+                "notes": "seed example",
+            },
+            {
+                "name": "North Sea Bravo",
+                "country": "GBR",
+                "basin": "North Sea",
+                "upstream_id": "seed-wm-0002",
+                "notes": "seed example",
+            },
+        ]
+        og_models: list[OilGasFieldSourceModel] = []
+        for resource, raw in zip(resources, raw_payloads):
+            domain = OilGasFieldBase.model_validate(raw)
+            model = OilGasFieldSourceModel(
+                created_by_id=user.id,
+                last_updated_by_id=user.id,
+            )
+            # Raw input (includes extra fields not in OilGasFieldBase)
+            model.original_payload = raw
+            # Canonical validated payload
+            model.payload = raw
+            model.name = domain.name
+            model.country = domain.country
+            model.basin = domain.basin
+            model.source = "dev-seed"
+            # Populate domain columns for queryability
+            og_models.append(model)
+        return og_models
     def seed_dev(engine) -> None:
@@ Expand All / @@ -376,39 +354,19 @@ def seed_dev(engine) -> None: @@
                 name=user_model.name,
             )
-            dev_entity = UserEntity(
-                id=dev_model.id,
-                sub=dev_model.sub,
-                email=dev_model.email,
-                name=dev_model.name,
-            )
-            gem_sources, wm_sources, rmi_sources, cc_sources = create_seed_sources()
-            session.add_all(gem_sources + wm_sources + rmi_sources + cc_sources)
             resources = create_seed_resources(user_entity)
-            resources = create_seed_resources(dev_entity)
             session.add_all(resources)
+            session.flush()
+            #
+            # Add sample OilGasField rows for the first two resources only
+            og_fields = create_seed_oil_gas_source_fields(user_entity, resources)
+            session.add_all(og_fields)
-            memberships = create_seed_memberships(
-                user_entity, resources, gem_sources, wm_sources, rmi_sources
-            )
+            memberships = create_seed_memberships(user_entity, resources, og_fields)
             session.add_all(memberships)
             session.commit()
-        reset_sequences(
-            engine,
-            tables=[
-                "users",
-                "gem_sources",
-                "wm_sources",
-                "rmi_manual_sources",
-                "resources",
-                "memberships",
-            ],
-        )
     def seed(engine, profile: SeedProfile | str) -> None:
         if profile == "dev":
@@ Expand Down @@

deployments/api/src/stitch/api/db/model/__init__.py

-Original file line number
+Diff line change
@@ -1,21 +1,13 @@
     from .common import Base as StitchBase
-    from .sources import (
-        GemSourceModel,
-        RMIManualSourceModel,
-        CCReservoirsSourceModel,
-        WMSourceModel,
-    )
+    from .oil_gas_field_source import OilGasFieldSourceModel
     from .resource import MembershipStatus, MembershipModel, ResourceModel
     from .user import User as UserModel
     __all__ = [
-        "CCReservoirsSourceModel",
-        "GemSourceModel",
         "MembershipModel",
         "MembershipStatus",
-        "RMIManualSourceModel",
         "ResourceModel",
         "StitchBase",
         "UserModel",
-        "WMSourceModel",
+        "OilGasFieldSourceModel",
     ]

deployments/api/src/stitch/api/db/model/mixins.py

-Original file line number
+Diff line change
@@ -1,12 +1,7 @@
     from datetime import datetime
-    from typing import Any, ClassVar, Generic, TypeVar, get_args, get_origin
-    from pydantic import TypeAdapter
-    from sqlalchemy import DateTime, ForeignKey, String, func
-    from sqlalchemy.ext.hybrid import hybrid_property
+    from sqlalchemy import DateTime, ForeignKey, func
     from sqlalchemy.orm import Mapped, declarative_mixin, mapped_column
-    from stitch.api.entities import SourceBase
-    from .types import StitchJson
     @declarative_mixin
@@ Expand All / @@ -28,45 +23,3 @@ class TimestampMixin: @@
     class UserAuditMixin:
         created_by_id: Mapped[int] = mapped_column(ForeignKey("users.id"))
         last_updated_by_id: Mapped[int] = mapped_column(ForeignKey("users.id"))
-    TPayload = TypeVar("TPayload", bound=SourceBase)
-    def _extract_payload_type(cls: type) -> type | None:
-        for base in getattr(cls, "__orig_bases__", []):
-            origin = get_origin(base)
-            if origin is PayloadMixin:
-                args = get_args(base)
-                if args:
-                    return args[0]
-    @declarative_mixin
-    class PayloadMixin(Generic[TPayload]):
-        __payload_adapter__: ClassVar[TypeAdapter]
-        source: Mapped[str] = mapped_column(String, nullable=False)  # "gem" | "woodmac"
-        _payload_data: Mapped[dict[str, Any]] = mapped_column(
-            "payload", StitchJson(), nullable=False
-        )
-        def __init_subclass__(cls, **kwargs: Any) -> None:
-            super().__init_subclass__(**kwargs)
-            payload_type = _extract_payload_type(cls)
-            if payload_type is not None:
-                cls.__payload_adapter__ = TypeAdapter(payload_type)
-        @hybrid_property
-        def payload(self) -> TPayload:
-            return self.__payload_adapter__.validate_python(self._payload_data)
-        @payload.inplace.setter
-        def _payload_setter(self, value: TPayload):
-            self.source = value.source
-            self._payload_data = value.model_dump(mode="json")
-        @payload.inplace.expression
-        @classmethod
-        def _payload_expression(cls):
-            return cls._payload_data

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Use ogsi model #32

Uh oh!

Diff view

Diff view

There are no files selected for viewing

Uh oh!

Uh oh!

Uh oh!

Use ogsi model #32

Are you sure you want to change the base?

Uh oh!

Use ogsi model #32

Uh oh!

Uh oh!

Diff view

Diff view

There are no files selected for viewing

Uh oh!

Uh oh!

Uh oh!