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

mlrun/api/db/__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.

mlrun/api/db/base.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.
@@ -12,11 +12,13 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 import datetime
+import typing
 import warnings
 from abc import ABC, abstractmethod
 from typing import Any, Dict, List, Optional, Tuple, Union
 
-from mlrun.api import schemas
+import mlrun.common.schemas
+import mlrun.model
 
 
 class DBError(Exception):
@@ -96,13 +98,14 @@ class DBInterface(ABC):
         start_time_to=None,
         last_update_time_from=None,
         last_update_time_to=None,
-        partition_by: schemas.RunPartitionByField = None,
+        partition_by: mlrun.common.schemas.RunPartitionByField = None,
         rows_per_partition: int = 1,
-        partition_sort_by: schemas.SortField = None,
-        partition_order: schemas.OrderType = schemas.OrderType.desc,
+        partition_sort_by: mlrun.common.schemas.SortField = None,
+        partition_order: mlrun.common.schemas.OrderType = mlrun.common.schemas.OrderType.desc,
         max_partitions: int = 0,
         requested_logs: bool = None,
         return_as_run_structs: bool = True,
+        with_notifications: bool = False,
     ):
         pass
 
@@ -119,7 +122,7 @@ class DBInterface(ABC):
         session,
         project: str,
         tag: str,
-        identifiers: List[schemas.ArtifactIdentifier],
+        identifiers: List[mlrun.common.schemas.ArtifactIdentifier],
     ):
         pass
 
@@ -128,7 +131,7 @@ class DBInterface(ABC):
         session,
         project: str,
         tag: str,
-        identifiers: List[schemas.ArtifactIdentifier],
+        identifiers: List[mlrun.common.schemas.ArtifactIdentifier],
     ):
         pass
 
@@ -137,7 +140,7 @@ class DBInterface(ABC):
         session,
         project: str,
         tag: str,
-        identifiers: List[schemas.ArtifactIdentifier],
+        identifiers: List[mlrun.common.schemas.ArtifactIdentifier],
     ):
         pass
 
@@ -169,7 +172,7 @@ class DBInterface(ABC):
         since=None,
         until=None,
         kind=None,
-        category: schemas.ArtifactCategories = None,
+        category: mlrun.common.schemas.ArtifactCategories = None,
         iter: int = None,
         best_iteration: bool = False,
         as_records: bool = False,
@@ -232,9 +235,9 @@ class DBInterface(ABC):
         session,
         project: str,
         name: str,
-        kind: schemas.ScheduleKinds,
+        kind: mlrun.common.schemas.ScheduleKinds,
         scheduled_object: Any,
-        cron_trigger: schemas.ScheduleCronTrigger,
+        cron_trigger: mlrun.common.schemas.ScheduleCronTrigger,
         concurrency_limit: int,
         labels: Dict = None,
         next_run_time: datetime.datetime = None,
@@ -248,7 +251,7 @@ class DBInterface(ABC):
         project: str,
         name: str,
         scheduled_object: Any = None,
-        cron_trigger: schemas.ScheduleCronTrigger = None,
+        cron_trigger: mlrun.common.schemas.ScheduleCronTrigger = None,
         labels: Dict = None,
         last_run_uri: str = None,
         concurrency_limit: int = None,
@@ -263,12 +266,14 @@ class DBInterface(ABC):
         project: str = None,
         name: str = None,
         labels: str = None,
-        kind: schemas.ScheduleKinds = None,
-    ) -> List[schemas.ScheduleRecord]:
+        kind: mlrun.common.schemas.ScheduleKinds = None,
+    ) -> List[mlrun.common.schemas.ScheduleRecord]:
         pass
 
     @abstractmethod
-    def get_schedule(self, session, project: str, name: str) -> schemas.ScheduleRecord:
+    def get_schedule(
+        self, session, project: str, name: str
+    ) -> mlrun.common.schemas.ScheduleRecord:
         pass
 
     @abstractmethod
@@ -282,7 +287,7 @@ class DBInterface(ABC):
     @abstractmethod
     def generate_projects_summaries(
         self, session, projects: List[str]
-    ) -> List[schemas.ProjectSummary]:
+    ) -> List[mlrun.common.schemas.ProjectSummary]:
         pass
 
     @abstractmethod
@@ -302,17 +307,17 @@ class DBInterface(ABC):
         self,
         session,
         owner: str = None,
-        format_: schemas.ProjectsFormat = schemas.ProjectsFormat.full,
+        format_: mlrun.common.schemas.ProjectsFormat = mlrun.common.schemas.ProjectsFormat.full,
         labels: List[str] = None,
-        state: schemas.ProjectState = None,
+        state: mlrun.common.schemas.ProjectState = None,
         names: Optional[List[str]] = None,
-    ) -> schemas.ProjectsOutput:
+    ) -> mlrun.common.schemas.ProjectsOutput:
         pass
 
     @abstractmethod
     def get_project(
         self, session, name: str = None, project_id: int = None
-    ) -> schemas.Project:
+    ) -> mlrun.common.schemas.Project:
         pass
 
     @abstractmethod
@@ -329,11 +334,11 @@ class DBInterface(ABC):
         pass
 
     @abstractmethod
-    def create_project(self, session, project: schemas.Project):
+    def create_project(self, session, project: mlrun.common.schemas.Project):
         pass
 
     @abstractmethod
-    def store_project(self, session, name: str, project: schemas.Project):
+    def store_project(self, session, name: str, project: mlrun.common.schemas.Project):
         pass
 
     @abstractmethod
@@ -342,7 +347,7 @@ class DBInterface(ABC):
         session,
         name: str,
         project: dict,
-        patch_mode: schemas.PatchMode = schemas.PatchMode.replace,
+        patch_mode: mlrun.common.schemas.PatchMode = mlrun.common.schemas.PatchMode.replace,
     ):
         pass
 
@@ -351,7 +356,7 @@ class DBInterface(ABC):
         self,
         session,
         name: str,
-        deletion_strategy: schemas.DeletionStrategy = schemas.DeletionStrategy.default(),
+        deletion_strategy: mlrun.common.schemas.DeletionStrategy = mlrun.common.schemas.DeletionStrategy.default(),
     ):
         pass
 
@@ -360,7 +365,7 @@ class DBInterface(ABC):
         self,
         session,
         project,
-        feature_set: schemas.FeatureSet,
+        feature_set: mlrun.common.schemas.FeatureSet,
         versioned=True,
     ) -> str:
         pass
@@ -371,7 +376,7 @@ class DBInterface(ABC):
         session,
         project,
         name,
-        feature_set: schemas.FeatureSet,
+        feature_set: mlrun.common.schemas.FeatureSet,
         tag=None,
         uid=None,
         versioned=True,
@@ -382,7 +387,7 @@ class DBInterface(ABC):
     @abstractmethod
     def get_feature_set(
         self, session, project: str, name: str, tag: str = None, uid: str = None
-    ) -> schemas.FeatureSet:
+    ) -> mlrun.common.schemas.FeatureSet:
         pass
 
     @abstractmethod
@@ -394,7 +399,7 @@ class DBInterface(ABC):
         tag: str = None,
         entities: List[str] = None,
         labels: List[str] = None,
-    ) -> schemas.FeaturesOutput:
+    ) -> mlrun.common.schemas.FeaturesOutput:
         pass
 
     @abstractmethod
@@ -405,7 +410,7 @@ class DBInterface(ABC):
         name: str = None,
         tag: str = None,
         labels: List[str] = None,
-    ) -> schemas.EntitiesOutput:
+    ) -> mlrun.common.schemas.EntitiesOutput:
         pass
 
     @abstractmethod
@@ -419,11 +424,11 @@ class DBInterface(ABC):
         entities: List[str] = None,
         features: List[str] = None,
         labels: List[str] = None,
-        partition_by: schemas.FeatureStorePartitionByField = None,
+        partition_by: mlrun.common.schemas.FeatureStorePartitionByField = None,
         rows_per_partition: int = 1,
-        partition_sort_by: schemas.SortField = None,
-        partition_order: schemas.OrderType = schemas.OrderType.desc,
-    ) -> schemas.FeatureSetsOutput:
+        partition_sort_by: mlrun.common.schemas.SortField = None,
+        partition_order: mlrun.common.schemas.OrderType = mlrun.common.schemas.OrderType.desc,
+    ) -> mlrun.common.schemas.FeatureSetsOutput:
         pass
 
     @abstractmethod
@@ -446,7 +451,7 @@ class DBInterface(ABC):
         feature_set_patch: dict,
         tag=None,
         uid=None,
-        patch_mode: schemas.PatchMode = schemas.PatchMode.replace,
+        patch_mode: mlrun.common.schemas.PatchMode = mlrun.common.schemas.PatchMode.replace,
     ) -> str:
         pass
 
@@ -459,7 +464,7 @@ class DBInterface(ABC):
         self,
         session,
         project,
-        feature_vector: schemas.FeatureVector,
+        feature_vector: mlrun.common.schemas.FeatureVector,
         versioned=True,
     ) -> str:
         pass
@@ -467,7 +472,7 @@ class DBInterface(ABC):
     @abstractmethod
     def get_feature_vector(
         self, session, project: str, name: str, tag: str = None, uid: str = None
-    ) -> schemas.FeatureVector:
+    ) -> mlrun.common.schemas.FeatureVector:
         pass
 
     @abstractmethod
@@ -479,11 +484,11 @@ class DBInterface(ABC):
         tag: str = None,
         state: str = None,
         labels: List[str] = None,
-        partition_by: schemas.FeatureStorePartitionByField = None,
+        partition_by: mlrun.common.schemas.FeatureStorePartitionByField = None,
         rows_per_partition: int = 1,
-        partition_sort_by: schemas.SortField = None,
-        partition_order: schemas.OrderType = schemas.OrderType.desc,
-    ) -> schemas.FeatureVectorsOutput:
+        partition_sort_by: mlrun.common.schemas.SortField = None,
+        partition_order: mlrun.common.schemas.OrderType = mlrun.common.schemas.OrderType.desc,
+    ) -> mlrun.common.schemas.FeatureVectorsOutput:
         pass
 
     @abstractmethod
@@ -503,7 +508,7 @@ class DBInterface(ABC):
         session,
         project,
         name,
-        feature_vector: schemas.FeatureVector,
+        feature_vector: mlrun.common.schemas.FeatureVector,
         tag=None,
         uid=None,
         versioned=True,
@@ -520,7 +525,7 @@ class DBInterface(ABC):
         feature_vector_update: dict,
         tag=None,
         uid=None,
-        patch_mode: schemas.PatchMode = schemas.PatchMode.replace,
+        patch_mode: mlrun.common.schemas.PatchMode = mlrun.common.schemas.PatchMode.replace,
     ) -> str:
         pass
 
@@ -536,29 +541,33 @@ class DBInterface(ABC):
         pass
 
     def list_artifact_tags(
-        self, session, project, category: Union[str, schemas.ArtifactCategories] = None
+        self,
+        session,
+        project,
+        category: Union[str, mlrun.common.schemas.ArtifactCategories] = None,
     ):
         return []
 
-    def create_marketplace_source(
-        self, session, ordered_source: schemas.IndexedMarketplaceSource
+    def create_hub_source(
+        self, session, ordered_source: mlrun.common.schemas.IndexedHubSource
     ):
         pass
 
-    def store_marketplace_source(
-        self, session, name, ordered_source: schemas.IndexedMarketplaceSource
+    def store_hub_source(
+        self,
+        session,
+        name,
+        ordered_source: mlrun.common.schemas.IndexedHubSource,
     ):
         pass
 
-    def list_marketplace_sources(
-        self, session
-    ) -> List[schemas.IndexedMarketplaceSource]:
+    def list_hub_sources(self, session) -> List[mlrun.common.schemas.IndexedHubSource]:
         pass
 
-    def delete_marketplace_source(self, session, name):
+    def delete_hub_source(self, session, name):
         pass
 
-    def get_marketplace_source(self, session, name) -> schemas.IndexedMarketplaceSource:
+    def get_hub_source(self, session, name) -> mlrun.common.schemas.IndexedHubSource:
         pass
 
     def store_background_task(
@@ -566,12 +575,51 @@ class DBInterface(ABC):
         session,
         name: str,
         project: str,
-        state: str = schemas.BackgroundTaskState.running,
+        state: str = mlrun.common.schemas.BackgroundTaskState.running,
         timeout: int = None,
     ):
         pass
 
     def get_background_task(
         self, session, name: str, project: str
-    ) -> schemas.BackgroundTask:
+    ) -> mlrun.common.schemas.BackgroundTask:
+        pass
+
+    @abstractmethod
+    def store_run_notifications(
+        self,
+        session,
+        notification_objects: typing.List[mlrun.model.Notification],
+        run_uid: str,
+        project: str,
+    ):
+        pass
+
+    @abstractmethod
+    def list_run_notifications(
+        self,
+        session,
+        run_uid: str,
+        project: str,
+    ) -> typing.List[mlrun.model.Notification]:
+        pass
+
+    def delete_run_notifications(
+        self,
+        session,
+        name: str = None,
+        run_uid: str = None,
+        project: str = None,
+        commit: bool = True,
+    ):
+        pass
+
+    def set_run_notifications(
+        self,
+        session,
+        project: str,
+        notifications: typing.List[mlrun.model.Notification],
+        identifiers: typing.List[mlrun.common.schemas.RunIdentifier],
+        **kwargs,
+    ):
         pass

mlrun/api/db/init_db.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.
@@ -12,13 +12,12 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 #
-from sqlalchemy.orm import Session
 
 from mlrun.api.db.sqldb.models import Base
 from mlrun.api.db.sqldb.session import get_engine
 from mlrun.config import config
 
 
-def init_db(db_session: Session) -> None:
+def init_db() -> None:
     if config.httpdb.db_type != "filedb":
         Base.metadata.create_all(bind=get_engine())

mlrun/api/db/session.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.
@@ -15,22 +15,14 @@
 from sqlalchemy.orm import Session
 
 from mlrun.api.db.sqldb.session import create_session as sqldb_create_session
-from mlrun.config import config
 
 
-def create_session(db_type=None) -> Session:
-    db_type = db_type or config.httpdb.db_type
-    if db_type == "filedb":
-        return None
-    else:
-        return sqldb_create_session()
+def create_session() -> Session:
+    return sqldb_create_session()
 
 
 def close_session(db_session):
-
-    # will be None when it's filedb session
-    if db_session is not None:
-        db_session.close()
+    db_session.close()
 
 
 def run_function_with_new_db_session(func):

mlrun/api/db/sqldb/__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.