corvic-engine 0.3.0rc76__cp38-abi3-win_amd64.whl → 0.3.0rc78__cp38-abi3-win_amd64.whl

corvic/orm/__init__.py

@@ -1,38 +1,36 @@
-"""Data model definitions; backed by an RDBMS."""
+"""Object-Relational Mappings."""
 
 from __future__ import annotations
 
-from datetime import datetime
+from collections.abc import Callable
+from datetime import UTC, datetime
+from typing import Any, ClassVar, cast
 
 import sqlalchemy as sa
+from sqlalchemy import event as sa_event
 from sqlalchemy import orm as sa_orm
+from sqlalchemy.ext import hybrid
 
+import corvic.context
 from corvic.orm._proto_columns import ProtoMessageDecorator
-from corvic.orm.base import Base, OrgBase
+from corvic.orm._soft_delete import (
+    BadDeleteError,
+    Session,
+    SoftDeleteMixin,
+    live_unique_constraint,
+)
 from corvic.orm.errors import (
     DeletedObjectError,
     InvalidORMIdentifierError,
     RequestedObjectsForNobodyError,
 )
+from corvic.orm.func import utc_now
 from corvic.orm.ids import (
-    AgentID,
-    AgentMessageID,
     BaseID,
     BaseIDFromInt,
-    CompletionModelID,
-    FeatureViewID,
-    FeatureViewSourceID,
+    BaseIDFromStr,
     IntIDDecorator,
-    MessageEntryID,
-    OrgID,
-    PipelineID,
-    ResourceID,
-    RoomID,
-    SourceID,
-    SpaceID,
-    SpaceParametersID,
-    SpaceRunID,
-    UserMessageID,
+    StrIDDecorator,
 )
 from corvic.orm.keys import (
     INT_PK_TYPE,
@@ -41,346 +39,243 @@ from corvic.orm.keys import (
     primary_key_identity_column,
     primary_key_uuid_column,
 )
-from corvic.orm.mixins import (
-    BelongsToOrgMixin,
-    Session,
-    SoftDeleteMixin,
-    live_unique_constraint,
-)
-from corvic_generated.orm.v1 import (
-    common_pb2,
-    completion_model_pb2,
-    feature_view_pb2,
-    pipeline_pb2,
-    space_pb2,
-    table_pb2,
-)
-from corvic_generated.status.v1 import event_pb2
-
-# NOTE: The only safe use of "sa_orm.relationship" uses the args:
-# `viewonly=True` and `init=False`. Writes quickly become
-# a complex mess when implementers of commit need to reason about
-# which sub-object should be updated.
-#
-# Rather, classes in corvic.model define their own commit protocols,
-# and if sub-orm-model updates are required they are explicit.
-
+from corvic.result import InvalidArgumentError
 
-class Org(SoftDeleteMixin, OrgBase, kw_only=True):
-    """An organization it a top level grouping of resources."""
 
-
-class Room(BelongsToOrgMixin, SoftDeleteMixin, Base, kw_only=True):
-    """A Room is a logical collection of Documents."""
-
-    __tablename__ = "room"
-    __table_args__ = (live_unique_constraint("name", "org_id"),)
-
-    name: sa_orm.Mapped[str] = sa_orm.mapped_column(sa.Text, default=None)
-    id: sa_orm.Mapped[RoomID | None] = primary_key_identity_column()
+class OrgID(BaseIDFromStr):
+    """A unique identifier for an organization."""
 
     @property
-    def room_key(self):
-        return self.name
-
-
-class BelongsToRoomMixin(sa_orm.MappedAsDataclass):
-    room_id: sa_orm.Mapped[RoomID | None] = sa_orm.mapped_column(
-        ForeignKey(Room).make(ondelete="CASCADE"),
-        nullable=True,
-        default=None,
-    )
-
-
-class DefaultObjects(Base, kw_only=True):
-    """Holds the identifiers for default objects."""
-
-    __tablename__ = "default_objects"
-    default_org: sa_orm.Mapped[OrgID | None] = sa_orm.mapped_column(
-        ForeignKey(Org).make(ondelete="CASCADE"),
-        nullable=False,
-    )
-    default_room: sa_orm.Mapped[RoomID | None] = sa_orm.mapped_column(
-        ForeignKey(Room).make(ondelete="CASCADE"), nullable=True, default=None
-    )
-    version: sa_orm.Mapped[int | None] = primary_key_identity_column(type_=INT_PK_TYPE)
-
-
-class Resource(BelongsToOrgMixin, BelongsToRoomMixin, Base, kw_only=True):
-    """A Resource is a reference to some durably stored file.
-
-    E.g., a document could be a PDF file, an image, or a text transcript of a
-    conversation
-    """
-
-    __tablename__ = "resource"
-
-    name: sa_orm.Mapped[str] = sa_orm.mapped_column(sa.Text)
-    mime_type: sa_orm.Mapped[str] = sa_orm.mapped_column(sa.Text)
-    url: sa_orm.Mapped[str] = sa_orm.mapped_column(sa.Text)
-    md5: sa_orm.Mapped[str] = sa_orm.mapped_column(sa.CHAR(32), nullable=True)
-    size: sa_orm.Mapped[int] = sa_orm.mapped_column(nullable=True)
-    original_path: sa_orm.Mapped[str] = sa_orm.mapped_column(nullable=True)
-    description: sa_orm.Mapped[str] = sa_orm.mapped_column(nullable=True)
-    id: sa_orm.Mapped[ResourceID | None] = primary_key_identity_column()
-    latest_event: sa_orm.Mapped[event_pb2.Event | None] = sa_orm.mapped_column(
-        default=None, nullable=True
-    )
-    is_terminal: sa_orm.Mapped[bool | None] = sa_orm.mapped_column(
-        default=None, nullable=True
-    )
-    pipeline_ref: sa_orm.Mapped[PipelineInput | None] = sa_orm.relationship(
-        init=False, viewonly=True
-    )
-
-
-class Source(BelongsToOrgMixin, BelongsToRoomMixin, Base, kw_only=True):
-    """A source."""
-
-    __tablename__ = "source"
-    __table_args__ = (sa.UniqueConstraint("name", "room_id"),)
-
-    name: sa_orm.Mapped[str] = sa_orm.mapped_column(sa.Text)
-    # protobuf describing the operations required to construct a table
-    table_op_graph: sa_orm.Mapped[table_pb2.TableComputeOp] = sa_orm.mapped_column()
-    id: sa_orm.Mapped[SourceID | None] = primary_key_identity_column()
-
-    source_files: sa_orm.Mapped[common_pb2.BlobUrlList | None] = sa_orm.mapped_column(
-        default=None
-    )
-    pipeline_ref: sa_orm.Mapped[PipelineOutput | None] = sa_orm.relationship(
-        init=False, viewonly=True
-    )
+    def is_super_user(self):
+        return self._value == corvic.context.SUPERUSER_ORG_ID
 
     @property
-    def source_key(self):
-        return self.name
+    def is_nobody(self):
+        return self._value == corvic.context.NOBODY_ORG_ID
 
 
-class Pipeline(BelongsToOrgMixin, BelongsToRoomMixin, Base, kw_only=True):
-    """A resource to source pipeline."""
+# A quick primer on SQLAlchemy (sa) hybrid methods:
+#
+# Hybrid just means functionality that is different at the class-level versus
+# the instance-level, and in the sa documentation, the authors really
+# want to stress that class-versus-instance (decorators) is orthgonal to an
+# ORM.
+#
+# However, this distinction is not particularly helpful for users of sa.
+# It is best to have a working model of instance-level means Python-world and
+# class-level means SQL-world. So, a hybrid method is something that has
+# a different Python representation from its SQL representation.
+#
+# Since sa already handles conversions like "Python str" to "SQL text",
+# certain representation differences between Python and SQL are already handled
+# (not considered differences at all).
+#
+# Hybrid methods are for cases where we want to do non-trivial transformations
+# between SQL and Python representations.
+#
+# The recipe is:
+#
+# 1. Define a hybrid_method / hybrid_property (wlog property) that produces the Python
+#    object you want.
+# 2. If the property doesn't need to be used in any sa query again, you are done.
+# 3. If the property is simple enough for sa to also use it to produce the SQL
+#    representation, you are also done. E.g., comparisons and bitwise operations
+#    on columns.
+# 4. Otherwise, you need to define a class-level function, hybrid_property.expression,
+#    which gives the SQL representation of your property when it is passed to
+#    a sa query.
+# 5. Because of how redefining decorators is supposed to work in Python [1], you
+#    should use @<property_method_name>.inplace.expression to define your
+#    class-level function that describes how the property should be represented
+#    in SQL.
+#
+# [1] https://docs.sqlalchemy.org/en/20/orm/extensions/hybrid.html#using-inplace-to-create-pep-484-compliant-hybrid-properties
 
-    __tablename__ = "pipeline"
-    __table_args__ = (sa.UniqueConstraint("name", "room_id"),)
 
-    transformation: sa_orm.Mapped[pipeline_pb2.PipelineTransformation] = (
-        sa_orm.mapped_column()
-    )
-    name: sa_orm.Mapped[str] = sa_orm.mapped_column()
-    description: sa_orm.Mapped[str | None] = sa_orm.mapped_column()
-    id: sa_orm.Mapped[PipelineID | None] = primary_key_identity_column()
+class Base(sa_orm.MappedAsDataclass, sa_orm.DeclarativeBase):
+    """Base class for all DB mapped classes."""
+
+    type_annotation_map: ClassVar = {
+        OrgID: StrIDDecorator(OrgID()),
+    }
 
-    outputs: sa_orm.Mapped[list[PipelineOutput]] = sa_orm.relationship(
-        viewonly=True,
+    _created_at: sa_orm.Mapped[datetime] = sa_orm.mapped_column(
+        "created_at",
+        sa.DateTime(timezone=True),
+        server_default=utc_now(),
         init=False,
-        default_factory=list,
+        index=True,
     )
 
-
-class PipelineInput(BelongsToOrgMixin, BelongsToRoomMixin, Base, kw_only=True):
-    """Pipeline input resources."""
-
-    __tablename__ = "pipeline_input"
-    __table_args__ = (sa.UniqueConstraint("name", "pipeline_id"),)
-
-    resource: sa_orm.Mapped[Resource] = sa_orm.relationship(viewonly=True, init=False)
-    name: sa_orm.Mapped[str]
-    """A name the pipeline uses to refer to this input."""
-
-    pipeline_id: sa_orm.Mapped[PipelineID] = primary_key_foreign_column(
-        ForeignKey(Pipeline).make(ondelete="CASCADE")
-    )
-    resource_id: sa_orm.Mapped[ResourceID] = primary_key_foreign_column(
-        ForeignKey(Resource).make(ondelete="CASCADE")
+    _updated_at: sa_orm.Mapped[datetime] = sa_orm.mapped_column(
+        "updated_at",
+        sa.DateTime(timezone=True),
+        onupdate=utc_now(),
+        server_default=utc_now(),
+        init=False,
+        nullable=True,
     )
 
+    @hybrid.hybrid_property
+    def created_at(self) -> datetime | None:
+        if not self._created_at:
+            return None
+        return self._created_at.replace(tzinfo=UTC)
 
-class PipelineOutput(BelongsToOrgMixin, BelongsToRoomMixin, Base, kw_only=True):
-    """Objects for tracking pipeline output sources."""
-
-    __tablename__ = "pipeline_output"
-    __table_args__ = (sa.UniqueConstraint("name", "pipeline_id"),)
+    @created_at.inplace.expression
+    @classmethod
+    def _created_at_expression(cls):
+        return cls._created_at
 
-    source: sa_orm.Mapped[Source] = sa_orm.relationship(viewonly=True, init=False)
-    name: sa_orm.Mapped[str]
-    """A name the pipeline uses to refer to this output."""
-
-    pipeline_id: sa_orm.Mapped[PipelineID] = primary_key_foreign_column(
-        ForeignKey(Pipeline).make(ondelete="CASCADE")
-    )
-    source_id: sa_orm.Mapped[SourceID] = primary_key_foreign_column(
-        ForeignKey(Source).make(ondelete="CASCADE")
-    )
+    @hybrid.hybrid_property
+    def updated_at(self) -> datetime | None:
+        if not self._updated_at:
+            return None
+        return self._updated_at.replace(tzinfo=UTC)
 
+    @updated_at.inplace.expression
+    @classmethod
+    def _updated_at_expression(cls):
+        return cls._updated_at
 
-class FeatureView(
-    SoftDeleteMixin, BelongsToOrgMixin, BelongsToRoomMixin, Base, kw_only=True
-):
-    """A FeatureView is a logical collection of sources used by various spaces."""
 
-    __tablename__ = "feature_view"
-    __table_args__ = (live_unique_constraint("name", "room_id"),)
+class Org(Base, kw_only=True):
+    """An organization is a top-level grouping of resources."""
 
-    id: sa_orm.Mapped[FeatureViewID | None] = primary_key_identity_column()
-    name: sa_orm.Mapped[str] = sa_orm.mapped_column(sa.Text, default=None)
-    description: sa_orm.Mapped[str] = sa_orm.mapped_column(sa.Text, default="")
+    __tablename__ = "org"
 
-    feature_view_output: sa_orm.Mapped[feature_view_pb2.FeatureViewOutput | None] = (
-        sa_orm.mapped_column(default_factory=feature_view_pb2.FeatureViewOutput)
-    )
+    id: sa_orm.Mapped[OrgID | None] = primary_key_uuid_column()
 
     @property
-    def feature_view_key(self):
-        return self.name
-
-    feature_view_sources: sa_orm.Mapped[list[FeatureViewSource]] = sa_orm.relationship(
-        viewonly=True,
-        init=False,
-        default_factory=list,
-    )
-
-
-class FeatureViewSource(BelongsToOrgMixin, BelongsToRoomMixin, Base, kw_only=True):
-    """A source inside of a feature view."""
-
-    __tablename__ = "feature_view_source"
-
-    table_op_graph: sa_orm.Mapped[table_pb2.TableComputeOp] = sa_orm.mapped_column()
-    feature_view_id: sa_orm.Mapped[FeatureViewID] = sa_orm.mapped_column(
-        ForeignKey(FeatureView).make(ondelete="CASCADE"),
-        nullable=False,
-    )
-    id: sa_orm.Mapped[FeatureViewSourceID | None] = primary_key_identity_column()
-    drop_disconnected: sa_orm.Mapped[bool] = sa_orm.mapped_column(default=False)
-    # this should be legal but pyright complains that it makes Source depend
-    # on itself
-    source_id: sa_orm.Mapped[SourceID] = sa_orm.mapped_column(
-        ForeignKey(Source).make(ondelete="CASCADE"),
+    def name(self) -> str:
+        return str(self.id)
+
+
+def _filter_org_objects(orm_execute_state: sa_orm.ORMExecuteState):
+    if all(
+        not issubclass(mapper.class_, BelongsToOrgMixin | Org)
+        for mapper in orm_execute_state.all_mappers
+    ):
+        # operation has nothing to do with models owned by org
+        return
+    if orm_execute_state.is_select:
+        requester = corvic.context.get_requester()
+        org_id = OrgID(requester.org_id)
+        if org_id.is_super_user:
+            return
+
+        if org_id.is_nobody:
+            raise RequestedObjectsForNobodyError(
+                "requester org from context was nobody"
+            )
+        # we need the real value in in expression world and
+        # because of sqlalchemys weird runtime parsing of this it
+        # needs to be a real local with a name
+        db_id = org_id.to_db().unwrap_or_raise()
+
+        # this goofy syntax doesn't typecheck well, but is the documented way to apply
+        # these operations to all subclasses (recursive). Sqlalchemy is inspecting the
+        # lambda rather than just executing it so a function won't work.
+        # https://docs.sqlalchemy.org/en/20/orm/queryguide/api.html#sqlalchemy.orm.with_loader_criteria
+        check_org_id_lambda: Callable[  # noqa: E731
+            [type[BelongsToOrgMixin]], sa.ColumnElement[bool]
+        ] = lambda cls: cls.org_id == db_id
+        orm_execute_state.statement = orm_execute_state.statement.options(
+            sa_orm.with_loader_criteria(
+                BelongsToOrgMixin,
+                cast(Any, check_org_id_lambda),
+                include_aliases=True,
+                track_closure_variables=False,
+            ),
+            sa_orm.with_loader_criteria(
+                Org,
+                Org.id == org_id,
+                include_aliases=True,
+                track_closure_variables=False,
+            ),
+        )
+
+
+class BelongsToOrgMixin(sa_orm.MappedAsDataclass):
+    """Mark models that should be subject to org level access control."""
+
+    @staticmethod
+    def _current_org_id_from_context():
+        requester = corvic.context.get_requester()
+        return OrgID(requester.org_id)
+
+    @staticmethod
+    def _make_org_id_default() -> OrgID | None:
+        org_id = BelongsToOrgMixin._current_org_id_from_context()
+
+        if org_id.is_nobody:
+            raise RequestedObjectsForNobodyError(
+                "the nobody org cannot change orm objects"
+            )
+
+        if org_id.is_super_user:
+            return None
+
+        return org_id
+
+    org_id: sa_orm.Mapped[OrgID | None] = sa_orm.mapped_column(
+        ForeignKey(Org).make(ondelete="CASCADE"),
         nullable=False,
-        default=None,
-    )
-    source: sa_orm.Mapped[Source] = sa_orm.relationship(
-        init=True, viewonly=True, default=None
+        default_factory=_make_org_id_default,
+        init=False,
     )
 
+    @sa_orm.validates("org_id")
+    def validate_org_id(self, _key: str, orm_id: OrgID | None):
+        expected_org_id = self._current_org_id_from_context()
+        if expected_org_id.is_nobody:
+            raise RequestedObjectsForNobodyError(
+                "the nobody org cannot change orm objects"
+            )
 
-class Space(BelongsToOrgMixin, BelongsToRoomMixin, Base, kw_only=True):
-    """A space is a named evaluation of space parameters."""
+        if expected_org_id.is_super_user:
+            return orm_id
 
-    __tablename__ = "space"
-    __table_args__ = (sa.UniqueConstraint("name", "room_id"),)
+        if orm_id != expected_org_id:
+            raise InvalidArgumentError(
+                "provided org_id must match the current org",
+                provided=orm_id,
+                expected=expected_org_id,
+            )
 
-    id: sa_orm.Mapped[SpaceID | None] = primary_key_identity_column()
-    name: sa_orm.Mapped[str] = sa_orm.mapped_column(sa.Text, default=None)
-    description: sa_orm.Mapped[str] = sa_orm.mapped_column(sa.Text, default="")
+        return orm_id
 
-    feature_view_id: sa_orm.Mapped[FeatureViewID] = sa_orm.mapped_column(
-        ForeignKey(FeatureView).make(ondelete="CASCADE"),
-        nullable=False,
-        default=None,
-    )
-    parameters: sa_orm.Mapped[space_pb2.SpaceParameters | None] = sa_orm.mapped_column(
-        default=None
-    )
-    auto_sync: sa_orm.Mapped[bool | None] = sa_orm.mapped_column(default=None)
-    feature_view: sa_orm.Mapped[FeatureView] = sa_orm.relationship(
-        init=False,
-        default=None,
-        viewonly=True,
-    )
+    @staticmethod
+    def register_session_event_listeners(session: type[sa_orm.Session]):
+        sa_event.listen(session, "do_orm_execute", _filter_org_objects)
 
-    @property
-    def space_key(self):
-        return self.name
 
+BelongsToOrgMixin.register_session_event_listeners(Session)
 
-class CompletionModel(SoftDeleteMixin, BelongsToOrgMixin, Base, kw_only=True):
-    """A customer's custom completion model definition."""
 
-    __tablename__ = "completion_model"
-    __table_args__ = (live_unique_constraint("name", "org_id"),)
-
-    id: sa_orm.Mapped[CompletionModelID | None] = primary_key_identity_column()
-    name: sa_orm.Mapped[str] = sa_orm.mapped_column(sa.Text, default=None)
-    description: sa_orm.Mapped[str] = sa_orm.mapped_column(sa.Text, default=None)
-    parameters: sa_orm.Mapped[completion_model_pb2.CompletionModelParameters | None] = (
-        sa_orm.mapped_column(default=None)
-    )
-    secret_api_key: sa_orm.Mapped[str] = sa_orm.mapped_column(sa.Text, default=None)
-    last_validation_time: sa_orm.Mapped[datetime | None] = sa_orm.mapped_column(
-        sa.DateTime(timezone=True),
-        server_default=None,
-        default=None,
-    )
-    last_successful_validation: sa_orm.Mapped[datetime | None] = sa_orm.mapped_column(
-        sa.DateTime(timezone=True),
-        server_default=None,
-        default=None,
-    )
-
-    @property
-    def model_key(self):
-        return self.name
-
-
-ID = (
-    AgentID
-    | AgentMessageID
-    | CompletionModelID
-    | FeatureViewID
-    | FeatureViewSourceID
-    | MessageEntryID
-    | OrgID
-    | PipelineID
-    | ResourceID
-    | RoomID
-    | SourceID
-    | SpaceID
-    | SpaceParametersID
-    | SpaceRunID
-    | UserMessageID
-)
+ID = OrgID
 
 
 __all__ = [
-    "AgentID",
-    "AgentMessageID",
+    "BadDeleteError",
     "Base",
     "BaseID",
     "BaseIDFromInt",
     "BelongsToOrgMixin",
-    "CompletionModel",
-    "CompletionModelID",
-    "DefaultObjects",
     "DeletedObjectError",
-    "FeatureView",
-    "FeatureViewID",
-    "FeatureViewSource",
-    "FeatureViewSourceID",
+    "ForeignKey",
     "ID",
+    "INT_PK_TYPE",
+    "IntIDDecorator",
     "InvalidORMIdentifierError",
-    "MessageEntryID",
     "Org",
     "OrgID",
-    "PipelineID",
-    "PipelineInput",
-    "PipelineOutput",
+    "ProtoMessageDecorator",
     "RequestedObjectsForNobodyError",
-    "Resource",
-    "ResourceID",
-    "Room",
-    "RoomID",
     "Session",
-    "Source",
-    "SourceID",
-    "Space",
-    "SpaceID",
-    "SpaceParametersID",
-    "SpaceRunID",
-    "UserMessageID",
+    "SoftDeleteMixin",
+    "live_unique_constraint",
     "primary_key_foreign_column",
     "primary_key_identity_column",
     "primary_key_uuid_column",
-    "ProtoMessageDecorator",
-    "IntIDDecorator",
 ]