mlrun 1.7.0rc41__py3-none-any.whl → 1.7.0rc43__py3-none-any.whl

mlrun/artifacts/manager.py

@@ -72,7 +72,12 @@ class ArtifactProducer:
         self.inputs = {}
 
     def get_meta(self) -> dict:
-        return {"kind": self.kind, "name": self.name, "tag": self.tag}
+        return {
+            "kind": self.kind,
+            "name": self.name,
+            "tag": self.tag,
+            "owner": self.owner,
+        }
 
     @property
     def uid(self):

mlrun/common/db/sql_session.py

@@ -11,13 +11,14 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
-#
 
 
 from sqlalchemy import create_engine
 from sqlalchemy.engine import Engine
 from sqlalchemy.orm import Session
-from sqlalchemy.orm import sessionmaker as SessionMaker
+from sqlalchemy.orm import (
+    sessionmaker as SessionMaker,  # noqa: N812 - `sessionmaker` is a class
+)
 
 from mlrun.config import config
 

mlrun/common/schemas/__init__.py

@@ -108,6 +108,7 @@ from .feature_store import (
     FeatureVectorsTagsOutput,
 )
 from .frontend_spec import (
+    ArtifactLimits,
     AuthenticationFeatureFlag,
     FeatureFlags,
     FrontendSpec,

mlrun/common/schemas/api_gateway.py

@@ -77,7 +77,7 @@ class APIGatewaySpec(_APIGatewayBaseModel):
     name: str
     description: Optional[str]
     path: Optional[str] = "/"
-    authenticationMode: Optional[APIGatewayAuthenticationMode] = (
+    authenticationMode: Optional[APIGatewayAuthenticationMode] = (  # noqa: N815 - for compatibility with Nuclio https://github.com/nuclio/nuclio/blob/672b8e36f9edd6e42b4685ec1d27cabae3c5f045/pkg/platform/types.go#L476
         APIGatewayAuthenticationMode.none
     )
     upstreams: list[APIGatewayUpstream]
@@ -103,11 +103,11 @@ class APIGateway(_APIGatewayBaseModel):
         ]
 
     def get_invoke_url(self):
-        return (
-            self.spec.host + self.spec.path
-            if self.spec.path and self.spec.host
-            else self.spec.host
-        ).rstrip("/")
+        if self.spec.host and self.spec.path:
+            return f"{self.spec.host.rstrip('/')}/{self.spec.path.lstrip('/')}".rstrip(
+                "/"
+            )
+        return self.spec.host.rstrip("/")
 
     def enrich_mlrun_names(self):
         self._enrich_api_gateway_mlrun_name()

mlrun/common/schemas/common.py

@@ -11,16 +11,16 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
-#
+
 import typing
 
 import pydantic
 
 
 class ImageBuilder(pydantic.BaseModel):
-    functionSourceCode: typing.Optional[str] = None
-    codeEntryType: typing.Optional[str] = None
-    codeEntryAttributes: typing.Optional[str] = None
+    functionSourceCode: typing.Optional[str] = None  # noqa: N815
+    codeEntryType: typing.Optional[str] = None  # noqa: N815
+    codeEntryAttributes: typing.Optional[str] = None  # noqa: N815
     source: typing.Optional[str] = None
     code_origin: typing.Optional[str] = None
     origin_filename: typing.Optional[str] = None

mlrun/common/schemas/frontend_spec.py

@@ -50,6 +50,12 @@ class FeatureFlags(pydantic.BaseModel):
     preemption_nodes: PreemptionNodesFeatureFlag
 
 
+class ArtifactLimits(pydantic.BaseModel):
+    max_chunk_size: int
+    max_preview_size: int
+    max_download_size: int
+
+
 class FrontendSpec(pydantic.BaseModel):
     jobs_dashboard_url: typing.Optional[str]
     model_monitoring_dashboard_url: typing.Optional[str]
@@ -71,3 +77,4 @@ class FrontendSpec(pydantic.BaseModel):
     allowed_artifact_path_prefixes_list: list[str]
     ce: typing.Optional[dict]
     internal_labels: list[str] = []
+    artifact_limits: ArtifactLimits

mlrun/common/schemas/notification.py

@@ -22,11 +22,38 @@ import mlrun.common.types
 
 
 class NotificationKind(mlrun.common.types.StrEnum):
-    console = "console"
-    git = "git"
-    ipython = "ipython"
-    slack = "slack"
-    webhook = "webhook"
+    """Currently, the supported notification kinds and their params are as follows:"""
+
+    console: str = "console"
+    """no params, local only"""
+
+    git: str = "git"
+    """
+    **token** - The git token to use for the git notification.\n
+    **repo** - The git repo to which to send the notification.\n
+    **issue** - The git issue to which to send the notification.\n
+    **merge_request** -
+     In GitLab (as opposed to GitHub), merge requests and issues are separate entities.
+     If using merge request, the issue will be ignored, and vice versa.\n
+    **server** - The git server to which to send the notification.\n
+    **gitlab** - (bool) Whether the git server is GitLab or not.\n
+    """
+
+    ipython: str = "ipython"
+    """no params, local only"""
+
+    slack: str = "slack"
+    """**webhook** - The slack webhook to which to send the notification."""
+
+    webhook: str = "webhook"
+    """
+    **url** - The webhook url to which to send the notification.\n
+    **method** - The http method to use when sending the notification (GET, POST, PUT, etc…).\n
+    **headers** - (dict) The http headers to send with the notification.\n
+    **override_body** - (dict) The body to send with the notification.\n
+    **verify_ssl** -
+     (bool) Whether SSL certificates are validated during HTTP requests or not, The default is set to True.
+    """
 
 
 class NotificationSeverity(mlrun.common.types.StrEnum):

mlrun/config.py

@@ -27,6 +27,7 @@ import copy
 import json
 import os
 import typing
+import warnings
 from collections.abc import Mapping
 from datetime import timedelta
 from distutils.util import strtobool
@@ -35,6 +36,7 @@ from threading import Lock
 
 import dotenv
 import semver
+import urllib3.exceptions
 import yaml
 
 import mlrun.common.constants
@@ -152,6 +154,11 @@ default_config = {
         "datasets": {
             "max_preview_columns": 100,
         },
+        "limits": {
+            "max_chunk_size": 1024 * 1024 * 1,  # 1MB
+            "max_preview_size": 1024 * 1024 * 10,  # 10MB
+            "max_download_size": 1024 * 1024 * 100,  # 100MB
+        },
     },
     # FIXME: Adding these defaults here so we won't need to patch the "installing component" (provazio-controller) to
     #  configure this values on field systems, for newer system this will be configured correctly
@@ -326,7 +333,7 @@ default_config = {
         "http": {
             # when True, the client will verify the server's TLS
             # set to False for backwards compatibility.
-            "verify": False,
+            "verify": True,
         },
         "db": {
             "commit_retry_timeout": 30,
@@ -1292,6 +1299,7 @@ def _do_populate(env=None, skip_errors=False):
     if data:
         config.update(data, skip_errors=skip_errors)
 
+    _configure_ssl_verification(config.httpdb.http.verify)
     _validate_config(config)
 
 
@@ -1351,6 +1359,16 @@ def _convert_str(value, typ):
     return typ(value)
 
 
+def _configure_ssl_verification(verify_ssl: bool) -> None:
+    """Configure SSL verification warnings based on the setting."""
+    if not verify_ssl:
+        urllib3.disable_warnings(urllib3.exceptions.InsecureRequestWarning)
+    else:
+        # If the user changes the `verify` setting to `True` at runtime using `mlrun.set_env_from_file` after
+        # importing `mlrun`, we need to reload the `mlrun` configuration and enable this warning.
+        warnings.simplefilter("default", urllib3.exceptions.InsecureRequestWarning)
+
+
 def read_env(env=None, prefix=env_prefix):
     """Read configuration from environment"""
     env = os.environ if env is None else env

mlrun/data_types/to_pandas.py

@@ -19,7 +19,7 @@ import pandas as pd
 import semver
 
 
-def _toPandas(spark_df):
+def _to_pandas(spark_df):
     """
     Modified version of spark DataFrame.toPandas() -
     https://github.com/apache/spark/blob/v3.2.3/python/pyspark/sql/pandas/conversion.py#L35
@@ -262,9 +262,9 @@ def spark_df_to_pandas(spark_df):
                 )
                 type_conversion_dict[field.name] = "datetime64[ns]"
 
-        df = _toPandas(spark_df)
+        df = _to_pandas(spark_df)
         if type_conversion_dict:
             df = df.astype(type_conversion_dict)
         return df
     else:
-        return _toPandas(spark_df)
+        return _to_pandas(spark_df)

mlrun/datastore/base.py

@@ -24,7 +24,6 @@ import pandas as pd
 import pyarrow
 import pytz
 import requests
-import urllib3
 from deprecated import deprecated
 
 import mlrun.config
@@ -745,8 +744,6 @@ class HttpStore(DataStore):
 
         verify_ssl = mlconf.httpdb.http.verify
         try:
-            if not verify_ssl:
-                urllib3.disable_warnings(urllib3.exceptions.InsecureRequestWarning)
             response = requests.get(url, headers=headers, auth=auth, verify=verify_ssl)
         except OSError as exc:
             raise OSError(f"error: cannot connect to {url}: {err_to_str(exc)}")

mlrun/datastore/storeytargets.py

@@ -137,7 +137,8 @@ class RedisNoSqlStoreyTarget(storey.NoSqlTarget):
     def __init__(self, *args, **kwargs):
         path = kwargs.pop("path")
         endpoint, uri = mlrun.datastore.targets.RedisNoSqlTarget.get_server_endpoint(
-            path
+            path,
+            kwargs.pop("credentials_prefix", None),
         )
         kwargs["path"] = endpoint + "/" + uri
         super().__init__(*args, **kwargs)

mlrun/datastore/targets.py

@@ -439,6 +439,12 @@ class BaseStoreTarget(DataTargetBase):
         self.storage_options = storage_options
         self.schema = schema or {}
         self.credentials_prefix = credentials_prefix
+        if credentials_prefix:
+            warnings.warn(
+                "The 'credentials_prefix' parameter is deprecated and will be removed in "
+                "1.9.0. Please use datastore profiles instead.",
+                FutureWarning,
+            )
 
         self._target = None
         self._resource = None
@@ -1479,7 +1485,7 @@ class RedisNoSqlTarget(NoSqlBaseTarget):
     writer_step_name = "RedisNoSqlTarget"
 
     @staticmethod
-    def get_server_endpoint(path):
+    def get_server_endpoint(path, credentials_prefix=None):
         endpoint, uri = parse_path(path)
         endpoint = endpoint or mlrun.mlconf.redis.url
         if endpoint.startswith("ds://"):
@@ -1497,7 +1503,9 @@ class RedisNoSqlTarget(NoSqlBaseTarget):
                 raise mlrun.errors.MLRunInvalidArgumentError(
                     "Provide Redis username and password only via secrets"
                 )
-            credentials_prefix = mlrun.get_secret_or_env(key="CREDENTIALS_PREFIX")
+            credentials_prefix = credentials_prefix or mlrun.get_secret_or_env(
+                key="CREDENTIALS_PREFIX"
+            )
             user = mlrun.get_secret_or_env(
                 "REDIS_USER", default="", prefix=credentials_prefix
             )
@@ -1517,7 +1525,9 @@ class RedisNoSqlTarget(NoSqlBaseTarget):
         from storey import Table
         from storey.redis_driver import RedisDriver
 
-        endpoint, uri = self.get_server_endpoint(self.get_target_path())
+        endpoint, uri = self.get_server_endpoint(
+            self.get_target_path(), self.credentials_prefix
+        )
 
         return Table(
             uri,
@@ -1526,7 +1536,9 @@ class RedisNoSqlTarget(NoSqlBaseTarget):
         )
 
     def get_spark_options(self, key_column=None, timestamp_key=None, overwrite=True):
-        endpoint, uri = self.get_server_endpoint(self.get_target_path())
+        endpoint, uri = self.get_server_endpoint(
+            self.get_target_path(), self.credentials_prefix
+        )
         parsed_endpoint = urlparse(endpoint)
         store, path_in_store, path = self._get_store_and_path()
         return {
@@ -1577,6 +1589,7 @@ class RedisNoSqlTarget(NoSqlBaseTarget):
             class_name="mlrun.datastore.storeytargets.RedisNoSqlStoreyTarget",
             columns=column_list,
             table=table,
+            credentials_prefix=self.credentials_prefix,
             **self.attributes,
         )
 

mlrun/errors.py

@@ -29,11 +29,14 @@ class MLRunBaseError(Exception):
     pass
 
 
-class MLRunTaskNotReady(MLRunBaseError):
+class MLRunTaskNotReadyError(MLRunBaseError):
     """indicate we are trying to read a value which is not ready
     or need to come from a job which is in progress"""
 
 
+MLRunTaskNotReady = MLRunTaskNotReadyError  # kept for BC only
+
+
 class MLRunHTTPError(MLRunBaseError, requests.HTTPError):
     def __init__(
         self,
@@ -205,15 +208,15 @@ class MLRunTimeoutError(MLRunHTTPStatusError, TimeoutError):
     error_status_code = HTTPStatus.GATEWAY_TIMEOUT.value
 
 
-class MLRunInvalidMMStoreType(MLRunHTTPStatusError, ValueError):
+class MLRunInvalidMMStoreTypeError(MLRunHTTPStatusError, ValueError):
     error_status_code = HTTPStatus.BAD_REQUEST.value
 
 
-class MLRunStreamConnectionFailure(MLRunHTTPStatusError, ValueError):
+class MLRunStreamConnectionFailureError(MLRunHTTPStatusError, ValueError):
     error_status_code = HTTPStatus.BAD_REQUEST.value
 
 
-class MLRunTSDBConnectionFailure(MLRunHTTPStatusError, ValueError):
+class MLRunTSDBConnectionFailureError(MLRunHTTPStatusError, ValueError):
     error_status_code = HTTPStatus.BAD_REQUEST.value
 
 

mlrun/execution.py

@@ -12,6 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import logging
 import os
 import uuid
 from copy import deepcopy
@@ -168,6 +169,8 @@ class MLClientCtx:
     @log_level.setter
     def log_level(self, value: str):
         """Set the logging level, e.g. 'debug', 'info', 'error'"""
+        level = logging.getLevelName(value.upper())
+        self._logger.set_logger_level(level)
         self._log_level = value
 
     @property
@@ -335,7 +338,7 @@ class MLClientCtx:
             "name": self.name,
             "kind": "run",
             "uri": uri,
-            "owner": get_in(self._labels, "owner"),
+            "owner": get_in(self._labels, mlrun_constants.MLRunInternalLabels.owner),
         }
         if mlrun_constants.MLRunInternalLabels.workflow in self._labels:
             resp[mlrun_constants.MLRunInternalLabels.workflow] = self._labels[

mlrun/feature_store/feature_vector.py

@@ -1086,7 +1086,9 @@ class OfflineVectorResponse:
     def to_dataframe(self, to_pandas=True):
         """return result as dataframe"""
         if self.status != "completed":
-            raise mlrun.errors.MLRunTaskNotReady("feature vector dataset is not ready")
+            raise mlrun.errors.MLRunTaskNotReadyError(
+                "feature vector dataset is not ready"
+            )
         return self._merger.get_df(to_pandas=to_pandas)
 
     def to_parquet(self, target_path, **kw):

mlrun/feature_store/retrieval/job.py

@@ -156,7 +156,9 @@ class RemoteVectorResponse:
 
     def _is_ready(self):
         if self.status != "completed":
-            raise mlrun.errors.MLRunTaskNotReady("feature vector dataset is not ready")
+            raise mlrun.errors.MLRunTaskNotReadyError(
+                "feature vector dataset is not ready"
+            )
         self.vector.reload()
 
     def to_dataframe(self, columns=None, df_module=None, **kwargs):

mlrun/frameworks/sklearn/mlrun_interface.py

@@ -97,7 +97,7 @@ class SKLearnMLRunInterface(MLRunInterface, ABC):
 
         def wrapper(
             self: SKLearnTypes.ModelType,
-            X: SKLearnTypes.DatasetType,
+            X: SKLearnTypes.DatasetType,  # noqa: N803 - should be lowercase "x", kept for BC
             y: SKLearnTypes.DatasetType = None,
             *args,
             **kwargs,
@@ -124,7 +124,12 @@ class SKLearnMLRunInterface(MLRunInterface, ABC):
 
         return wrapper
 
-    def mlrun_predict(self, X: SKLearnTypes.DatasetType, *args, **kwargs):
+    def mlrun_predict(
+        self,
+        X: SKLearnTypes.DatasetType,  # noqa: N803 - should be lowercase "x", kept for BC
+        *args,
+        **kwargs,
+    ):
         """
         MLRun's wrapper for the common ML API predict method.
         """
@@ -136,7 +141,12 @@ class SKLearnMLRunInterface(MLRunInterface, ABC):
 
         return y_pred
 
-    def mlrun_predict_proba(self, X: SKLearnTypes.DatasetType, *args, **kwargs):
+    def mlrun_predict_proba(
+        self,
+        X: SKLearnTypes.DatasetType,  # noqa: N803 - should be lowercase "x", kept for BC
+        *args,
+        **kwargs,
+    ):
         """
         MLRun's wrapper for the common ML API predict_proba method.
         """

mlrun/k8s_utils.py

@@ -12,6 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 import re
+import warnings
 
 import kubernetes.client
 
@@ -133,7 +134,7 @@ def sanitize_label_value(value: str) -> str:
     return re.sub(r"([^a-zA-Z0-9_.-]|^[^a-zA-Z0-9]|[^a-zA-Z0-9]$)", "-", value[:63])
 
 
-def verify_label_key(key: str):
+def verify_label_key(key: str, allow_k8s_prefix: bool = False):
     """
     Verify that the label key is valid for Kubernetes.
     Refer to https://kubernetes.io/docs/concepts/overview/working-with-objects/labels/#syntax-and-character-set
@@ -146,6 +147,10 @@ def verify_label_key(key: str):
         name = parts[0]
     elif len(parts) == 2:
         prefix, name = parts
+        if len(name) == 0:
+            raise mlrun.errors.MLRunInvalidArgumentError(
+                "Label key name cannot be empty when a prefix is set"
+            )
         if len(prefix) == 0:
             raise mlrun.errors.MLRunInvalidArgumentError(
                 "Label key prefix cannot be empty"
@@ -173,7 +178,13 @@ def verify_label_key(key: str):
         mlrun.utils.regex.qualified_name,
     )
 
-    if key.startswith("k8s.io/") or key.startswith("kubernetes.io/"):
+    # Allow the use of Kubernetes reserved prefixes ('k8s.io/' or 'kubernetes.io/')
+    # only when setting node selectors, not when adding new labels.
+    if (
+        key.startswith("k8s.io/")
+        or key.startswith("kubernetes.io/")
+        and not allow_k8s_prefix
+    ):
         raise mlrun.errors.MLRunInvalidArgumentError(
             "Labels cannot start with 'k8s.io/' or 'kubernetes.io/'"
         )
@@ -185,3 +196,38 @@ def verify_label_value(value, label_key):
         value,
         mlrun.utils.regex.label_value,
     )
+
+
+def validate_node_selectors(
+    node_selectors: dict[str, str], raise_on_error: bool = True
+) -> bool:
+    """
+    Ensures that user-defined node selectors adhere to Kubernetes label standards:
+    - Validates that each key conforms to Kubernetes naming conventions, with specific rules for name and prefix.
+    - Ensures values comply with Kubernetes label value rules.
+    - If raise_on_error is True, raises errors for invalid selectors.
+    - If raise_on_error is False, logs warnings for invalid selectors.
+    """
+
+    # Helper function for handling errors or warnings
+    def handle_invalid(message):
+        if raise_on_error:
+            raise
+        else:
+            warnings.warn(
+                f"{message}\n"
+                f"The node selector you’ve set does not meet the validation rules for the current Kubernetes version. "
+                f"Please note that invalid node selectors may cause issues with function scheduling."
+            )
+
+    node_selectors = node_selectors or {}
+    for key, value in node_selectors.items():
+        try:
+            verify_label_key(key, allow_k8s_prefix=True)
+            verify_label_value(value, label_key=key)
+        except mlrun.errors.MLRunInvalidArgumentError as err:
+            # An error or warning is raised by handle_invalid due to validation failure.
+            # Returning False indicates validation failed, allowing us to exit the function.
+            handle_invalid(str(err))
+            return False
+    return True

mlrun/model.py

@@ -487,7 +487,7 @@ class ImageBuilder(ModelObj):
 
     def __init__(
         self,
-        functionSourceCode=None,
+        functionSourceCode=None,  # noqa: N803 - should be "snake_case", kept for BC
         source=None,
         image=None,
         base_image=None,
@@ -681,7 +681,8 @@ class ImageBuilder(ModelObj):
 class Notification(ModelObj):
     """Notification object
 
-    :param kind: notification implementation kind - slack, webhook, etc.
+    :param kind: notification implementation kind - slack, webhook, etc. See
+        :py:class:`mlrun.common.schemas.notification.NotificationKind`
     :param name: for logging and identification
     :param message: message content in the notification
     :param severity: severity to display in the notification

mlrun/model_monitoring/db/stores/__init__.py

@@ -63,7 +63,7 @@ class ObjectStoreFactory(enum.Enum):
         :param value: Provided enum (invalid) value.
         """
         valid_values = list(cls.__members__.keys())
-        raise mlrun.errors.MLRunInvalidMMStoreType(
+        raise mlrun.errors.MLRunInvalidMMStoreTypeError(
             f"{value} is not a valid endpoint store, please choose a valid value: %{valid_values}."
         )
 
@@ -101,7 +101,7 @@ def get_store_object(
 
     :return: `StoreBase` object. Using this object, the user can apply different operations such as write, update, get
              and delete a model endpoint record.
-    :raise: `MLRunInvalidMMStoreType` if the user didn't provide store connection
+    :raise: `MLRunInvalidMMStoreTypeError` if the user didn't provide store connection
              or the provided store connection is invalid.
     """
 
@@ -123,7 +123,7 @@ def get_store_object(
             mlrun.common.schemas.model_monitoring.ModelEndpointTarget.V3IO_NOSQL
         )
     else:
-        raise mlrun.errors.MLRunInvalidMMStoreType(
+        raise mlrun.errors.MLRunInvalidMMStoreTypeError(
             "You must provide a valid store connection by using "
             "set_model_monitoring_credentials API."
         )

mlrun/model_monitoring/db/tsdb/__init__.py

@@ -57,7 +57,7 @@ class ObjectTSDBFactory(enum.Enum):
         :param value: Provided enum (invalid) value.
         """
         valid_values = list(cls.__members__.keys())
-        raise mlrun.errors.MLRunInvalidMMStoreType(
+        raise mlrun.errors.MLRunInvalidMMStoreTypeError(
             f"{value} is not a valid tsdb, please choose a valid value: %{valid_values}."
         )
 
@@ -76,7 +76,7 @@ def get_tsdb_connector(
 
     :return: `TSDBConnector` object. The main goal of this object is to handle different operations on the
              TSDB connector such as updating drift metrics or write application record result.
-    :raise: `MLRunInvalidMMStoreType` if the user didn't provide TSDB connection
+    :raise: `MLRunInvalidMMStoreTypeError` if the user didn't provide TSDB connection
             or the provided TSDB connection is invalid.
     """
 
@@ -93,7 +93,7 @@ def get_tsdb_connector(
     elif tsdb_connection_string and tsdb_connection_string == "v3io":
         tsdb_connector_type = mlrun.common.schemas.model_monitoring.TSDBTarget.V3IO_TSDB
     else:
-        raise mlrun.errors.MLRunInvalidMMStoreType(
+        raise mlrun.errors.MLRunInvalidMMStoreTypeError(
             "You must provide a valid tsdb store connection by using "
             "set_model_monitoring_credentials API."
         )

mlrun/model_monitoring/db/tsdb/tdengine/tdengine_connector.py

@@ -68,7 +68,7 @@ class TDEngineConnector(TSDBConnector):
         try:
             conn.execute(f"USE {self.database}")
         except taosws.QueryError as e:
-            raise mlrun.errors.MLRunTSDBConnectionFailure(
+            raise mlrun.errors.MLRunTSDBConnectionFailureError(
                 f"Failed to use TDEngine database {self.database}, {mlrun.errors.err_to_str(e)}"
             )
         return conn
@@ -91,7 +91,7 @@ class TDEngineConnector(TSDBConnector):
         """Create TDEngine supertables."""
         for table in self.tables:
             create_table_query = self.tables[table]._create_super_table_query()
-            self._connection.execute(create_table_query)
+            self.connection.execute(create_table_query)
 
     def write_application_event(
         self,
@@ -135,10 +135,10 @@ class TDEngineConnector(TSDBConnector):
         create_table_query = table._create_subtable_query(
             subtable=table_name, values=event
         )
-        self._connection.execute(create_table_query)
+        self.connection.execute(create_table_query)
 
         insert_statement = table._insert_subtable_query(
-            self._connection,
+            self.connection,
             subtable=table_name,
             values=event,
         )
@@ -204,12 +204,12 @@ class TDEngineConnector(TSDBConnector):
             get_subtable_names_query = self.tables[table]._get_subtables_query(
                 values={mm_schemas.EventFieldType.PROJECT: self.project}
             )
-            subtables = self._connection.query(get_subtable_names_query)
+            subtables = self.connection.query(get_subtable_names_query)
             for subtable in subtables:
                 drop_query = self.tables[table]._drop_subtable_query(
                     subtable=subtable[0]
                 )
-                self._connection.execute(drop_query)
+                self.connection.execute(drop_query)
         logger.info(
             f"Deleted all project resources in the TSDB connector for project {self.project}"
         )
@@ -281,7 +281,7 @@ class TDEngineConnector(TSDBConnector):
             database=self.database,
         )
         try:
-            query_result = self._connection.query(full_query)
+            query_result = self.connection.query(full_query)
         except taosws.QueryError as e:
             raise mlrun.errors.MLRunInvalidArgumentError(
                 f"Failed to query table {table} in database {self.database}, {str(e)}"