mlrun 1.3.3__py3-none-any.whl → 1.4.0__py3-none-any.whl

mlrun/api/db/sqldb/helpers.py

@@ -1,4 +1,4 @@
-# Copyright 2018 Iguazio
+# Copyright 2023 Iguazio
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.

mlrun/api/db/sqldb/models/__init__.py

@@ -1,4 +1,4 @@
-# Copyright 2018 Iguazio
+# Copyright 2023 Iguazio
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
@@ -20,10 +20,10 @@ if mysql_dsn_data:
     from .models_mysql import *  # noqa
 
     # importing private variables as well
-    from .models_mysql import _classes, _labeled, _table2cls, _tagged  # noqa # isort:skip
+    from .models_mysql import _classes, _labeled, _table2cls, _tagged, _with_notifications  # noqa # isort:skip
 else:
     from .models_sqlite import *  # noqa
 
     # importing private variables as well
-    from .models_sqlite import _classes, _labeled, _table2cls, _tagged  # noqa # isort:skip
+    from .models_sqlite import _classes, _labeled, _table2cls, _tagged, _with_notifications  # noqa # isort:skip
 # fmt: on

mlrun/api/db/sqldb/models/models_mysql.py

@@ -1,4 +1,4 @@
-# Copyright 2019 Iguazio
+# Copyright 2023 Iguazio
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
@@ -30,9 +30,10 @@ from sqlalchemy import (
     UniqueConstraint,
 )
 from sqlalchemy.ext.declarative import declarative_base
-from sqlalchemy.orm import class_mapper, relationship
+from sqlalchemy.orm import relationship
 
-from mlrun.api import schemas
+import mlrun.common.schemas
+import mlrun.utils.db
 from mlrun.api.utils.db.sql_collation import SQLCollationUtil
 
 Base = declarative_base()
@@ -40,42 +41,8 @@ NULL = None  # Avoid flake8 issuing warnings when comparing in filter
 run_time_fmt = "%Y-%m-%dT%H:%M:%S.%fZ"
 
 
-class BaseModel:
-    def to_dict(self, exclude=None):
-        """
-        NOTE - this function (currently) does not handle serializing relationships
-        """
-        exclude = exclude or []
-        mapper = class_mapper(self.__class__)
-        columns = [column.key for column in mapper.columns if column.key not in exclude]
-        get_key_value = (
-            lambda c: (c, getattr(self, c).isoformat())
-            if isinstance(getattr(self, c), datetime)
-            else (c, getattr(self, c))
-        )
-        return dict(map(get_key_value, columns))
-
-
-class HasStruct(BaseModel):
-    @property
-    def struct(self):
-        return pickle.loads(self.body)
-
-    @struct.setter
-    def struct(self, value):
-        self.body = pickle.dumps(value)
-
-    def to_dict(self, exclude=None):
-        """
-        NOTE - this function (currently) does not handle serializing relationships
-        """
-        exclude = exclude or []
-        exclude.append("body")
-        return super().to_dict(exclude)
-
-
 def make_label(table):
-    class Label(Base, BaseModel):
+    class Label(Base, mlrun.utils.db.BaseModel):
         __tablename__ = f"{table}_labels"
         __table_args__ = (
             UniqueConstraint("name", "parent", name=f"_{table}_labels_uc"),
@@ -90,7 +57,7 @@ def make_label(table):
 
 
 def make_tag(table):
-    class Tag(Base, BaseModel):
+    class Tag(Base, mlrun.utils.db.BaseModel):
         __tablename__ = f"{table}_tags"
         __table_args__ = (
             UniqueConstraint("project", "name", "obj_id", name=f"_{table}_tags_uc"),
@@ -107,7 +74,7 @@ def make_tag(table):
 # TODO: don't want to refactor everything in one PR so splitting this function to 2 versions - eventually only this one
 #  should be used
 def make_tag_v2(table):
-    class Tag(Base, BaseModel):
+    class Tag(Base, mlrun.utils.db.BaseModel):
         __tablename__ = f"{table}_tags"
         __table_args__ = (
             UniqueConstraint("project", "name", "obj_name", name=f"_{table}_tags_uc"),
@@ -122,11 +89,57 @@ def make_tag_v2(table):
     return Tag
 
 
+def make_notification(table):
+    class Notification(Base, mlrun.utils.db.BaseModel):
+        __tablename__ = f"{table}_notifications"
+        __table_args__ = (
+            UniqueConstraint("name", "parent_id", name=f"_{table}_notifications_uc"),
+        )
+
+        id = Column(Integer, primary_key=True)
+        project = Column(String(255, collation=SQLCollationUtil.collation()))
+        name = Column(
+            String(255, collation=SQLCollationUtil.collation()), nullable=False
+        )
+        kind = Column(
+            String(255, collation=SQLCollationUtil.collation()), nullable=False
+        )
+        message = Column(
+            String(255, collation=SQLCollationUtil.collation()), nullable=False
+        )
+        severity = Column(
+            String(255, collation=SQLCollationUtil.collation()), nullable=False
+        )
+        when = Column(
+            String(255, collation=SQLCollationUtil.collation()), nullable=False
+        )
+        condition = Column(
+            String(255, collation=SQLCollationUtil.collation()), nullable=False
+        )
+        params = Column("params", JSON)
+        parent_id = Column(Integer, ForeignKey(f"{table}.id"))
+
+        # TODO: Separate table for notification state.
+        #   Currently, we are only supporting one notification being sent per DB row (either on completion or on error).
+        #   In the future, we might want to support multiple notifications per DB row, and we might want to support on
+        #   start, therefore we need to separate the state from the notification itself (e.g. this table can be  table
+        #   with notification_id, state, when, last_sent, etc.). This will require some refactoring in the code.
+        sent_time = Column(
+            sqlalchemy.dialects.mysql.TIMESTAMP(fsp=3),
+            nullable=True,
+        )
+        status = Column(
+            String(255, collation=SQLCollationUtil.collation()), nullable=False
+        )
+
+    return Notification
+
+
 # quell SQLAlchemy warnings on duplicate class name (Label)
 with warnings.catch_warnings():
     warnings.simplefilter("ignore")
 
-    class Artifact(Base, HasStruct):
+    class Artifact(Base, mlrun.utils.db.HasStruct):
         __tablename__ = "artifacts"
         __table_args__ = (
             UniqueConstraint("uid", "project", "key", name="_artifacts_uc"),
@@ -140,7 +153,7 @@ with warnings.catch_warnings():
         project = Column(String(255, collation=SQLCollationUtil.collation()))
         uid = Column(String(255, collation=SQLCollationUtil.collation()))
         updated = Column(sqlalchemy.dialects.mysql.TIMESTAMP(fsp=3))
-        # TODO: change to JSON, see mlrun/api/schemas/function.py::FunctionState for reasoning
+        # TODO: change to JSON, see mlrun/common/schemas/function.py::FunctionState for reasoning
         body = Column(sqlalchemy.dialects.mysql.MEDIUMBLOB)
 
         labels = relationship(Label, cascade="all, delete-orphan")
@@ -149,7 +162,7 @@ with warnings.catch_warnings():
         def get_identifier_string(self) -> str:
             return f"{self.project}/{self.key}/{self.uid}"
 
-    class Function(Base, HasStruct):
+    class Function(Base, mlrun.utils.db.HasStruct):
         __tablename__ = "functions"
         __table_args__ = (
             UniqueConstraint("name", "project", "uid", name="_functions_uc"),
@@ -162,7 +175,7 @@ with warnings.catch_warnings():
         name = Column(String(255, collation=SQLCollationUtil.collation()))
         project = Column(String(255, collation=SQLCollationUtil.collation()))
         uid = Column(String(255, collation=SQLCollationUtil.collation()))
-        # TODO: change to JSON, see mlrun/api/schemas/function.py::FunctionState for reasoning
+        # TODO: change to JSON, see mlrun/common/schemas/function.py::FunctionState for reasoning
         body = Column(sqlalchemy.dialects.mysql.MEDIUMBLOB)
         updated = Column(sqlalchemy.dialects.mysql.TIMESTAMP(fsp=3))
 
@@ -172,19 +185,19 @@ with warnings.catch_warnings():
         def get_identifier_string(self) -> str:
             return f"{self.project}/{self.name}/{self.uid}"
 
-    class Log(Base, BaseModel):
+    class Log(Base, mlrun.utils.db.BaseModel):
         __tablename__ = "logs"
 
         id = Column(Integer, primary_key=True)
         uid = Column(String(255, collation=SQLCollationUtil.collation()))
         project = Column(String(255, collation=SQLCollationUtil.collation()))
-        # TODO: change to JSON, see mlrun/api/schemas/function.py::FunctionState for reasoning
+        # TODO: change to JSON, see mlrun/common/schemas/function.py::FunctionState for reasoning
         body = Column(sqlalchemy.dialects.mysql.MEDIUMBLOB)
 
         def get_identifier_string(self) -> str:
             return f"{self.project}/{self.uid}"
 
-    class Run(Base, HasStruct):
+    class Run(Base, mlrun.utils.db.HasStruct):
         __tablename__ = "runs"
         __table_args__ = (
             UniqueConstraint("uid", "project", "iteration", name="_runs_uc"),
@@ -192,6 +205,7 @@ with warnings.catch_warnings():
 
         Label = make_label(__tablename__)
         Tag = make_tag(__tablename__)
+        Notification = make_notification(__tablename__)
 
         id = Column(Integer, primary_key=True)
         uid = Column(String(255, collation=SQLCollationUtil.collation()))
@@ -201,7 +215,7 @@ with warnings.catch_warnings():
         )
         iteration = Column(Integer)
         state = Column(String(255, collation=SQLCollationUtil.collation()))
-        # TODO: change to JSON, see mlrun/api/schemas/function.py::FunctionState for reasoning
+        # TODO: change to JSON, see mlrun/common/schemas/function.py::FunctionState for reasoning
         body = Column(sqlalchemy.dialects.mysql.MEDIUMBLOB)
         start_time = Column(sqlalchemy.dialects.mysql.TIMESTAMP(fsp=3))
         updated = Column(
@@ -215,11 +229,12 @@ with warnings.catch_warnings():
 
         labels = relationship(Label, cascade="all, delete-orphan")
         tags = relationship(Tag, cascade="all, delete-orphan")
+        notifications = relationship(Notification, cascade="all, delete-orphan")
 
         def get_identifier_string(self) -> str:
             return f"{self.project}/{self.uid}/{self.iteration}"
 
-    class BackgroundTask(Base, BaseModel):
+    class BackgroundTask(Base, mlrun.utils.db.BaseModel):
         __tablename__ = "background_tasks"
         __table_args__ = (
             UniqueConstraint("name", "project", name="_background_tasks_uc"),
@@ -243,7 +258,7 @@ with warnings.catch_warnings():
         state = Column(String(255, collation=SQLCollationUtil.collation()))
         timeout = Column(Integer)
 
-    class Schedule(Base, BaseModel):
+    class Schedule(Base, mlrun.utils.db.BaseModel):
         __tablename__ = "schedules_v2"
         __table_args__ = (UniqueConstraint("project", "name", name="_schedules_v2_uc"),)
 
@@ -262,7 +277,7 @@ with warnings.catch_warnings():
         creation_time = Column(sqlalchemy.dialects.mysql.TIMESTAMP(fsp=3))
         cron_trigger_str = Column(String(255, collation=SQLCollationUtil.collation()))
         last_run_uri = Column(String(255, collation=SQLCollationUtil.collation()))
-        # TODO: change to JSON, see mlrun/api/schemas/function.py::FunctionState for reasoning
+        # TODO: change to JSON, see mlrun/common/schemas/function.py::FunctionState for reasoning
         struct = Column(sqlalchemy.dialects.mysql.MEDIUMBLOB)
         labels = relationship(Label, cascade="all, delete-orphan")
         concurrency_limit = Column(Integer, nullable=False)
@@ -280,11 +295,11 @@ with warnings.catch_warnings():
             self.struct = pickle.dumps(value)
 
         @property
-        def cron_trigger(self) -> schemas.ScheduleCronTrigger:
+        def cron_trigger(self) -> mlrun.common.schemas.ScheduleCronTrigger:
             return orjson.loads(self.cron_trigger_str)
 
         @cron_trigger.setter
-        def cron_trigger(self, trigger: schemas.ScheduleCronTrigger):
+        def cron_trigger(self, trigger: mlrun.common.schemas.ScheduleCronTrigger):
             self.cron_trigger_str = orjson.dumps(trigger.dict(exclude_unset=True))
 
     # Define "many to many" users/projects
@@ -295,14 +310,14 @@ with warnings.catch_warnings():
         Column("user_id", Integer, ForeignKey("users.id")),
     )
 
-    class User(Base, BaseModel):
+    class User(Base, mlrun.utils.db.BaseModel):
         __tablename__ = "users"
         __table_args__ = (UniqueConstraint("name", name="_users_uc"),)
 
         id = Column(Integer, primary_key=True)
         name = Column(String(255, collation=SQLCollationUtil.collation()))
 
-    class Project(Base, BaseModel):
+    class Project(Base, mlrun.utils.db.BaseModel):
         __tablename__ = "projects"
         # For now since we use project name a lot
         __table_args__ = (UniqueConstraint("name", name="_projects_uc"),)
@@ -314,7 +329,7 @@ with warnings.catch_warnings():
         source = Column(String(255, collation=SQLCollationUtil.collation()))
         # the attribute name used to be _spec which is just a wrong naming, the attribute was renamed to _full_object
         # leaving the column as is to prevent redundant migration
-        # TODO: change to JSON, see mlrun/api/schemas/function.py::FunctionState for reasoning
+        # TODO: change to JSON, see mlrun/common/schemas/function.py::FunctionState for reasoning
         _full_object = Column("spec", sqlalchemy.dialects.mysql.MEDIUMBLOB)
         created = Column(
             sqlalchemy.dialects.mysql.TIMESTAMP(fsp=3), default=datetime.utcnow
@@ -338,7 +353,7 @@ with warnings.catch_warnings():
         def full_object(self, value):
             self._full_object = pickle.dumps(value)
 
-    class Feature(Base, BaseModel):
+    class Feature(Base, mlrun.utils.db.BaseModel):
         __tablename__ = "features"
         id = Column(Integer, primary_key=True)
         feature_set_id = Column(Integer, ForeignKey("feature_sets.id"))
@@ -352,7 +367,7 @@ with warnings.catch_warnings():
         def get_identifier_string(self) -> str:
             return f"{self.project}/{self.name}"
 
-    class Entity(Base, BaseModel):
+    class Entity(Base, mlrun.utils.db.BaseModel):
         __tablename__ = "entities"
         id = Column(Integer, primary_key=True)
         feature_set_id = Column(Integer, ForeignKey("feature_sets.id"))
@@ -366,7 +381,7 @@ with warnings.catch_warnings():
         def get_identifier_string(self) -> str:
             return f"{self.project}/{self.name}"
 
-    class FeatureSet(Base, BaseModel):
+    class FeatureSet(Base, mlrun.utils.db.BaseModel):
         __tablename__ = "feature_sets"
         __table_args__ = (
             UniqueConstraint("name", "project", "uid", name="_feature_set_uc"),
@@ -410,7 +425,7 @@ with warnings.catch_warnings():
             # TODO - convert to pickle, to avoid issues with non-json serializable fields such as datetime
             self._full_object = json.dumps(value, default=str)
 
-    class FeatureVector(Base, BaseModel):
+    class FeatureVector(Base, mlrun.utils.db.BaseModel):
         __tablename__ = "feature_vectors"
         __table_args__ = (
             UniqueConstraint("name", "project", "uid", name="_feature_vectors_uc"),
@@ -451,9 +466,9 @@ with warnings.catch_warnings():
             # TODO - convert to pickle, to avoid issues with non-json serializable fields such as datetime
             self._full_object = json.dumps(value, default=str)
 
-    class MarketplaceSource(Base, BaseModel):
-        __tablename__ = "marketplace_sources"
-        __table_args__ = (UniqueConstraint("name", name="_marketplace_sources_uc"),)
+    class HubSource(Base, mlrun.utils.db.BaseModel):
+        __tablename__ = "hub_sources"
+        __table_args__ = (UniqueConstraint("name", name="_hub_sources_uc"),)
 
         id = Column(Integer, primary_key=True)
         name = Column(String(255, collation=SQLCollationUtil.collation()))
@@ -482,7 +497,7 @@ with warnings.catch_warnings():
             # TODO - convert to pickle, to avoid issues with non-json serializable fields such as datetime
             self._full_object = json.dumps(value, default=str)
 
-    class DataVersion(Base, BaseModel):
+    class DataVersion(Base, mlrun.utils.db.BaseModel):
         __tablename__ = "data_versions"
         __table_args__ = (UniqueConstraint("version", name="_versions_uc"),)
 
@@ -497,5 +512,8 @@ with warnings.catch_warnings():
 # Must be after all table definitions
 _tagged = [cls for cls in Base.__subclasses__() if hasattr(cls, "Tag")]
 _labeled = [cls for cls in Base.__subclasses__() if hasattr(cls, "Label")]
+_with_notifications = [
+    cls for cls in Base.__subclasses__() if hasattr(cls, "Notification")
+]
 _classes = [cls for cls in Base.__subclasses__()]
 _table2cls = {cls.__table__.name: cls for cls in Base.__subclasses__()}