[Model Monitoring] Fix SQL store with multiple filter criteria (#5629)

mlrun/model_monitoring/db/stores/sqldb/models/__init__.py

@@ -12,7 +12,8 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from typing import Optional, Union
+from functools import partial
+from typing import Optional, TypeVar, Union
 
 from .mysql import ApplicationResultTable as MySQLApplicationResultTable
 from .mysql import ModelEndpointsTable as MySQLModelEndpointsTable
@@ -21,29 +22,42 @@
 from .sqlite import ModelEndpointsTable as SQLiteModelEndpointsTable
 from .sqlite import MonitoringSchedulesTable as SQLiteMonitoringSchedulesTable
 
+MySQLTableType = TypeVar("MySQLTableType")
+SQLiteTableType = TypeVar("SQLiteTableType")
 
-def _get_model_endpoints_table(
-    connection_string: Optional[str] = None,
-) -> Union[type[MySQLModelEndpointsTable], type[SQLiteModelEndpointsTable]]:
-    """Return ModelEndpointsTable based on the provided connection string"""
-    if connection_string and "mysql:" in connection_string:
-        return MySQLModelEndpointsTable
-    return SQLiteModelEndpointsTable
-
-
-def _get_application_result_table(
-    connection_string: Optional[str] = None,
-) -> Union[type[MySQLApplicationResultTable], type[SQLiteApplicationResultTable]]:
-    """Return ModelEndpointsTable based on the provided connection string"""
-    if connection_string and "mysql:" in connection_string:
-        return MySQLApplicationResultTable
-    return SQLiteApplicationResultTable
+_MYSQL_SCHEME = "mysql:"
 
 
-def _get_monitoring_schedules_table(
+def _get_sql_table(
+    *,
+    mysql_table: MySQLTableType,
+    sqlite_table: SQLiteTableType,
     connection_string: Optional[str] = None,
-) -> Union[type[MySQLMonitoringSchedulesTable], type[SQLiteMonitoringSchedulesTable]]:
-    """Return ModelEndpointsTable based on the provided connection string"""
-    if connection_string and "mysql:" in connection_string:
-        return MySQLMonitoringSchedulesTable
-    return SQLiteMonitoringSchedulesTable
+) -> Union[MySQLTableType, SQLiteTableType]:
+    """
+    Return a SQLAlchemy table for MySQL or SQLite according to the connection string.
+    Note: this function should not be directly used in other modules.
+    """
+    if connection_string and _MYSQL_SCHEME in connection_string:
+        return mysql_table
+    return sqlite_table
+
+
+_get_model_endpoints_table = partial(
+    _get_sql_table,
+    mysql_table=MySQLModelEndpointsTable,
+    sqlite_table=SQLiteModelEndpointsTable,
+)
+
+
+_get_application_result_table = partial(
+    _get_sql_table,
+    mysql_table=MySQLApplicationResultTable,
+    sqlite_table=SQLiteApplicationResultTable,
+)
+
+_get_monitoring_schedules_table = partial(
+    _get_sql_table,
+    mysql_table=MySQLMonitoringSchedulesTable,
+    sqlite_table=SQLiteMonitoringSchedulesTable,
+)