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

corvic/model/__init__.py

@@ -2,7 +2,7 @@
 
 import corvic.model._feature_type as feature_type
 from corvic.model._agent import Agent, AgentID
-from corvic.model._base_model import BaseModel
+from corvic.model._base_model import BaseModel, HasProtoSelf, UsesOrmID
 from corvic.model._completion_model import (
     CompletionModel,
     CompletionModelID,
@@ -47,7 +47,9 @@ from corvic.model._space import (
     SemanticSpace,
     Space,
     SpecificSpace,
+    SpecificSpaceParameters,
     TabularSpace,
+    UnknownSpace,
     embedding_model_proto_to_name,
     image_model_proto_to_name,
 )
@@ -70,6 +72,7 @@ __all__ = [
     "FeatureView",
     "FeatureViewEdgeTableMetadata",
     "FeatureViewRelationshipsMetadata",
+    "HasProtoSelf",
     "ImageSpace",
     "Node2VecParameters",
     "OcrPdfsPipeline",
@@ -87,9 +90,12 @@ __all__ = [
     "Space",
     "SpecificPipeline",
     "SpecificSpace",
+    "SpecificSpaceParameters",
     "TabularSpace",
     "UNCOMMITTED_ID_PREFIX",
     "UnknownTransformationPipeline",
+    "UnknownSpace",
+    "UsesOrmID",
     "embedding_model_proto_to_name",
     "feature_type",
     "image_model_proto_to_name",

corvic/model/_base_model.py

@@ -5,7 +5,7 @@ import datetime
 import functools
 import uuid
 from collections.abc import Callable, Iterable, Iterator, Sequence
-from typing import Any, Final, Generic, Protocol, TypeVar
+from typing import Final, Generic, Protocol, TypeVar
 
 import sqlalchemy as sa
 import sqlalchemy.orm as sa_orm
@@ -14,7 +14,7 @@ from google.protobuf import timestamp_pb2
 from typing_extensions import Self
 
 from corvic import orm, system
-from corvic.model._proto_orm_convert import UNCOMMITTED_ID_PREFIX
+from corvic.model._proto_orm_convert import ID, UNCOMMITTED_ID_PREFIX, OrmObj
 from corvic.result import InvalidArgumentError, NotFoundError, Ok
 
 _logger = structlog.get_logger()
@@ -26,21 +26,6 @@ class _ModelProto(Protocol):
 
 
 _ProtoObj = TypeVar("_ProtoObj", bound=_ModelProto)
-_ID = TypeVar("_ID", bound=orm.BaseID[Any])
-
-
-class _OrmModel(Protocol[_ID]):
-    id: sa_orm.Mapped[_ID | None]
-
-    @sa.ext.hybrid.hybrid_property
-    def created_at(self) -> datetime.datetime | None: ...
-
-    @created_at.inplace.expression
-    @classmethod
-    def _created_at_expression(cls): ...
-
-
-_OrmObj = TypeVar("_OrmObj", bound=_OrmModel[Any])
 
 
 def _generate_uncommitted_id_str():
@@ -58,55 +43,63 @@ def _create_or_join_session(
             yield session
 
 
-class BaseModel(Generic[_ID, _ProtoObj, _OrmObj], abc.ABC):
-    """Base for orm wrappers providing a unified update mechanism."""
-
+class HasProtoSelf(Generic[_ProtoObj], abc.ABC):
     client: Final[system.Client]
     proto_self: Final[_ProtoObj]
 
     def __init__(self, client: system.Client, proto_self: _ProtoObj):
-        if not proto_self.id:
-            proto_self.id = _generate_uncommitted_id_str()
         self.proto_self = proto_self
         self.client = client
 
+    @property
+    def created_at(self) -> datetime.datetime | None:
+        if self.proto_self.created_at:
+            return self.proto_self.created_at.ToDatetime(tzinfo=datetime.timezone.utc)
+        return None
+
+
+class UsesOrmID(Generic[ID, _ProtoObj], HasProtoSelf[_ProtoObj]):
+    def __init__(self, client: system.Client, proto_self: _ProtoObj):
+        if not proto_self.id:
+            proto_self.id = _generate_uncommitted_id_str()
+        super().__init__(client, proto_self)
+
     @classmethod
     @abc.abstractmethod
-    def orm_class(cls) -> type[_OrmObj]: ...
+    def id_class(cls) -> type[ID]: ...
+
+    @functools.cached_property
+    def id(self) -> ID:
+        return self.id_class().from_str(self.proto_self.id)
+
+
+class BaseModel(Generic[ID, _ProtoObj, OrmObj], UsesOrmID[ID, _ProtoObj]):
+    """Base for orm wrappers providing a unified update mechanism."""
 
     @classmethod
     @abc.abstractmethod
-    def id_class(cls) -> type[_ID]: ...
+    def orm_class(cls) -> type[OrmObj]: ...
 
     @classmethod
     @abc.abstractmethod
-    def orm_to_proto(cls, orm_obj: _OrmObj) -> _ProtoObj: ...
+    def orm_to_proto(cls, orm_obj: OrmObj) -> _ProtoObj: ...
 
     @classmethod
     @abc.abstractmethod
     def proto_to_orm(
         cls, proto_obj: _ProtoObj, session: orm.Session
-    ) -> Ok[_OrmObj] | InvalidArgumentError: ...
+    ) -> Ok[OrmObj] | InvalidArgumentError: ...
+
     @classmethod
     @abc.abstractmethod
     def delete_by_ids(
-        cls, ids: Sequence[_ID], session: orm.Session
+        cls, ids: Sequence[ID], session: orm.Session
     ) -> Ok[None] | InvalidArgumentError: ...
 
-    @functools.cached_property
-    def id(self) -> _ID:
-        return self.id_class().from_str(self.proto_self.id)
-
-    @property
-    def created_at(self) -> datetime.datetime | None:
-        if self.proto_self.created_at:
-            return self.proto_self.created_at.ToDatetime(tzinfo=datetime.timezone.utc)
-        return None
-
     @classmethod
     def load_proto_for(
         cls,
-        obj_id: _ID,
+        obj_id: ID,
         client: system.Client,
         existing_session: sa_orm.Session | None = None,
     ) -> Ok[_ProtoObj] | NotFoundError:
@@ -120,8 +113,8 @@ class BaseModel(Generic[_ID, _ProtoObj, _OrmObj], abc.ABC):
 
     @classmethod
     def _generate_query_results(
-        cls, query: sa.Select[tuple[_OrmObj]], session: sa_orm.Session
-    ) -> Iterator[_OrmObj]:
+        cls, query: sa.Select[tuple[OrmObj]], session: sa_orm.Session
+    ) -> Iterator[OrmObj]:
         it = iter(session.scalars(query))
         while True:
             try:
@@ -147,9 +140,9 @@ class BaseModel(Generic[_ID, _ProtoObj, _OrmObj], abc.ABC):
         limit: int | None = None,
         room_id: orm.RoomID | None = None,
         created_before: datetime.datetime | None = None,
-        ids: Iterable[_ID] | None = None,
+        ids: Iterable[ID] | None = None,
         additional_query_transform: Callable[
-            [sa.Select[tuple[_OrmObj]]], sa.Select[tuple[_OrmObj]]
+            [sa.Select[tuple[OrmObj]]], sa.Select[tuple[OrmObj]]
         ]
         | None = None,
         existing_session: sa_orm.Session | None = None,

corvic/model/_feature_view.py

@@ -18,15 +18,12 @@ from sqlalchemy.orm.interfaces import LoaderOption
 from typing_extensions import Self
 
 from corvic import op_graph, orm, system
-from corvic.model._base_model import BaseModel
+from corvic.model._base_model import BaseModel, UsesOrmID
 from corvic.model._defaults import Defaults
 from corvic.model._proto_orm_convert import (
     feature_view_delete_orms,
     feature_view_orm_to_proto,
     feature_view_proto_to_orm,
-    feature_view_source_delete_orms,
-    feature_view_source_orm_to_proto,
-    feature_view_source_proto_to_orm,
 )
 from corvic.model._source import Source, SourceID
 from corvic.result import InvalidArgumentError, NotFoundError, Ok
@@ -541,36 +538,31 @@ class RelationshipPath:
         return f"end-{uuid.uuid4()}"
 
 
-class FeatureViewSource(
-    BaseModel[FeatureViewSourceID, models_pb2.FeatureViewSource, orm.FeatureViewSource]
-):
+class FeatureViewSource(UsesOrmID[FeatureViewSourceID, models_pb2.FeatureViewSource]):
     """A table from a source with some extra operations defined by a feature view."""
 
-    @classmethod
-    def orm_class(cls):
-        return orm.FeatureViewSource
-
     @classmethod
     def id_class(cls):
         return FeatureViewSourceID
 
     @classmethod
-    def orm_to_proto(
-        cls, orm_obj: orm.FeatureViewSource
-    ) -> models_pb2.FeatureViewSource:
-        return feature_view_source_orm_to_proto(orm_obj)
-
-    @classmethod
-    def proto_to_orm(
-        cls, proto_obj: models_pb2.FeatureViewSource, session: orm.Session
-    ) -> Ok[orm.FeatureViewSource] | InvalidArgumentError:
-        return feature_view_source_proto_to_orm(proto_obj, session)
-
-    @classmethod
-    def delete_by_ids(
-        cls, ids: Sequence[FeatureViewSourceID], session: orm.Session
-    ) -> Ok[None] | InvalidArgumentError:
-        return feature_view_source_delete_orms(ids, session)
+    def create(
+        cls,
+        table: Table,
+        source: Source,
+        *,
+        drop_disconnected: bool,
+        room_id: orm.RoomID,
+    ):
+        return cls(
+            source.client,
+            models_pb2.FeatureViewSource(
+                source=source.proto_self,
+                table_op_graph=table.op_graph.to_proto(),
+                drop_disconnected=drop_disconnected,
+                room_id=str(room_id),
+            ),
+        )
 
     def with_table(self, table: Table | op_graph.Op) -> FeatureViewSource:
         if isinstance(table, Table):
@@ -713,12 +705,8 @@ class FeatureView(BaseModel[FeatureViewID, models_pb2.FeatureView, orm.FeatureVi
         return [
             sa_orm.selectinload(orm.FeatureView.feature_view_sources)
             .selectinload(orm.FeatureViewSource.source)
-            .selectinload(orm.Source.resource_associations)
-            .selectinload(orm.SourceResourceAssociation.resource),
-            sa_orm.selectinload(orm.FeatureView.feature_view_sources)
-            .selectinload(orm.FeatureViewSource.source)
-            .selectinload(orm.Source.pipeline_output_refs)
-            .selectinload(orm.PipelineOutput.pipeline),
+            .selectinload(orm.Source.pipeline_ref)
+            .selectinload(orm.PipelineOutput.source)
         ]
 
     @classmethod
@@ -775,12 +763,19 @@ class FeatureView(BaseModel[FeatureViewID, models_pb2.FeatureView, orm.FeatureVi
         ]
 
     @functools.cached_property
-    def output_sources(self) -> set[SourceID]:
+    def output_source_ids(self) -> set[SourceID]:
         return {
             SourceID(output_source.source_id)
             for output_source in self.proto_self.feature_view_output.output_sources
         }
 
+    @functools.cached_property
+    def output_sources(self) -> Sequence[Source]:
+        return [
+            self.source_id_to_feature_view_source[source_id].source
+            for source_id in self.output_source_ids
+        ]
+
     @functools.cached_property
     def source_id_to_feature_view_source(self) -> Mapping[SourceID, FeatureViewSource]:
         return {fvs.source.id: fvs for fvs in self.feature_view_sources}
@@ -846,7 +841,7 @@ class FeatureView(BaseModel[FeatureViewID, models_pb2.FeatureView, orm.FeatureVi
         paths_between_outputs = list(
             flatten(
                 self._calculate_paths_to_outputs(output, self.relationships)
-                for output in self.output_sources
+                for output in self.output_source_ids
             )
         )
         return [RelationshipPath(path) for path in paths_between_outputs]
@@ -857,7 +852,7 @@ class FeatureView(BaseModel[FeatureViewID, models_pb2.FeatureView, orm.FeatureVi
         paths: list[list[Relationship]] = []
         for rel in rels:
             if rel.start_source.id == output:
-                if rel.end_source.id in self.output_sources:
+                if rel.end_source.id in self.output_source_ids:
                     paths.append([rel])
                 else:
                     child_paths = self._calculate_paths_to_outputs(
@@ -944,7 +939,7 @@ class FeatureView(BaseModel[FeatureViewID, models_pb2.FeatureView, orm.FeatureVi
             )
         node_tables = {
             source_id: self.source_id_to_feature_view_source[source_id].table
-            for source_id in self.output_sources
+            for source_id in self.output_source_ids
         }
         node_id_columns = _validate_deepgnn_nodes(node_tables).unwrap_or_raise()
 
@@ -1034,6 +1029,11 @@ class FeatureView(BaseModel[FeatureViewID, models_pb2.FeatureView, orm.FeatureVi
         new_proto_self.name = name
         return self.__class__(self.client, proto_self=new_proto_self)
 
+    def with_description(self, description: str) -> Self:
+        new_proto_self = copy.deepcopy(self.proto_self)
+        new_proto_self.description = description
+        return self.__class__(self.client, proto_self=new_proto_self)
+
     def with_source(
         self,
         source: SourceID | Source,
@@ -1074,15 +1074,15 @@ class FeatureView(BaseModel[FeatureViewID, models_pb2.FeatureView, orm.FeatureVi
                 FeatureViewSourceColumnRenames(column_renames=renames)
             )
 
-        proto_feature_view_source = models_pb2.FeatureViewSource(
-            room_id=str(source.room_id),
-            table_op_graph=new_table.op_graph.to_proto(),
+        fvs = FeatureViewSource.create(
+            room_id=source.room_id,
+            table=new_table,
             drop_disconnected=drop_disconnected,
-            source=source.proto_self,
+            source=source,
         )
 
         new_proto_self = copy.deepcopy(self.proto_self)
-        new_proto_self.feature_view_sources.append(proto_feature_view_source)
+        new_proto_self.feature_view_sources.append(fvs.proto_self)
 
         if output:
             primary_key = source.table.schema.get_primary_key()
@@ -1336,7 +1336,7 @@ class FeatureView(BaseModel[FeatureViewID, models_pb2.FeatureView, orm.FeatureVi
         del new_proto_self.feature_view_sources[:]
 
         for fvs in self.feature_view_sources:
-            if fvs.source.id in self.output_sources:
+            if fvs.source.id in self.output_source_ids:
                 row_id_column = f"row-id-{uuid.uuid4()}"
                 # TODO(Patrick): make row ids stable
                 new_fvs = fvs.with_table(

corvic/model/_pipeline.py

@@ -124,22 +124,10 @@ class Pipeline(BaseModel[PipelineID, models_pb2.Pipeline, orm.Pipeline]):
         return [
             sa_orm.selectinload(orm.Pipeline.inputs)
             .selectinload(orm.PipelineInput.resource)
-            .selectinload(orm.Resource.source_associations),
-            sa_orm.selectinload(orm.Pipeline.inputs)
-            .selectinload(orm.PipelineInput.resource)
-            .selectinload(orm.Resource.pipeline_input_refs),
-            sa_orm.selectinload(orm.Pipeline.inputs).selectinload(
-                orm.PipelineInput.pipeline
-            ),
-            sa_orm.selectinload(orm.Pipeline.outputs).selectinload(
-                orm.PipelineOutput.pipeline
-            ),
-            sa_orm.selectinload(orm.Pipeline.outputs)
-            .selectinload(orm.PipelineOutput.source)
-            .selectinload(orm.Source.resource_associations),
+            .selectinload(orm.Resource.pipeline_ref),
             sa_orm.selectinload(orm.Pipeline.outputs)
             .selectinload(orm.PipelineOutput.source)
-            .selectinload(orm.Source.pipeline_output_refs),
+            .selectinload(orm.Source.pipeline_ref),
         ]
 
     @classmethod