corvic-engine 0.3.0rc43__cp38-abi3-win_amd64.whl → 0.3.0rc45__cp38-abi3-win_amd64.whl

corvic/model/_proto_orm_convert.py

@@ -1,7 +1,6 @@
-import dataclasses
 import datetime
-from collections.abc import Sequence
-from typing import Generic, TypeVar, cast
+from collections.abc import Callable, Sequence
+from typing import Any, Protocol, TypeVar, cast
 
 import sqlalchemy as sa
 import sqlalchemy.orm as sa_orm
@@ -9,54 +8,64 @@ from google.protobuf import timestamp_pb2
 from more_itertools import flatten
 
 from corvic import orm
-from corvic.result import InvalidArgumentError, Ok
+from corvic.result import InternalError, InvalidArgumentError, Ok
 from corvic_generated.model.v1alpha import models_pb2
+from corvic_generated.orm.v1 import feature_view_pb2
 
 UNCOMMITTED_ID_PREFIX = "__uncommitted_object-"
 
-_Proto = (
-    models_pb2.Resource
-    | models_pb2.Source
-    | models_pb2.FeatureView
-    | models_pb2.Space
-    | models_pb2.FeatureViewSource
-    | models_pb2.Agent
-    | models_pb2.Pipeline
-    | models_pb2.Room
-    | models_pb2.CompletionModel
-)
-_ID = TypeVar(
-    "_ID",
-    orm.ResourceID,
-    orm.SourceID,
-    orm.FeatureViewID,
-    orm.SpaceID,
-    orm.FeatureViewSourceID,
-    orm.AgentID,
-    orm.PipelineID,
-    orm.RoomID,
-    orm.CompletionModelID,
+_Proto = TypeVar(
+    "_Proto",
+    models_pb2.Resource,
+    models_pb2.Source,
+    models_pb2.FeatureView,
+    models_pb2.Space,
+    models_pb2.FeatureViewSource,
+    models_pb2.Agent,
+    models_pb2.Pipeline,
+    models_pb2.Room,
+    models_pb2.CompletionModel,
 )
+ID = TypeVar("ID", bound=orm.BaseID[Any])
+
+
+class _OrmModel(Protocol[ID]):
+    id: sa_orm.Mapped[ID | None]
+
+    org_id: sa_orm.Mapped[orm.OrgID | None]
 
+    @sa.ext.hybrid.hybrid_property
+    def created_at(self) -> datetime.datetime | None: ...
 
-@dataclasses.dataclass
-class _OrmIDs(Generic[_ID]):
-    obj_id: _ID | None
-    room_id: orm.RoomID | None
+    @created_at.inplace.expression
+    @classmethod
+    def _created_at_expression(cls): ...
+
+
+OrmObj = TypeVar("OrmObj", bound=_OrmModel[Any])
+
+
+def _translate_orm_id(
+    obj_id: str, id_class: type[ID]
+) -> Ok[ID | None] | orm.InvalidORMIdentifierError:
+    if obj_id.startswith(UNCOMMITTED_ID_PREFIX):
+        return Ok(None)
+    parsed_obj_id = id_class(obj_id)
+    match parsed_obj_id.to_db():
+        case orm.InvalidORMIdentifierError() as err:
+            return err
+        case Ok():
+            return Ok(parsed_obj_id)
 
 
 def _translate_orm_ids(
-    proto_obj: _Proto, obj_id_class: type[_ID]
-) -> Ok[_OrmIDs[_ID]] | orm.InvalidORMIdentifierError:
-    if proto_obj.id.startswith(UNCOMMITTED_ID_PREFIX):
-        obj_id = None
-    else:
-        obj_id = obj_id_class(proto_obj.id)
-        match obj_id.to_db():
-            case orm.InvalidORMIdentifierError() as err:
-                return err
-            case Ok():
-                pass
+    proto_obj: _Proto, obj_id_class: type[ID]
+) -> Ok[tuple[ID | None, orm.RoomID | None]] | orm.InvalidORMIdentifierError:
+    match _translate_orm_id(proto_obj.id, obj_id_class):
+        case orm.InvalidORMIdentifierError() as err:
+            return err
+        case Ok(obj_id):
+            pass
 
     match proto_obj:
         case (
@@ -66,6 +75,7 @@ def _translate_orm_ids(
             | models_pb2.Space()
             | models_pb2.Agent()
             | models_pb2.Pipeline()
+            | models_pb2.FeatureViewSource()
         ):
             room_id = orm.RoomID(proto_obj.room_id)
             match room_id.to_db():
@@ -73,12 +83,12 @@ def _translate_orm_ids(
                     return err
                 case Ok():
                     pass
-        case models_pb2.FeatureViewSource() | models_pb2.CompletionModel():
+        case models_pb2.CompletionModel():
             room_id = None
         case models_pb2.Room():
             room_id = cast(orm.RoomID, obj_id)
 
-    return Ok(_OrmIDs(obj_id, room_id))
+    return Ok((obj_id, room_id))
 
 
 def timestamp_orm_to_proto(
@@ -103,12 +113,11 @@ def resource_orm_to_proto(resource_orm: orm.Resource) -> models_pb2.Resource:
         size=resource_orm.size,
         original_path=resource_orm.original_path,
         room_id=str(resource_orm.room_id),
-        source_ids=[str(val.source_id) for val in resource_orm.source_associations],
         org_id=str(resource_orm.org_id),
         recent_events=[resource_orm.latest_event] if resource_orm.latest_event else [],
-        pipeline_id=str(resource_orm.pipeline_input_refs[-1].pipeline.id)
-        if resource_orm.pipeline_input_refs
-        else None,
+        pipeline_id=str(resource_orm.pipeline_ref.pipeline_id)
+        if resource_orm.pipeline_ref
+        else "",
         created_at=timestamp_orm_to_proto(resource_orm.created_at),
     )
 
@@ -119,13 +128,10 @@ def source_orm_to_proto(source_orm: orm.Source) -> models_pb2.Source:
         name=source_orm.name,
         table_op_graph=source_orm.table_op_graph,
         room_id=str(source_orm.room_id),
-        resource_id=str(source_orm.resource_associations[0].resource_id)
-        if source_orm.resource_associations
-        else "",
         org_id=str(source_orm.org_id),
-        pipeline_id=str(source_orm.pipeline_output_refs[-1].pipeline.id)
-        if source_orm.pipeline_output_refs
-        else None,
+        pipeline_id=str(source_orm.pipeline_ref.pipeline_id)
+        if source_orm.pipeline_ref
+        else "",
         created_at=timestamp_orm_to_proto(source_orm.created_at),
     )
 
@@ -144,10 +150,17 @@ def agent_orm_to_proto(agent_orm: orm.Agent) -> models_pb2.Agent:
 def feature_view_source_orm_to_proto(
     feature_view_source_orm: orm.FeatureViewSource,
 ) -> models_pb2.FeatureViewSource:
+    if feature_view_source_orm.table_op_graph.WhichOneof("op") is not None:
+        op = feature_view_source_orm.table_op_graph
+    else:
+        # some legacy feature views were stored without op graphs fill those
+        # with the source's opgraph
+        op = feature_view_source_orm.source.table_op_graph
     return models_pb2.FeatureViewSource(
         id=str(feature_view_source_orm.id),
+        room_id=str(feature_view_source_orm.room_id),
         source=source_orm_to_proto(feature_view_source_orm.source),
-        table_op_graph=feature_view_source_orm.table_op_graph,
+        table_op_graph=op,
         drop_disconnected=feature_view_source_orm.drop_disconnected,
         org_id=str(feature_view_source_orm.org_id),
         created_at=timestamp_orm_to_proto(feature_view_source_orm.created_at),
@@ -201,7 +214,7 @@ def space_orm_to_proto(space_orm: orm.Space) -> models_pb2.Space:
         description=space_orm.description,
         room_id=str(space_orm.room_id),
         space_parameters=space_orm.parameters,
-        feature_view_id=str(space_orm.feature_view_id),
+        feature_view=feature_view_orm_to_proto(space_orm.feature_view),
         auto_sync=space_orm.auto_sync if space_orm.auto_sync is not None else False,
         org_id=str(space_orm.org_id),
         created_at=timestamp_orm_to_proto(space_orm.created_at),
@@ -231,27 +244,31 @@ def completion_model_orm_to_proto(
     )
 
 
+def _add_orm_to_session(
+    orm_obj: OrmObj, org_id: str, session: sa_orm.Session
+) -> OrmObj:
+    if org_id:
+        orm_obj.org_id = orm.OrgID(org_id)
+    if not orm_obj.id:
+        session.add(orm_obj)
+    else:
+        orm_obj = session.merge(orm_obj)
+    return orm_obj
+
+
 def resource_proto_to_orm(
     proto_obj: models_pb2.Resource, session: sa_orm.Session
 ) -> Ok[orm.Resource] | InvalidArgumentError:
     match _translate_orm_ids(proto_obj, orm.ResourceID):
         case orm.InvalidORMIdentifierError() as err:
             return err
-        case Ok(ids):
+        case Ok((obj_id, room_id)):
             pass
-    if not ids.room_id:
-        return InvalidArgumentError("room id required to commit resource")
-
-    source_ids = list[orm.SourceID]()
-    for source_id in proto_obj.source_ids:
-        orm_id = orm.SourceID(source_id)
-        match orm_id.to_db():
-            case orm.InvalidORMIdentifierError() as err:
-                return err
-            case Ok():
-                source_ids.append(orm_id)
+    if not room_id:
+        return InvalidArgumentError("room id required to commit")
+
     orm_obj = orm.Resource(
-        id=ids.obj_id,
+        id=obj_id,
         name=proto_obj.name,
         description=proto_obj.description,
         mime_type=proto_obj.mime_type,
@@ -260,23 +277,31 @@ def resource_proto_to_orm(
         size=proto_obj.size,
         original_path=proto_obj.original_path,
         latest_event=proto_obj.recent_events[-1] if proto_obj.recent_events else None,
-        room_id=ids.room_id,
-        source_associations=[
-            orm.SourceResourceAssociation(source_id=src_id, resource_id=ids.obj_id)
-            for src_id in source_ids
-        ],
+        room_id=room_id,
     )
-    if proto_obj.org_id:
-        orm_obj.org_id = orm.OrgID(proto_obj.org_id)
-        for assn in orm_obj.source_associations:
-            assn.org_id = orm_obj.org_id
+    return Ok(_add_orm_to_session(orm_obj, proto_obj.org_id, session))
 
-    if ids.obj_id:
-        orm_obj = session.merge(orm_obj)
-    else:
-        session.add(orm_obj)
 
-    return Ok(orm_obj)
+def _ensure_id(
+    proto_obj: _Proto,
+    proto_to_orm: Callable[[_Proto, sa_orm.Session], Ok[OrmObj] | InvalidArgumentError],
+    id_type: type[ID],
+    session: sa_orm.Session,
+) -> Ok[ID] | orm.InvalidORMIdentifierError | InvalidArgumentError:
+    match _translate_orm_id(proto_obj.id, id_type):
+        case orm.InvalidORMIdentifierError() as err:
+            return err
+        case Ok(orm_id):
+            if orm_id:
+                return Ok(orm_id)
+    match proto_to_orm(proto_obj, session):
+        case orm.InvalidORMIdentifierError() | InvalidArgumentError() as err:
+            return err
+        case Ok(orm_obj):
+            session.flush()
+            if not orm_obj.id:
+                raise InternalError("internal assertion did not hold")
+            return Ok(orm_obj.id)
 
 
 def pipeline_proto_to_orm(  # noqa: C901
@@ -285,60 +310,59 @@ def pipeline_proto_to_orm(  # noqa: C901
     match _translate_orm_ids(proto_obj, orm.PipelineID):
         case orm.InvalidORMIdentifierError() as err:
             return err
-        case Ok(ids):
+        case Ok((obj_id, room_id)):
             pass
-    if not ids.room_id:
-        return InvalidArgumentError("room id required to commit pipeline")
+    if not room_id:
+        return InvalidArgumentError("room id required to commit")
 
+    inputs = list[orm.PipelineInput]()
     orm_obj = orm.Pipeline(
-        id=ids.obj_id,
+        id=obj_id,
         name=proto_obj.name,
-        room_id=ids.room_id,
+        room_id=room_id,
         transformation=proto_obj.pipeline_transformation,
         description=proto_obj.description,
     )
+    orm_obj = _add_orm_to_session(orm_obj, proto_obj.org_id, session)
+    session.flush()
 
-    if ids.obj_id:
-        orm_obj = session.merge(orm_obj)
-    else:
-        session.add(orm_obj)
+    if not orm_obj.id:
+        raise InternalError("internal assertion did not hold")
 
     for name, val in proto_obj.resource_inputs.items():
-        if any(input.name == name for input in orm_obj.inputs):
-            continue
-        match resource_proto_to_orm(val, session):
+        match _ensure_id(val, resource_proto_to_orm, orm.ResourceID, session):
             case orm.InvalidORMIdentifierError() | InvalidArgumentError() as err:
                 return err
-            case Ok(resource_orm):
-                if resource_orm.id:
-                    resource_orm = session.merge(resource_orm)
-                else:
-                    session.add(resource_orm)
-                session.merge(
-                    orm.PipelineInput(
-                        pipeline=orm_obj, resource=resource_orm, name=name
-                    )
-                )
+            case Ok(resource_id):
+                pass
+        inputs.append(
+            orm.PipelineInput(
+                resource_id=resource_id,
+                name=name,
+                pipeline_id=orm_obj.id,
+                room_id=room_id,
+            )
+        )
 
+    outputs = list[orm.PipelineOutput]()
     for name, val in proto_obj.source_outputs.items():
-        if any(output.name == name for output in orm_obj.outputs):
-            continue
-        match source_proto_to_orm(val, session):
+        match _ensure_id(val, source_proto_to_orm, orm.SourceID, session):
             case orm.InvalidORMIdentifierError() | InvalidArgumentError() as err:
                 return err
-            case Ok(source_orm):
-                if source_orm.id:
-                    source_orm = session.merge(source_orm)
-                else:
-                    session.add(source_orm)
-                session.merge(
-                    orm.PipelineOutput(pipeline=orm_obj, source=source_orm, name=name)
-                )
+            case Ok(source_id):
+                pass
+        outputs.append(
+            orm.PipelineOutput(
+                source_id=source_id, name=name, pipeline_id=orm_obj.id, room_id=room_id
+            )
+        )
+
     if proto_obj.org_id:
         org_id = orm.OrgID(proto_obj.org_id)
-        orm_obj.org_id = org_id
-        for obj in flatten((orm_obj.inputs, orm_obj.outputs)):
-            obj.org_id = orm.OrgID(proto_obj.org_id)
+        for obj in flatten((inputs, outputs)):
+            obj.org_id = org_id
+    for obj in flatten((inputs, outputs)):
+        session.merge(obj)
     return Ok(orm_obj)
 
 
@@ -348,35 +372,18 @@ def source_proto_to_orm(
     match _translate_orm_ids(proto_obj, orm.SourceID):
         case orm.InvalidORMIdentifierError() as err:
             return err
-        case Ok(ids):
+        case Ok((obj_id, room_id)):
             pass
-    if not ids.room_id:
-        return InvalidArgumentError("room id required to commit resource")
-    resource_id = orm.ResourceID(proto_obj.resource_id)
-    if resource_id:
-        associations = [
-            orm.SourceResourceAssociation(source_id=ids.obj_id, resource_id=resource_id)
-        ]
-    else:
-        associations = list[orm.SourceResourceAssociation]()
+    if not room_id:
+        return InvalidArgumentError("room id required to commit")
 
     orm_obj = orm.Source(
-        id=ids.obj_id,
+        id=obj_id,
         name=proto_obj.name,
         table_op_graph=proto_obj.table_op_graph,
-        room_id=ids.room_id,
-        resource_associations=associations,
+        room_id=room_id,
     )
-    if proto_obj.org_id:
-        orm_obj.org_id = orm.OrgID(proto_obj.org_id)
-        for assn in orm_obj.resource_associations:
-            assn.org_id = orm.OrgID(proto_obj.org_id)
-    if ids.obj_id:
-        orm_obj = session.merge(orm_obj)
-    else:
-        session.add(orm_obj)
-
-    return Ok(orm_obj)
+    return Ok(_add_orm_to_session(orm_obj, proto_obj.org_id, session))
 
 
 def agent_proto_to_orm(
@@ -385,26 +392,18 @@ def agent_proto_to_orm(
     match _translate_orm_ids(proto_obj, orm.AgentID):
         case orm.InvalidORMIdentifierError() as err:
             return err
-        case Ok(ids):
+        case Ok((obj_id, room_id)):
             pass
-    if not ids.room_id:
-        return InvalidArgumentError("room id required to commit resource")
+    if not room_id:
+        return InvalidArgumentError("room id required to commit")
 
     orm_obj = orm.Agent(
-        id=ids.obj_id,
+        id=obj_id,
         name=proto_obj.name,
         parameters=proto_obj.agent_parameters,
-        room_id=ids.room_id,
+        room_id=room_id,
     )
-
-    if proto_obj.org_id:
-        orm_obj.org_id = orm.OrgID(proto_obj.org_id)
-    if ids.obj_id:
-        orm_obj = session.merge(orm_obj)
-    else:
-        session.add(orm_obj)
-
-    return Ok(orm_obj)
+    return Ok(_add_orm_to_session(orm_obj, proto_obj.org_id, session))
 
 
 def space_proto_to_orm(
@@ -413,98 +412,129 @@ def space_proto_to_orm(
     match _translate_orm_ids(proto_obj, orm.SpaceID):
         case orm.InvalidORMIdentifierError() as err:
             return err
-        case Ok(ids):
+        case Ok((obj_id, room_id)):
             pass
-    if not ids.room_id:
-        return InvalidArgumentError("room id required to commit resource")
+    if not room_id:
+        return InvalidArgumentError("room id required to commit")
 
-    feature_view_id = orm.FeatureViewID(proto_obj.feature_view_id)
-    feature_view = session.get(orm.FeatureView, feature_view_id)
-    if not feature_view:
-        return InvalidArgumentError("feature view required to commit resource")
+    match _ensure_id(
+        proto_obj.feature_view, feature_view_proto_to_orm, orm.FeatureViewID, session
+    ):
+        case orm.InvalidORMIdentifierError() | InvalidArgumentError() as err:
+            return err
+        case Ok(feature_view_id):
+            pass
+
+    if not feature_view_id:
+        raise InternalError("internal assertion did not hold")
 
     orm_obj = orm.Space(
-        id=ids.obj_id,
+        id=obj_id,
         name=proto_obj.name,
         description=proto_obj.description,
-        room_id=ids.room_id,
+        room_id=room_id,
         feature_view_id=feature_view_id,
         parameters=proto_obj.space_parameters,
         auto_sync=proto_obj.auto_sync,
-        feature_view=feature_view,
     )
-    if proto_obj.org_id:
-        orm_obj.org_id = orm.OrgID(proto_obj.org_id)
-
-    if ids.obj_id:
-        orm_obj = session.merge(orm_obj)
-    else:
-        session.add(orm_obj)
-    return Ok(orm_obj)
+    return Ok(_add_orm_to_session(orm_obj, proto_obj.org_id, session))
 
 
-def feature_view_proto_to_orm(
+def feature_view_proto_to_orm(  # noqa: C901
     proto_obj: models_pb2.FeatureView, session: sa_orm.Session
 ) -> Ok[orm.FeatureView] | orm.InvalidORMIdentifierError | InvalidArgumentError:
     match _translate_orm_ids(proto_obj, orm.FeatureViewID):
         case orm.InvalidORMIdentifierError() as err:
             return err
-        case Ok(ids):
+        case Ok((obj_id, room_id)):
             pass
-    if not ids.room_id:
-        return InvalidArgumentError("room id required to commit resource")
-
-    feature_view_sources = list[orm.FeatureViewSource]()
-    for fvs in proto_obj.feature_view_sources:
-        match feature_view_source_proto_to_orm(fvs, session):
-            case orm.InvalidORMIdentifierError() | InvalidArgumentError() as err:
-                return err
-            case Ok(orm_fvs):
-                if orm_fvs.id:
-                    orm_fvs = session.merge(orm_fvs)
-                else:
-                    session.add(orm_fvs)
-                feature_view_sources.append(orm_fvs)
+    if not room_id:
+        return InvalidArgumentError("room id required to commit")
 
     orm_obj = orm.FeatureView(
-        id=ids.obj_id,
+        id=obj_id,
         name=proto_obj.name,
         description=proto_obj.description,
-        room_id=ids.room_id,
-        feature_view_output=proto_obj.feature_view_output,
-        feature_view_sources=feature_view_sources,
+        room_id=room_id,
     )
     if proto_obj.org_id:
         orm_obj.org_id = orm.OrgID(proto_obj.org_id)
 
-    if ids.obj_id:
-        orm_obj = session.merge(orm_obj)
-    else:
-        session.add(orm_obj)
+    orm_obj = _add_orm_to_session(orm_obj, proto_obj.org_id, session)
+    session.flush()
+
+    if not orm_obj.id:
+        raise InternalError("internal assertion did not hold")
+
+    new_fv_sources = list[orm.FeatureViewSource]()
+    for fvs in proto_obj.feature_view_sources:
+        match _feature_view_source_proto_to_orm(fvs, orm_obj.id, session):
+            case orm.InvalidORMIdentifierError() | InvalidArgumentError() as err:
+                return err
+            case Ok(fvs_orm):
+                new_fv_sources.append(fvs_orm)
+
+    old_to_new_source_id = dict[str, str]()
+    for old_fvs, new_fvs in zip(
+        proto_obj.feature_view_sources, new_fv_sources, strict=True
+    ):
+        if old_fvs.source.id != str(new_fvs.source_id):
+            old_to_new_source_id[old_fvs.source.id] = str(new_fvs.source_id)
+
+    orm_obj.feature_view_output = (
+        feature_view_pb2.FeatureViewOutput(
+            relationships=[
+                feature_view_pb2.FeatureViewRelationship(
+                    start_source_id=old_to_new_source_id.get(
+                        old_rel.start_source_id, old_rel.start_source_id
+                    ),
+                    end_source_id=old_to_new_source_id.get(
+                        old_rel.end_source_id, old_rel.end_source_id
+                    ),
+                )
+                for old_rel in proto_obj.feature_view_output.relationships
+            ],
+            output_sources=[
+                feature_view_pb2.OutputSource(
+                    source_id=old_to_new_source_id.get(
+                        old_output_source.source_id, old_output_source.source_id
+                    )
+                )
+                for old_output_source in proto_obj.feature_view_output.output_sources
+            ],
+        )
+        if old_to_new_source_id
+        else proto_obj.feature_view_output
+    )
+    orm_obj = session.merge(orm_obj)
     return Ok(orm_obj)
 
 
-def feature_view_source_proto_to_orm(
-    proto_obj: models_pb2.FeatureViewSource, session: sa_orm.Session
+def _feature_view_source_proto_to_orm(
+    proto_obj: models_pb2.FeatureViewSource,
+    feature_view_id: orm.FeatureViewID,
+    session: sa_orm.Session,
 ) -> Ok[orm.FeatureViewSource] | orm.InvalidORMIdentifierError | InvalidArgumentError:
-    match source_proto_to_orm(proto_obj.source, session):
+    match _ensure_id(proto_obj.source, source_proto_to_orm, orm.SourceID, session):
         case orm.InvalidORMIdentifierError() | InvalidArgumentError() as err:
             return err
-        case Ok(source):
+        case Ok(source_id):
+            pass
+    match _translate_orm_id(proto_obj.id, orm.FeatureViewSourceID):
+        case orm.InvalidORMIdentifierError() as err:
+            return err
+        case Ok(obj_id):
             pass
 
-    if source.id:
-        source = session.merge(source)
-    else:
-        session.add(source)
     orm_obj = orm.FeatureViewSource(
+        room_id=orm.RoomID(proto_obj.room_id),
+        id=obj_id,
         table_op_graph=proto_obj.table_op_graph,
         drop_disconnected=proto_obj.drop_disconnected,
-        source=source,
+        source_id=source_id,
+        feature_view_id=feature_view_id,
     )
-    if proto_obj.org_id:
-        orm_obj.org_id = orm.OrgID(proto_obj.org_id)
-    return Ok(orm_obj)
+    return Ok(_add_orm_to_session(orm_obj, proto_obj.org_id, session))
 
 
 def room_proto_to_orm(
@@ -513,20 +543,14 @@ def room_proto_to_orm(
     match _translate_orm_ids(proto_obj, orm.RoomID):
         case orm.InvalidORMIdentifierError() as err:
             return err
-        case Ok(ids):
+        case Ok((obj_id, _)):
             pass
 
     orm_obj = orm.Room(
-        id=ids.obj_id,
+        id=obj_id,
         name=proto_obj.name,
     )
-
-    if ids.obj_id:
-        orm_obj = session.merge(orm_obj)
-    else:
-        session.add(orm_obj)
-
-    return Ok(orm_obj)
+    return Ok(_add_orm_to_session(orm_obj, proto_obj.org_id, session))
 
 
 def completion_model_proto_to_orm(
@@ -535,23 +559,17 @@ def completion_model_proto_to_orm(
     match _translate_orm_ids(proto_obj, orm.CompletionModelID):
         case orm.InvalidORMIdentifierError() as err:
             return err
-        case Ok(ids):
+        case Ok((obj_id, _)):
             pass
 
     orm_obj = orm.CompletionModel(
-        id=ids.obj_id,
+        id=obj_id,
         name=proto_obj.name,
         description=proto_obj.description,
         parameters=proto_obj.parameters,
         secret_api_key=proto_obj.secret_api_key,
     )
-
-    if ids.obj_id:
-        orm_obj = session.merge(orm_obj)
-    else:
-        session.add(orm_obj)
-
-    return Ok(orm_obj)
+    return Ok(_add_orm_to_session(orm_obj, proto_obj.org_id, session))
 
 
 def source_delete_orms(