mlrun 1.10.0rc40__py3-none-any.whl → 1.11.0rc16__py3-none-any.whl

mlrun/datastore/storeytargets.py

@@ -18,11 +18,12 @@ from mergedeep import merge
 from storey import V3ioDriver
 
 import mlrun
+import mlrun.common.model_monitoring.helpers
 from mlrun.datastore.base import DataStore
 from mlrun.datastore.datastore_profile import (
     DatastoreProfileKafkaStream,
     DatastoreProfileKafkaTarget,
-    DatastoreProfileTDEngine,
+    DatastoreProfilePostgreSQL,
     datastore_profile_read,
 )
 
@@ -48,17 +49,22 @@ def get_url_and_storage_options(path, external_storage_options=None):
     return url, DataStore._sanitize_options(storage_options)
 
 
-class TDEngineStoreyTarget(storey.TDEngineTarget):
+class TimescaleDBStoreyTarget(storey.TimescaleDBTarget):
     def __init__(self, *args, url: str, **kwargs):
         if url.startswith("ds://"):
             datastore_profile = datastore_profile_read(url)
-            if not isinstance(datastore_profile, DatastoreProfileTDEngine):
-                raise ValueError(
-                    f"Unexpected datastore profile type:{datastore_profile.type}."
-                    "Only DatastoreProfileTDEngine is supported"
+            if not isinstance(datastore_profile, DatastoreProfilePostgreSQL):
+                raise mlrun.errors.MLRunInvalidArgumentError(
+                    f"Unexpected datastore profile type: {type(datastore_profile)}. "
+                    "Only DatastoreProfilePostgreSQL is supported"
                 )
-            url = datastore_profile.dsn()
-        super().__init__(*args, url=url, **kwargs)
+            # Use the shared helper to determine the correct database name
+            # This ensures consistency with TimescaleDBConnector's database naming
+            database = mlrun.common.model_monitoring.helpers.get_tsdb_database_name(
+                datastore_profile.database
+            )
+            url = datastore_profile.dsn(database=database)
+        super().__init__(*args, dsn=url, **kwargs)
 
 
 class StoreyTargetUtils:
@@ -137,7 +143,7 @@ class KafkaStoreyTarget(storey.KafkaTarget):
             datastore_profile = datastore_profile_read(path)
             if not isinstance(
                 datastore_profile,
-                (DatastoreProfileKafkaStream, DatastoreProfileKafkaTarget),
+                DatastoreProfileKafkaStream | DatastoreProfileKafkaTarget,
             ):
                 raise mlrun.errors.MLRunInvalidArgumentError(
                     f"Unsupported datastore profile type: {type(datastore_profile)}"
@@ -174,7 +180,7 @@ class RedisNoSqlStoreyTarget(storey.NoSqlTarget):
         endpoint, uri = mlrun.datastore.targets.RedisNoSqlTarget.get_server_endpoint(
             path
         )
-        kwargs["path"] = endpoint + "/" + uri
+        kwargs["path"] = f"{endpoint}/{uri}"
         super().__init__(*args, **kwargs)
 
 

mlrun/datastore/targets.py

@@ -532,7 +532,7 @@ class BaseStoreTarget(DataTargetBase):
             if (
                 file_system.protocol == "file"
                 # fsspec 2023.10.0 changed protocol from "file" to ("file", "local")
-                or isinstance(file_system.protocol, (tuple, list))
+                or isinstance(file_system.protocol, tuple | list)
                 and "file" in file_system.protocol
             ):
                 dir = os.path.dirname(target_path)

mlrun/datastore/utils.py

@@ -190,12 +190,12 @@ def validate_additional_filters(additional_filters):
     for filter_tuple in additional_filters:
         if filter_tuple == () or filter_tuple == []:
             continue
-        if not isinstance(filter_tuple, (list, tuple)):
+        if not isinstance(filter_tuple, list | tuple):
             raise mlrun.errors.MLRunInvalidArgumentError(
                 f"mlrun supports additional_filters only as a list of tuples."
                 f" Current additional_filters: {additional_filters}"
             )
-        if isinstance(filter_tuple[0], (list, tuple)):
+        if isinstance(filter_tuple[0], list | tuple):
             raise mlrun.errors.MLRunInvalidArgumentError(
                 f"additional_filters does not support nested list inside filter tuples except in -in- logic."
                 f" Current filter_tuple: {filter_tuple}."
@@ -208,7 +208,7 @@ def validate_additional_filters(additional_filters):
         col_name, op, value = filter_tuple
         if isinstance(value, float) and math.isnan(value):
             raise mlrun.errors.MLRunInvalidArgumentError(nan_error_message)
-        elif isinstance(value, (list, tuple)):
+        elif isinstance(value, list | tuple):
             for sub_value in value:
                 if isinstance(sub_value, float) and math.isnan(sub_value):
                     raise mlrun.errors.MLRunInvalidArgumentError(nan_error_message)
@@ -345,3 +345,16 @@ def parse_url(url):
 def accepts_param(func: callable, param_name):
     sig = inspect.signature(func)
     return param_name in sig.parameters
+
+
+def parse_s3_bucket_and_key(s3_path: str) -> tuple[str, str]:
+    try:
+        path_parts = s3_path.replace("s3://", "").split("/")
+        bucket = path_parts.pop(0)
+        key = "/".join(path_parts)
+    except Exception as exc:
+        raise mlrun.errors.MLRunInvalidArgumentError(
+            "failed to parse s3 bucket and key"
+        ) from exc
+
+    return bucket, key

mlrun/datastore/v3io.py

@@ -57,7 +57,7 @@ class V3ioStore(DataStore):
         self.auth = None
         self.token = token
         if token:
-            self.headers = {"X-v3io-session-key": token}
+            self.headers = {mlrun.common.schemas.HeaderNames.v3io_session_key: token}
         elif username and password:
             self.headers = basic_auth_header(username, password)
 

mlrun/db/base.py

@@ -188,7 +188,6 @@ class RunDBInterface(ABC):
         tree: Optional[str] = None,
         parent: Optional[str] = None,
         format_: mlrun.common.formatters.ArtifactFormat = mlrun.common.formatters.ArtifactFormat.full,
-        limit: Optional[int] = None,
         partition_by: Optional[
             Union[mlrun.common.schemas.ArtifactPartitionByField, str]
         ] = None,
@@ -675,17 +674,6 @@ class RunDBInterface(ABC):
     ):
         pass
 
-    @abstractmethod
-    def create_user_secrets(
-        self,
-        user: str,
-        provider: Union[
-            str, mlrun.common.schemas.SecretProviderName
-        ] = mlrun.common.schemas.SecretProviderName.vault,
-        secrets: Optional[dict] = None,
-    ):
-        pass
-
     @abstractmethod
     def create_model_endpoint(
         self,
@@ -1173,3 +1161,39 @@ class RunDBInterface(ABC):
         end: Optional[datetime.datetime] = None,
     ) -> mlrun.common.schemas.model_monitoring.ModelEndpointDriftValues:
         pass
+
+    @abstractmethod
+    def store_secret_token(
+        self,
+        secret_token: mlrun.common.schemas.SecretToken,
+        log_warning: bool = True,
+        force: bool = False,
+    ) -> mlrun.common.schemas.StoreSecretTokensResponse:
+        pass
+
+    @abstractmethod
+    def store_secret_tokens(
+        self,
+        secret_tokens: list[mlrun.common.schemas.SecretToken],
+        log_warning: bool = True,
+        force: bool = False,
+    ) -> mlrun.common.schemas.StoreSecretTokensResponse:
+        pass
+
+    @abstractmethod
+    def list_secret_tokens(
+        self,
+    ) -> mlrun.common.schemas.ListSecretTokensResponse:
+        pass
+
+    @abstractmethod
+    def revoke_secret_token(self, token_name: str) -> None:
+        pass
+
+    @abstractmethod
+    def get_secret_token(
+        self,
+        token_name: str,
+        username: Optional[str] = None,
+    ) -> mlrun.common.schemas.SecretToken:
+        pass