mlrun 1.7.0rc18__py3-none-any.whl → 1.7.0rc19__py3-none-any.whl

mlrun/model_monitoring/db/stores/sqldb/sql_store.py

@@ -19,6 +19,9 @@ import uuid
 
 import pandas as pd
 import sqlalchemy
+import sqlalchemy.exc
+import sqlalchemy.orm
+from sqlalchemy.sql.elements import BinaryExpression
 
 import mlrun.common.model_monitoring.helpers
 import mlrun.common.schemas.model_monitoring as mm_schemas
@@ -26,6 +29,7 @@ import mlrun.model_monitoring.db
 import mlrun.model_monitoring.db.stores.sqldb.models
 import mlrun.model_monitoring.helpers
 from mlrun.common.db.sql_session import create_session, get_engine
+from mlrun.utils import datetime_now, logger
 
 
 class SQLStoreBase(mlrun.model_monitoring.db.StoreBase):
@@ -35,7 +39,6 @@ class SQLStoreBase(mlrun.model_monitoring.db.StoreBase):
     data, the user needs to provide a valid connection string for the database.
     """
 
-    _engine = None
     _tables = {}
 
     def __init__(
@@ -63,6 +66,7 @@ class SQLStoreBase(mlrun.model_monitoring.db.StoreBase):
     def _init_tables(self):
         self._init_model_endpoints_table()
         self._init_application_results_table()
+        self._init_application_metrics_table()
         self._init_monitoring_schedules_table()
 
     def _init_model_endpoints_table(self):
@@ -76,13 +80,21 @@ class SQLStoreBase(mlrun.model_monitoring.db.StoreBase):
         )
 
     def _init_application_results_table(self):
-        self.ApplicationResultsTable = (
+        self.application_results_table = (
             mlrun.model_monitoring.db.stores.sqldb.models._get_application_result_table(
                 connection_string=self._sql_connection_string
             )
         )
         self._tables[mm_schemas.FileTargetKind.APP_RESULTS] = (
-            self.ApplicationResultsTable
+            self.application_results_table
+        )
+
+    def _init_application_metrics_table(self) -> None:
+        self.application_metrics_table = mlrun.model_monitoring.db.stores.sqldb.models._get_application_metrics_table(
+            connection_string=self._sql_connection_string
+        )
+        self._tables[mm_schemas.FileTargetKind.APP_METRICS] = (
+            self.application_metrics_table
         )
 
     def _init_monitoring_schedules_table(self):
@@ -93,61 +105,61 @@ class SQLStoreBase(mlrun.model_monitoring.db.StoreBase):
             self.MonitoringSchedulesTable
         )
 
-    def _write(self, table: str, event: dict[str, typing.Any]):
+    def _write(self, table_name: str, event: dict[str, typing.Any]) -> None:
         """
         Create a new record in the SQL table.
 
-        :param table: Target table name.
-        :param event: Event dictionary that will be written into the DB.
+        :param table_name: Target table name.
+        :param event:      Event dictionary that will be written into the DB.
         """
-
         with self._engine.connect() as connection:
             # Convert the result into a pandas Dataframe and write it into the database
             event_df = pd.DataFrame([event])
-
-            event_df.to_sql(table, con=connection, index=False, if_exists="append")
+            event_df.to_sql(table_name, con=connection, index=False, if_exists="append")
 
     def _update(
         self,
         attributes: dict[str, typing.Any],
         table: sqlalchemy.orm.decl_api.DeclarativeMeta,
-        **filtered_values,
-    ):
+        criteria: list[BinaryExpression],
+    ) -> None:
         """
         Update a record in the SQL table.
 
         :param attributes:  Dictionary of attributes that will be used for update the record. Note that the keys
                             of the attributes dictionary should exist in the SQL table.
         :param table:       SQLAlchemy declarative table.
-
+        :param criteria:    A list of binary expressions that filter the query.
         """
-        filter_query_ = []
-        for _filter in filtered_values:
-            filter_query_.append(f"{_filter} = '{filtered_values[_filter]}'")
-
         with create_session(dsn=self._sql_connection_string) as session:
             # Generate and commit the update session query
-            session.query(table).filter(sqlalchemy.sql.text(*filter_query_)).update(
-                attributes, synchronize_session=False
-            )
+            session.query(
+                table  # pyright: ignore[reportOptionalCall]
+            ).filter(*criteria).update(attributes, synchronize_session=False)
             session.commit()
 
-    def _get(self, table: sqlalchemy.orm.decl_api.DeclarativeMeta, **filtered_values):
+    def _get(
+        self,
+        table: sqlalchemy.orm.decl_api.DeclarativeMeta,
+        criteria: list[BinaryExpression],
+    ):
         """
         Get a record from the SQL table.
 
-        param table: SQLAlchemy declarative table.
+        param table:     SQLAlchemy declarative table.
+        :param criteria: A list of binary expressions that filter the query.
         """
-
-        filter_query_ = []
-        for _filter in filtered_values:
-            filter_query_.append(f"{_filter} = '{filtered_values[_filter]}'")
         with create_session(dsn=self._sql_connection_string) as session:
             try:
+                logger.debug(
+                    "Querying the DB",
+                    table=table.__name__,
+                    criteria=[str(criterion) for criterion in criteria],
+                )
                 # Generate the get query
                 return (
-                    session.query(table)
-                    .filter(sqlalchemy.sql.text(*filter_query_))
+                    session.query(table)  # pyright: ignore[reportOptionalCall]
+                    .filter(*criteria)
                     .one_or_none()
                 )
             except sqlalchemy.exc.ProgrammingError:
@@ -156,21 +168,21 @@ class SQLStoreBase(mlrun.model_monitoring.db.StoreBase):
                 return
 
     def _delete(
-        self, table: sqlalchemy.orm.decl_api.DeclarativeMeta, **filtered_values
-    ):
+        self,
+        table: sqlalchemy.orm.decl_api.DeclarativeMeta,
+        criteria: list[BinaryExpression],
+    ) -> None:
         """
         Delete records from the SQL table.
 
-        param table: SQLAlchemy declarative table.
+        param table:     SQLAlchemy declarative table.
+        :param criteria: A list of binary expressions that filter the query.
         """
-        filter_query_ = []
-        for _filter in filtered_values:
-            filter_query_.append(f"{_filter} = '{filtered_values[_filter]}'")
         with create_session(dsn=self._sql_connection_string) as session:
             # Generate and commit the delete query
-            session.query(table).filter(sqlalchemy.sql.text(*filter_query_)).delete(
-                synchronize_session=False
-            )
+            session.query(
+                table  # pyright: ignore[reportOptionalCall]
+            ).filter(*criteria).delete(synchronize_session=False)
             session.commit()
 
     def write_model_endpoint(self, endpoint: dict[str, typing.Any]):
@@ -184,11 +196,10 @@ class SQLStoreBase(mlrun.model_monitoring.db.StoreBase):
         # Adjust timestamps fields
         endpoint[mm_schemas.EventFieldType.FIRST_REQUEST] = (endpoint)[
             mm_schemas.EventFieldType.LAST_REQUEST
-        ] = mlrun.utils.datetime_now()
+        ] = datetime_now()
 
         self._write(
-            table=mm_schemas.EventFieldType.MODEL_ENDPOINTS,
-            event=endpoint,
+            table_name=mm_schemas.EventFieldType.MODEL_ENDPOINTS, event=endpoint
         )
 
     def update_model_endpoint(
@@ -206,23 +217,24 @@ class SQLStoreBase(mlrun.model_monitoring.db.StoreBase):
 
         attributes.pop(mm_schemas.EventFieldType.ENDPOINT_ID, None)
 
-        filter_endpoint = {mm_schemas.EventFieldType.UID: endpoint_id}
-
         self._update(
-            attributes=attributes, table=self.ModelEndpointsTable, **filter_endpoint
+            attributes=attributes,
+            table=self.ModelEndpointsTable,
+            criteria=[self.ModelEndpointsTable.uid == endpoint_id],
         )
 
-    def delete_model_endpoint(self, endpoint_id: str):
+    def delete_model_endpoint(self, endpoint_id: str) -> None:
         """
         Deletes the SQL record of a given model endpoint id.
 
         :param endpoint_id: The unique id of the model endpoint.
         """
         self._init_model_endpoints_table()
-
-        filter_endpoint = {mm_schemas.EventFieldType.UID: endpoint_id}
         # Delete the model endpoint record using sqlalchemy ORM
-        self._delete(table=self.ModelEndpointsTable, **filter_endpoint)
+        self._delete(
+            table=self.ModelEndpointsTable,
+            criteria=[self.ModelEndpointsTable.uid == endpoint_id],
+        )
 
     def get_model_endpoint(
         self,
@@ -239,9 +251,11 @@ class SQLStoreBase(mlrun.model_monitoring.db.StoreBase):
         """
         self._init_model_endpoints_table()
 
-        # Get the model endpoint record using sqlalchemy ORM
-        filter_endpoint = {mm_schemas.EventFieldType.UID: endpoint_id}
-        endpoint_record = self._get(table=self.ModelEndpointsTable, **filter_endpoint)
+        # Get the model endpoint record
+        endpoint_record = self._get(
+            table=self.ModelEndpointsTable,
+            criteria=[self.ModelEndpointsTable.uid == endpoint_id],
+        )
 
         if not endpoint_record:
             raise mlrun.errors.MLRunNotFoundError(f"Endpoint {endpoint_id} not found")
@@ -277,7 +291,7 @@ class SQLStoreBase(mlrun.model_monitoring.db.StoreBase):
         endpoint_list = []
 
         model_endpoints_table = (
-            self.ModelEndpointsTable.__table__  # pyright: ignore[reportGeneralTypeIssues]
+            self.ModelEndpointsTable.__table__  # pyright: ignore[reportAttributeAccessIssue]
         )
 
         # Get the model endpoints records using sqlalchemy ORM
@@ -343,66 +357,81 @@ class SQLStoreBase(mlrun.model_monitoring.db.StoreBase):
         """
         Write a new application event in the target table.
 
-        :param event: An event dictionary that represents the application result, should be corresponded to the
-                      schema defined in the :py:class:`~mm_constants.constants.WriterEvent`
-                      object.
+        :param event: An event dictionary that represents the application result or metric,
+                      should be corresponded to the schema defined in the
+                      :py:class:`~mm_constants.constants.WriterEvent` object.
         :param kind: The type of the event, can be either "result" or "metric".
         """
 
         if kind == mm_schemas.WriterEventKind.METRIC:
-            # TODO : Implement the logic for writing metrics to MySQL
-            return
-
-        self._init_application_results_table()
+            self._init_application_metrics_table()
+            table = self.application_metrics_table
+            table_name = mm_schemas.FileTargetKind.APP_METRICS
+        elif kind == mm_schemas.WriterEventKind.RESULT:
+            self._init_application_results_table()
+            table = self.application_results_table
+            table_name = mm_schemas.FileTargetKind.APP_RESULTS
+        else:
+            raise ValueError(f"Invalid {kind = }")
 
-        application_filter_dict = {
-            mm_schemas.EventFieldType.UID: self._generate_application_result_uid(event)
-        }
+        application_result_uid = self._generate_application_result_uid(event, kind=kind)
+        criteria = [table.uid == application_result_uid]
 
-        application_record = self._get(
-            table=self.ApplicationResultsTable, **application_filter_dict
-        )
+        application_record = self._get(table=table, criteria=criteria)
         if application_record:
             self._convert_to_datetime(
-                event=event,
-                key=mm_schemas.WriterEvent.START_INFER_TIME,
+                event=event, key=mm_schemas.WriterEvent.START_INFER_TIME
             )
             self._convert_to_datetime(
-                event=event,
-                key=mm_schemas.WriterEvent.END_INFER_TIME,
+                event=event, key=mm_schemas.WriterEvent.END_INFER_TIME
             )
             # Update an existing application result
-            self._update(
-                attributes=event,
-                table=self.ApplicationResultsTable,
-                **application_filter_dict,
-            )
+            self._update(attributes=event, table=table, criteria=criteria)
         else:
             # Write a new application result
-            event[mm_schemas.EventFieldType.UID] = application_filter_dict[
-                mm_schemas.EventFieldType.UID
-            ]
-
-            self._write(
-                table=mm_schemas.FileTargetKind.APP_RESULTS,
-                event=event,
-            )
+            event[mm_schemas.EventFieldType.UID] = application_result_uid
+            self._write(table_name=table_name, event=event)
 
     @staticmethod
-    def _convert_to_datetime(event: dict[str, typing.Any], key: str):
+    def _convert_to_datetime(event: dict[str, typing.Any], key: str) -> None:
         if isinstance(event[key], str):
             event[key] = datetime.datetime.fromisoformat(event[key])
+        event[key] = event[key].astimezone(tz=datetime.timezone.utc)
 
     @staticmethod
-    def _generate_application_result_uid(event: dict[str, typing.Any]) -> str:
-        return (
-            event[mm_schemas.WriterEvent.ENDPOINT_ID]
-            + "_"
-            + event[mm_schemas.WriterEvent.APPLICATION_NAME]
-            + "_"
-            + event[mm_schemas.ResultData.RESULT_NAME]
+    def _generate_application_result_uid(
+        event: dict[str, typing.Any],
+        kind: mm_schemas.WriterEventKind = mm_schemas.WriterEventKind.RESULT,
+    ) -> str:
+        if kind == mm_schemas.WriterEventKind.RESULT:
+            name = event[mm_schemas.ResultData.RESULT_NAME]
+        else:
+            name = event[mm_schemas.MetricData.METRIC_NAME]
+        return "_".join(
+            [
+                event[mm_schemas.WriterEvent.ENDPOINT_ID],
+                event[mm_schemas.WriterEvent.APPLICATION_NAME],
+                name,
+            ]
         )
 
+    @staticmethod
+    def _get_filter_criteria(
+        *,
+        table: sqlalchemy.orm.decl_api.DeclarativeMeta,
+        endpoint_id: str,
+        application_name: typing.Optional[str] = None,
+    ) -> list[BinaryExpression]:
+        """
+        Return the filter criteria for the given endpoint_id and application_name.
+        Note: the table object must include the relevant columns:
+        `endpoint_id` and `application_name`.
+        """
+        criteria = [table.endpoint_id == endpoint_id]
+        if application_name is not None:
+            criteria.append(table.application_name == application_name)
+        return criteria
+
     def get_last_analyzed(self, endpoint_id: str, application_name: str) -> int:
         """
         Get the last analyzed time for the provided model endpoint and application.
@@ -412,14 +441,15 @@ class SQLStoreBase(mlrun.model_monitoring.db.StoreBase):
 
         :return: Timestamp as a Unix time.
         :raise:  MLRunNotFoundError if last analyzed value is not found.
-
         """
         self._init_monitoring_schedules_table()
-        application_filter_dict = self.filter_endpoint_and_application_name(
-            endpoint_id=endpoint_id, application_name=application_name
-        )
         monitoring_schedule_record = self._get(
-            table=self.MonitoringSchedulesTable, **application_filter_dict
+            table=self.MonitoringSchedulesTable,
+            criteria=self._get_filter_criteria(
+                table=self.MonitoringSchedulesTable,
+                endpoint_id=endpoint_id,
+                application_name=application_name,
+            ),
         )
         if not monitoring_schedule_record:
             raise mlrun.errors.MLRunNotFoundError(
@@ -441,16 +471,18 @@ class SQLStoreBase(mlrun.model_monitoring.db.StoreBase):
         """
         self._init_monitoring_schedules_table()
 
-        application_filter_dict = self.filter_endpoint_and_application_name(
-            endpoint_id=endpoint_id, application_name=application_name
+        criteria = self._get_filter_criteria(
+            table=self.MonitoringSchedulesTable,
+            endpoint_id=endpoint_id,
+            application_name=application_name,
         )
         monitoring_schedule_record = self._get(
-            table=self.MonitoringSchedulesTable, **application_filter_dict
+            table=self.MonitoringSchedulesTable, criteria=criteria
         )
         if not monitoring_schedule_record:
-            # Add a new record with empty last analyzed value
+            # Add a new record with last analyzed value
             self._write(
-                table=mm_schemas.FileTargetKind.MONITORING_SCHEDULES,
+                table_name=mm_schemas.FileTargetKind.MONITORING_SCHEDULES,
                 event={
                     mm_schemas.SchedulingKeys.UID: uuid.uuid4().hex,
                     mm_schemas.SchedulingKeys.APPLICATION_NAME: application_name,
@@ -462,30 +494,44 @@ class SQLStoreBase(mlrun.model_monitoring.db.StoreBase):
         self._update(
             attributes={mm_schemas.SchedulingKeys.LAST_ANALYZED: last_analyzed},
             table=self.MonitoringSchedulesTable,
-            **application_filter_dict,
+            criteria=criteria,
         )
 
-    def _delete_last_analyzed(self, endpoint_id: str = "", application_name: str = ""):
+    def _delete_last_analyzed(
+        self, endpoint_id: str, application_name: typing.Optional[str] = None
+    ) -> None:
         self._init_monitoring_schedules_table()
-
-        application_filter_dict = self.filter_endpoint_and_application_name(
-            endpoint_id=endpoint_id, application_name=application_name
+        criteria = self._get_filter_criteria(
+            table=self.MonitoringSchedulesTable,
+            endpoint_id=endpoint_id,
+            application_name=application_name,
         )
-
         # Delete the model endpoint record using sqlalchemy ORM
-        self._delete(table=self.MonitoringSchedulesTable, **application_filter_dict)
+        self._delete(table=self.MonitoringSchedulesTable, criteria=criteria)
 
     def _delete_application_result(
-        self, endpoint_id: str = "", application_name: str = ""
-    ):
+        self, endpoint_id: str, application_name: typing.Optional[str] = None
+    ) -> None:
         self._init_application_results_table()
-
-        application_filter_dict = self.filter_endpoint_and_application_name(
-            endpoint_id=endpoint_id, application_name=application_name
+        criteria = self._get_filter_criteria(
+            table=self.application_results_table,
+            endpoint_id=endpoint_id,
+            application_name=application_name,
         )
+        # Delete the relevant records from the results table
+        self._delete(table=self.application_results_table, criteria=criteria)
 
-        # Delete the model endpoint record using sqlalchemy ORM
-        self._delete(table=self.ApplicationResultsTable, **application_filter_dict)
+    def _delete_application_metrics(
+        self, endpoint_id: str, application_name: typing.Optional[str] = None
+    ) -> None:
+        self._init_application_metrics_table()
+        criteria = self._get_filter_criteria(
+            table=self.application_metrics_table,
+            endpoint_id=endpoint_id,
+            application_name=application_name,
+        )
+        # Delete the relevant records from the metrics table
+        self._delete(table=self.application_metrics_table, criteria=criteria)
 
     def _create_tables_if_not_exist(self):
         self._init_tables()
@@ -493,9 +539,7 @@ class SQLStoreBase(mlrun.model_monitoring.db.StoreBase):
         for table in self._tables:
             # Create table if not exist. The `metadata` contains the `ModelEndpointsTable`
             if not self._engine.has_table(table):
-                self._tables[table].metadata.create_all(  # pyright: ignore[reportGeneralTypeIssues]
-                    bind=self._engine
-                )
+                self._tables[table].metadata.create_all(bind=self._engine)
 
     @staticmethod
     def _filter_values(
@@ -572,31 +616,12 @@ class SQLStoreBase(mlrun.model_monitoring.db.StoreBase):
 
         return True
 
-    @staticmethod
-    def filter_endpoint_and_application_name(
-        endpoint_id: str, application_name: str
-    ) -> dict[str, str]:
-        """Generate a dictionary filter for endpoint id and application name"""
-        if not endpoint_id and not application_name:
-            raise mlrun.errors.MLRunBadRequestError(
-                "Please provide a valid endpoint_id and/or application_name"
-            )
-        application_filter_dict = {}
-        if endpoint_id:
-            application_filter_dict[mm_schemas.SchedulingKeys.ENDPOINT_ID] = endpoint_id
-        if application_name:
-            application_filter_dict[mm_schemas.SchedulingKeys.APPLICATION_NAME] = (
-                application_name
-            )
-        return application_filter_dict
-
-    def delete_model_endpoints_resources(self):
+    def delete_model_endpoints_resources(self) -> None:
         """
-        Delete all model endpoints resources in both SQL and the time series DB.
-
+        Delete all the model monitoring resources of the project in the SQL tables.
         """
-
         endpoints = self.list_model_endpoints()
+        logger.debug("Deleting model monitoring resources", project=self.project)
 
         for endpoint_dict in endpoints:
             endpoint_id = endpoint_dict[mm_schemas.EventFieldType.UID]
@@ -604,8 +629,9 @@ class SQLStoreBase(mlrun.model_monitoring.db.StoreBase):
             # Delete last analyzed records
             self._delete_last_analyzed(endpoint_id=endpoint_id)
 
-            # Delete application results records
+            # Delete application results and metrics records
             self._delete_application_result(endpoint_id=endpoint_id)
+            self._delete_application_metrics(endpoint_id=endpoint_id)
 
             # Delete model endpoint record
             self.delete_model_endpoint(endpoint_id=endpoint_id)
@@ -613,4 +639,46 @@ class SQLStoreBase(mlrun.model_monitoring.db.StoreBase):
     def get_model_endpoint_metrics(
         self, endpoint_id: str, type: mm_schemas.ModelEndpointMonitoringMetricType
     ) -> list[mm_schemas.ModelEndpointMonitoringMetric]:
-        raise NotImplementedError
+        """
+        Fetch the model endpoint metrics or results (according to `type`) for the
+        requested endpoint.
+        """
+        logger.debug(
+            "Fetching metrics for model endpoint",
+            project=self.project,
+            endpoint_id=endpoint_id,
+            type=type,
+        )
+        if type == mm_schemas.ModelEndpointMonitoringMetricType.METRIC:
+            self._init_application_metrics_table()
+            table = self.application_metrics_table
+            name_col = mm_schemas.MetricData.METRIC_NAME
+        else:
+            self._init_application_results_table()
+            table = self.application_results_table
+            name_col = mm_schemas.ResultData.RESULT_NAME
+
+        # Note: the block below does not use self._get, as we need here all the
+        # results, not only `one_or_none`.
+        with sqlalchemy.orm.Session(self._engine) as session:
+            metric_rows = (
+                session.query(table)  # pyright: ignore[reportOptionalCall]
+                .filter(table.endpoint_id == endpoint_id)
+                .all()
+            )
+
+        return [
+            mm_schemas.ModelEndpointMonitoringMetric(
+                project=self.project,
+                app=metric_row.application_name,
+                type=type,
+                name=getattr(metric_row, name_col),
+                full_name=mlrun.model_monitoring.helpers._compose_full_name(
+                    project=self.project,
+                    app=metric_row.application_name,
+                    type=type,
+                    name=getattr(metric_row, name_col),
+                ),
+            )
+            for metric_row in metric_rows
+        ]

mlrun/model_monitoring/db/tsdb/__init__.py

@@ -65,7 +65,7 @@ class ObjectTSDBFactory(enum.Enum):
 def get_tsdb_connector(
     project: str,
     tsdb_connector_type: str = "",
-    secret_provider: typing.Callable = None,
+    secret_provider: typing.Optional[typing.Callable] = None,
     **kwargs,
 ) -> TSDBConnector:
     """