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

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

@@ -1,344 +0,0 @@
-# Copyright 2024 Iguazio
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#   http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# 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 datetime
-from dataclasses import dataclass
-from io import StringIO
-from typing import Optional, Union
-
-import taosws
-
-import mlrun.common.schemas.model_monitoring as mm_schemas
-import mlrun.common.types
-
-_MODEL_MONITORING_DATABASE = "mlrun_model_monitoring"
-
-
-class _TDEngineColumnType:
-    def __init__(self, data_type: str, length: Optional[int] = None):
-        self.data_type = data_type
-        self.length = length
-
-    def values_to_column(self, values):
-        raise NotImplementedError()
-
-    def __str__(self):
-        if self.length is not None:
-            return f"{self.data_type}({self.length})"
-        else:
-            return self.data_type
-
-
-class _TDEngineColumn(mlrun.common.types.StrEnum):
-    TIMESTAMP = _TDEngineColumnType("TIMESTAMP")
-    FLOAT = _TDEngineColumnType("FLOAT")
-    INT = _TDEngineColumnType("INT")
-    BINARY_40 = _TDEngineColumnType("BINARY", 40)
-    BINARY_64 = _TDEngineColumnType("BINARY", 64)
-    BINARY_1000 = _TDEngineColumnType("BINARY", 1000)
-
-
-def values_to_column(values, column_type):
-    if column_type == _TDEngineColumn.TIMESTAMP:
-        timestamps = [round(timestamp.timestamp() * 1000) for timestamp in values]
-        return taosws.millis_timestamps_to_column(timestamps)
-    if column_type == _TDEngineColumn.FLOAT:
-        return taosws.floats_to_column(values)
-    if column_type == _TDEngineColumn.INT:
-        return taosws.ints_to_column(values)
-    if column_type == _TDEngineColumn.BINARY_40:
-        return taosws.binary_to_column(values)
-    if column_type == _TDEngineColumn.BINARY_64:
-        return taosws.binary_to_column(values)
-    if column_type == _TDEngineColumn.BINARY_1000:
-        return taosws.binary_to_column(values)
-
-    raise mlrun.errors.MLRunInvalidArgumentError(
-        f"unsupported column type '{column_type}'"
-    )
-
-
-@dataclass
-class TDEngineSchema:
-    """
-    A class to represent a supertable schema in TDengine. Using this schema, you can generate the relevant queries to
-    create, insert, delete and query data from TDengine. At the moment, there are 3 schemas: AppResultTable,
-    Metrics, and Predictions.
-    """
-
-    def __init__(
-        self,
-        super_table: str,
-        columns: dict[str, _TDEngineColumn],
-        tags: dict[str, str],
-        project: str,
-        database: Optional[str] = None,
-    ):
-        self.super_table = f"{super_table}_{project.replace('-', '_')}"
-        self.columns = columns
-        self.tags = tags
-        self.database = database or _MODEL_MONITORING_DATABASE
-
-    def _create_super_table_query(self) -> str:
-        columns = ", ".join(f"{col} {val}" for col, val in self.columns.items())
-        tags = ", ".join(f"{col} {val}" for col, val in self.tags.items())
-        return f"CREATE STABLE if NOT EXISTS {self.database}.{self.super_table} ({columns}) TAGS ({tags});"
-
-    def _create_subtable_sql(
-        self,
-        subtable: str,
-        values: dict[str, Union[str, int, float, datetime.datetime]],
-    ) -> str:
-        try:
-            tags = ", ".join(f"'{values[val]}'" for val in self.tags)
-        except KeyError:
-            raise mlrun.errors.MLRunInvalidArgumentError(
-                f"values must contain all tags: {self.tags.keys()}"
-            )
-        return f"CREATE TABLE if NOT EXISTS {self.database}.{subtable} USING {self.super_table} TAGS ({tags});"
-
-    def _delete_subtable_query(
-        self,
-        subtable: str,
-        values: dict[str, Union[str, int, float, datetime.datetime]],
-    ) -> str:
-        values = " AND ".join(
-            f"{val} LIKE '{values[val]}'" for val in self.tags if val in values
-        )
-        if not values:
-            raise mlrun.errors.MLRunInvalidArgumentError(
-                f"values must contain at least one tag: {self.tags.keys()}"
-            )
-        return f"DELETE FROM {self.database}.{subtable} WHERE {values};"
-
-    def drop_subtable_query(
-        self,
-        subtable: str,
-    ) -> str:
-        return f"DROP TABLE if EXISTS {self.database}.`{subtable}`;"
-
-    def drop_supertable_query(self) -> str:
-        return f"DROP STABLE if EXISTS {self.database}.{self.super_table};"
-
-    def _get_subtables_query_by_tag(
-        self,
-        filter_tag: str,
-        filter_values: list[str],
-        operator: str = "OR",
-    ) -> str:
-        if filter_tag not in self.tags:
-            raise mlrun.errors.MLRunInvalidArgumentError(
-                f"`filter_tag` must be one of the tags: {self.tags.keys()}"
-            )
-
-        values = f" {operator} ".join(
-            f"{filter_tag} LIKE '{val}'" for val in filter_values
-        )
-
-        return f"SELECT DISTINCT tbname FROM {self.database}.{self.super_table} WHERE {values};"
-
-    @staticmethod
-    def _get_records_query(
-        table: str,
-        start: datetime.datetime,
-        end: datetime.datetime,
-        columns_to_filter: Optional[list[str]] = None,
-        filter_query: Optional[str] = None,
-        interval: Optional[str] = None,
-        limit: int = 0,
-        agg_funcs: Optional[list] = None,
-        sliding_window_step: Optional[str] = None,
-        timestamp_column: str = "time",
-        database: str = _MODEL_MONITORING_DATABASE,
-        group_by: Optional[Union[list[str], str]] = None,
-        preform_agg_funcs_columns: Optional[list[str]] = None,
-        order_by: Optional[str] = None,
-        desc: Optional[bool] = None,
-        partition_by: Optional[str] = None,
-    ) -> str:
-        if agg_funcs and not columns_to_filter:
-            raise mlrun.errors.MLRunInvalidArgumentError(
-                "`columns_to_filter` must be provided when using aggregate functions"
-            )
-
-        # if aggregate function or interval is provided, the other must be provided as well
-        if interval and not agg_funcs:
-            raise mlrun.errors.MLRunInvalidArgumentError(
-                "`agg_funcs` must be provided when using interval"
-            )
-        if partition_by and not agg_funcs:
-            raise mlrun.errors.MLRunInvalidArgumentError(
-                "`agg_funcs` must be provided when using partition by"
-            )
-        if sliding_window_step and not interval:
-            raise mlrun.errors.MLRunInvalidArgumentError(
-                "`interval` must be provided when using sliding window"
-            )
-        if group_by and not agg_funcs:
-            raise mlrun.errors.MLRunInvalidArgumentError(
-                "aggregate functions must be provided when using group by"
-            )
-        if desc and not order_by:
-            raise mlrun.errors.MLRunInvalidArgumentError(
-                "`order_by` must be provided when using descending"
-            )
-
-        with StringIO() as query:
-            query.write("SELECT ")
-            if interval:
-                query.write("_wstart, _wend, ")
-            if agg_funcs:
-                preform_agg_funcs_columns = (
-                    columns_to_filter
-                    if preform_agg_funcs_columns is None
-                    else preform_agg_funcs_columns
-                )
-                query.write(
-                    ", ".join(
-                        [
-                            f"{a}({col})"
-                            if col.upper()
-                            in map(
-                                str.upper, preform_agg_funcs_columns
-                            )  # Case-insensitive check
-                            else f"{col}"
-                            for a in agg_funcs
-                            for col in columns_to_filter
-                        ]
-                    )
-                )
-            elif columns_to_filter:
-                query.write(", ".join(columns_to_filter))
-            else:
-                query.write("*")
-            query.write(f" FROM {database}.{table}")
-
-            if any([filter_query, start, end]):
-                query.write(" WHERE ")
-                if filter_query:
-                    query.write(f"{filter_query} AND ")
-                if start:
-                    query.write(f"{timestamp_column} >= '{start}' AND ")
-                if end:
-                    query.write(f"{timestamp_column} <= '{end}'")
-            if group_by:
-                if isinstance(group_by, list):
-                    group_by = ", ".join(group_by)
-                query.write(f" GROUP BY {group_by}")
-            if partition_by:
-                query.write(f" PARTITION BY {partition_by}")
-            if order_by:
-                desc = " DESC" if desc else ""
-                query.write(f" ORDER BY {order_by}{desc}")
-            if interval:
-                query.write(f" INTERVAL({interval})")
-            if sliding_window_step:
-                query.write(f" SLIDING({sliding_window_step})")
-            if limit:
-                query.write(f" LIMIT {limit}")
-            query.write(";")
-            return query.getvalue()
-
-
-@dataclass
-class AppResultTable(TDEngineSchema):
-    def __init__(self, project: str, database: Optional[str] = None):
-        super_table = mm_schemas.TDEngineSuperTables.APP_RESULTS
-        columns = {
-            mm_schemas.WriterEvent.END_INFER_TIME: _TDEngineColumn.TIMESTAMP,
-            mm_schemas.WriterEvent.START_INFER_TIME: _TDEngineColumn.TIMESTAMP,
-            mm_schemas.ResultData.RESULT_VALUE: _TDEngineColumn.FLOAT,
-            mm_schemas.ResultData.RESULT_STATUS: _TDEngineColumn.INT,
-            mm_schemas.ResultData.RESULT_EXTRA_DATA: _TDEngineColumn.BINARY_1000,
-        }
-        tags = {
-            mm_schemas.WriterEvent.ENDPOINT_ID: _TDEngineColumn.BINARY_64,
-            mm_schemas.WriterEvent.APPLICATION_NAME: _TDEngineColumn.BINARY_64,
-            mm_schemas.ResultData.RESULT_NAME: _TDEngineColumn.BINARY_64,
-            mm_schemas.ResultData.RESULT_KIND: _TDEngineColumn.INT,
-        }
-        super().__init__(
-            super_table=super_table,
-            columns=columns,
-            tags=tags,
-            database=database,
-            project=project,
-        )
-
-
-@dataclass
-class Metrics(TDEngineSchema):
-    def __init__(self, project: str, database: Optional[str] = None):
-        super_table = mm_schemas.TDEngineSuperTables.METRICS
-        columns = {
-            mm_schemas.WriterEvent.END_INFER_TIME: _TDEngineColumn.TIMESTAMP,
-            mm_schemas.WriterEvent.START_INFER_TIME: _TDEngineColumn.TIMESTAMP,
-            mm_schemas.MetricData.METRIC_VALUE: _TDEngineColumn.FLOAT,
-        }
-        tags = {
-            mm_schemas.WriterEvent.ENDPOINT_ID: _TDEngineColumn.BINARY_64,
-            mm_schemas.WriterEvent.APPLICATION_NAME: _TDEngineColumn.BINARY_64,
-            mm_schemas.MetricData.METRIC_NAME: _TDEngineColumn.BINARY_64,
-        }
-        super().__init__(
-            super_table=super_table,
-            columns=columns,
-            tags=tags,
-            database=database,
-            project=project,
-        )
-
-
-@dataclass
-class Predictions(TDEngineSchema):
-    def __init__(self, project: str, database: Optional[str] = None):
-        super_table = mm_schemas.TDEngineSuperTables.PREDICTIONS
-        columns = {
-            mm_schemas.EventFieldType.TIME: _TDEngineColumn.TIMESTAMP,
-            mm_schemas.EventFieldType.LATENCY: _TDEngineColumn.FLOAT,
-            mm_schemas.EventKeyMetrics.CUSTOM_METRICS: _TDEngineColumn.BINARY_1000,
-            mm_schemas.EventFieldType.ESTIMATED_PREDICTION_COUNT: _TDEngineColumn.FLOAT,
-            mm_schemas.EventFieldType.EFFECTIVE_SAMPLE_COUNT: _TDEngineColumn.INT,
-        }
-        tags = {
-            mm_schemas.WriterEvent.ENDPOINT_ID: _TDEngineColumn.BINARY_64,
-        }
-        super().__init__(
-            super_table=super_table,
-            columns=columns,
-            tags=tags,
-            database=database,
-            project=project,
-        )
-
-
-@dataclass
-class Errors(TDEngineSchema):
-    def __init__(self, project: str, database: Optional[str] = None):
-        super_table = mm_schemas.TDEngineSuperTables.ERRORS
-        columns = {
-            mm_schemas.EventFieldType.TIME: _TDEngineColumn.TIMESTAMP,
-            mm_schemas.EventFieldType.MODEL_ERROR: _TDEngineColumn.BINARY_1000,
-        }
-        tags = {
-            mm_schemas.WriterEvent.ENDPOINT_ID: _TDEngineColumn.BINARY_64,
-            mm_schemas.EventFieldType.ERROR_TYPE: _TDEngineColumn.BINARY_64,
-        }
-        super().__init__(
-            super_table=super_table,
-            columns=columns,
-            tags=tags,
-            database=database,
-            project=project,
-        )

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

@@ -1,75 +0,0 @@
-# Copyright 2024 Iguazio
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#   http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# 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 json
-from datetime import datetime
-
-import mlrun.feature_store.steps
-from mlrun.common.schemas.model_monitoring import (
-    EventFieldType,
-    EventKeyMetrics,
-)
-from mlrun.utils import logger
-
-
-class ProcessBeforeTDEngine(mlrun.feature_store.steps.MapClass):
-    def __init__(self, **kwargs):
-        """
-        Process the data before writing to TDEngine. This step create the relevant keys for the TDEngine table,
-        including project name, custom metrics, time column, and table name column.
-
-        :returns: Event as a dictionary which will be written into the TDEngine Predictions table.
-        """
-        super().__init__(**kwargs)
-
-    def do(self, event):
-        event[EventFieldType.PROJECT] = event[EventFieldType.FUNCTION_URI].split("/")[0]
-        event[EventKeyMetrics.CUSTOM_METRICS] = json.dumps(
-            event.get(EventFieldType.METRICS, {})
-        )
-        event[EventFieldType.TIME] = event.get(EventFieldType.TIMESTAMP)
-        event[EventFieldType.TABLE_COLUMN] = "_" + event.get(EventFieldType.ENDPOINT_ID)
-
-        return event
-
-
-class ErrorExtractor(mlrun.feature_store.steps.MapClass):
-    def __init__(self, **kwargs):
-        """
-        Prepare the event for insertion into the TDEngine error table
-        """
-        super().__init__(**kwargs)
-
-    def do(self, event):
-        error = str(event.get("error"))
-        if len(error) > 1000:
-            error = error[-1000:]
-            logger.warning(
-                f"Error message exceeds 1000 chars: The error message writen to TSDB will be it last "
-                f"1000 chars, Error:  {error}",
-                event=event,
-            )
-        timestamp = datetime.fromisoformat(event.get("when"))
-        endpoint_id = event[EventFieldType.ENDPOINT_ID]
-        event = {
-            EventFieldType.MODEL_ERROR: error,
-            EventFieldType.ERROR_TYPE: EventFieldType.INFER_ERROR,
-            EventFieldType.ENDPOINT_ID: endpoint_id,
-            EventFieldType.TIME: timestamp,
-            EventFieldType.PROJECT: event[EventFieldType.FUNCTION_URI].split("/")[0],
-            EventFieldType.TABLE_COLUMN: "_err_"
-            + event.get(EventFieldType.ENDPOINT_ID),
-        }
-        logger.info("Write error to errors TSDB table", event=event)
-        return event

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

@@ -1,281 +0,0 @@
-# Copyright 2025 Iguazio
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#   http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# 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 time
-from collections.abc import Callable
-from enum import Enum
-from typing import Any, Final, Optional, Union
-
-import taosws
-from taosws import TaosStmt
-
-import mlrun
-from mlrun.utils import logger
-
-
-class _StrEnum(str, Enum):
-    pass
-
-
-class TimestampPrecision(_StrEnum):
-    ms = "ms"  # milliseconds
-    us = "us"  # microseconds
-    ns = "ns"  # nanoseconds
-
-
-_TS_PRECISION_TO_FACTOR_AND_FUNC: Final[
-    dict[TimestampPrecision, tuple[int, Callable[[list[int]], taosws.PyColumnView]]]
-] = {
-    TimestampPrecision.ms: (10**3, taosws.millis_timestamps_to_column),
-    TimestampPrecision.us: (10**6, taosws.micros_timestamps_to_column),
-    TimestampPrecision.ns: (10**9, taosws.nanos_timestamps_to_column),
-}
-
-
-class QueryResult:
-    def __init__(self, data, fields):
-        self.data = data
-        self.fields = fields
-
-    def __eq__(self, other):
-        return self.data == other.data and self.fields == other.fields
-
-    def __repr__(self):
-        return f"QueryResult({self.data}, {self.fields})"
-
-
-class Field:
-    def __init__(self, name, type, bytes):
-        self.name = name
-        self.type = type
-        self.bytes = bytes
-
-    def __eq__(self, other):
-        return (
-            self.name == other.name
-            and self.type == other.type
-            and self.bytes == other.bytes
-        )
-
-    def __repr__(self):
-        return f"Field({self.name}, {self.type}, {self.bytes})"
-
-
-class TDEngineError(Exception):
-    pass
-
-
-class ErrorResult:
-    def __init__(self, tb, err):
-        self.tb = tb
-        self.err = err
-
-
-def _get_timestamp_column(
-    values: list, timestamp_precision: TimestampPrecision
-) -> taosws.PyColumnView:
-    factor, to_col_func = _TS_PRECISION_TO_FACTOR_AND_FUNC[timestamp_precision]
-    timestamps = [round(timestamp.timestamp() * factor) for timestamp in values]
-    return to_col_func(timestamps)
-
-
-def values_to_column(
-    values: list,
-    column_type: str,
-    timestamp_precision: TimestampPrecision = TimestampPrecision.ms,
-) -> taosws.PyColumnView:
-    if column_type == "TIMESTAMP":
-        return _get_timestamp_column(values, timestamp_precision)
-    if column_type == "FLOAT":
-        return taosws.floats_to_column(values)
-    if column_type == "INT":
-        return taosws.ints_to_column(values)
-    if column_type.startswith("BINARY"):
-        return taosws.binary_to_column(values)
-
-    raise NotImplementedError(f"Unsupported column type '{column_type}'")
-
-
-class Statement:
-    def __init__(
-        self,
-        columns: dict[str, str],
-        subtable: str,
-        values: dict[str, Any],
-        timestamp_precision: str = TimestampPrecision.ms,
-    ) -> None:
-        self.columns = columns
-        self.subtable = subtable
-        self.values = values
-        self.timestamp_precision = TimestampPrecision[timestamp_precision]
-
-    def prepare(self, statement: TaosStmt) -> TaosStmt:
-        question_marks = ", ".join("?" * len(self.columns))
-        statement.prepare(f"INSERT INTO ? VALUES ({question_marks});")
-        statement.set_tbname(self.subtable)
-
-        bind_params = []
-
-        for col_name, col_type in self.columns.items():
-            val = self.values[col_name]
-            bind_params.append(
-                values_to_column(
-                    [val], col_type, timestamp_precision=self.timestamp_precision
-                )
-            )
-
-        statement.bind_param(bind_params)
-        statement.add_batch()
-        return statement
-
-
-class TDEngineConnection:
-    def __init__(self, connection_string, max_retries=3, retry_delay=0.5):
-        self._connection_string = connection_string
-        self.prefix_statements = []
-        self._max_retries = max_retries
-        self._retry_delay = retry_delay
-
-        self._conn = self._create_connection()
-
-    def _create_connection(self):
-        """Create a new TDEngine connection."""
-        return taosws.connect(self._connection_string)
-
-    def _reconnect(self):
-        """Close current connection and create a new one."""
-        try:
-            if hasattr(self, "_conn") and self._conn:
-                self._conn.close()
-        except Exception as e:
-            logger.warning(f"Error closing connection during reconnect: {e}")
-
-        self._conn = self._create_connection()
-        logger.info("Successfully reconnected to TDEngine")
-
-    def _execute_with_retry(self, operation, operation_name, *args, **kwargs):
-        """
-        Execute an operation with retry logic for connection failures.
-
-        :param operation: The function to execute
-        :param operation_name: Name of the operation for logging
-        :param args: Arguments to pass to the operation
-        :param kwargs: Keyword arguments to pass to the operation
-        :return: Result of the operation
-        """
-        last_exception = None
-
-        for attempt in range(self._max_retries + 1):  # +1 for initial attempt
-            try:
-                return operation(*args, **kwargs)
-
-            except taosws.Error as e:
-                last_exception = e
-
-                if attempt < self._max_retries:
-                    logger.warning(
-                        f"Connection error during {operation_name} "
-                        f"(attempt {attempt + 1}/{self._max_retries + 1}): {e}. "
-                        f"Retrying in {self._retry_delay} seconds..."
-                    )
-
-                    # Wait before retrying
-                    time.sleep(self._retry_delay)
-
-                    # Reconnect
-                    try:
-                        self._reconnect()
-                    except Exception as reconnect_error:
-                        logger.error(f"Failed to reconnect: {reconnect_error}")
-                        if attempt == self._max_retries - 1:
-                            # Last attempt, raise the reconnection error
-                            raise TDEngineError(
-                                f"Failed to reconnect after {operation_name} failure: {reconnect_error}"
-                            ) from reconnect_error
-                        continue
-                else:
-                    # Max retries exceeded
-                    logger.error(
-                        f"Max retries ({self._max_retries}) exceeded for {operation_name}"
-                    )
-                    break
-
-            except Exception as e:
-                # Non-TDEngine error, don't retry
-                raise TDEngineError(
-                    f"Unexpected error during {operation_name}: {e}"
-                ) from e
-
-        # If we get here, all retries failed
-        raise TDEngineError(
-            f"Failed to {operation_name} after {self._max_retries} retries: {last_exception}"
-        ) from last_exception
-
-    def _execute_statement(self, statement):
-        """Execute a single statement (string or Statement object)."""
-        if isinstance(statement, Statement):
-            prepared_statement = statement.prepare(self._conn.statement())
-            prepared_statement.execute()
-        else:
-            self._conn.execute(statement)
-
-    def _execute_query(self, query):
-        """Execute a query and return the result."""
-        return self._conn.query(query)
-
-    def run(
-        self,
-        statements: Optional[Union[str, Statement, list[Union[str, Statement]]]] = None,
-        query: Optional[str] = None,
-    ) -> Optional[QueryResult]:
-        statements = statements or []
-        if not isinstance(statements, list):
-            statements = [statements]
-
-        # Execute all statements with retry logic
-        all_statements = self.prefix_statements + statements
-        for i, statement in enumerate(all_statements):
-            operation_name = f"execute statement {i + 1}/{len(all_statements)}"
-            if isinstance(statement, Statement):
-                operation_name += " (prepared)"
-            else:
-                operation_name += f" `{statement}`"
-
-            self._execute_with_retry(self._execute_statement, operation_name, statement)
-
-        if not query:
-            return None
-
-        # Execute query with retry logic
-        res = self._execute_with_retry(
-            self._execute_query, f"execute query `{query}`", query
-        )
-
-        # Process results
-        fields = [
-            Field(field.name(), field.type(), field.bytes()) for field in res.fields
-        ]
-
-        return QueryResult(list(res), fields)
-
-    def close(self):
-        """Close the connection."""
-        try:
-            if self._conn:
-                self._conn.close()
-                logger.debug("TDEngine connection closed")
-                self._conn = None
-        except Exception as e:
-            logger.warning(
-                f"Error closing TDEngine connection: {mlrun.errors.err_to_str(e)}"
-            )