MindsDB 25.4.5.0__py3-none-any.whl → 25.5.4.0__py3-none-any.whl

mindsdb/api/mysql/mysql_proxy/mysql_proxy.py

@@ -21,24 +21,17 @@ import sys
 import tempfile
 import traceback
 from functools import partial
-from typing import Dict, List, Optional
+from typing import List
 from dataclasses import dataclass
 
-from numpy import dtype as np_dtype
-from pandas.api import types as pd_types
-
-from mindsdb.api.mysql.mysql_proxy.libs.constants.mysql import NULL_VALUE
 from mindsdb.api.mysql.mysql_proxy.data_types.mysql_datum import Datum
-
 import mindsdb.utilities.hooks as hooks
 import mindsdb.utilities.profiler as profiler
+from mindsdb.utilities.sql import clear_sql
 from mindsdb.api.mysql.mysql_proxy.classes.client_capabilities import ClentCapabilities
 from mindsdb.api.mysql.mysql_proxy.classes.server_capabilities import (
     server_capabilities,
 )
-from mindsdb.api.mysql.mysql_proxy.classes.sql_statement_parser import (
-    SqlStatementParser,
-)
 from mindsdb.api.executor.controllers import SessionController
 from mindsdb.api.mysql.mysql_proxy.data_types.mysql_packet import Packet
 from mindsdb.api.mysql.mysql_proxy.data_types.mysql_packets import (
@@ -63,14 +56,14 @@ from mindsdb.api.mysql.mysql_proxy.external_libs.mysql_scramble import (
     scramble as scramble_func,
 )
 from mindsdb.api.mysql.mysql_proxy.libs.constants.mysql import (
-    CAPABILITIES,
+    DEFAULT_AUTH_METHOD,
     CHARSET_NUMBERS,
+    SERVER_STATUS,
+    CAPABILITIES,
+    NULL_VALUE,
     COMMANDS,
-    DEFAULT_AUTH_METHOD,
     ERR,
-    SERVER_STATUS,
-    TYPES,
-    getConstName,
+    getConstName
 )
 from mindsdb.api.executor.data_types.answer import ExecuteAnswer
 from mindsdb.api.executor.data_types.response_type import RESPONSE_TYPE
@@ -81,12 +74,14 @@ from mindsdb.api.mysql.mysql_proxy.utilities import (
 from mindsdb.api.executor import exceptions as exec_exc
 
 from mindsdb.api.common.check_auth import check_auth
-from mindsdb.api.mysql.mysql_proxy.utilities.lightwood_dtype import dtype
+from mindsdb.api.mysql.mysql_proxy.libs.constants.mysql import MYSQL_DATA_TYPE
+from mindsdb.api.executor.sql_query.result_set import Column, ResultSet
 from mindsdb.utilities import log
 from mindsdb.utilities.config import config
 from mindsdb.utilities.context import context as ctx
-from mindsdb.utilities.otel.metric_handlers import get_query_request_counter
+from mindsdb.utilities.otel import increment_otel_query_request_counter
 from mindsdb.utilities.wizards import make_ssl_cert
+from mindsdb.api.mysql.mysql_proxy.utilities.dump import dump_result_set_to_mysql, column_to_mysql_column_dict
 
 logger = log.getLogger(__name__)
 
@@ -98,18 +93,43 @@ def empty_fn():
 @dataclass
 class SQLAnswer:
     resp_type: RESPONSE_TYPE = RESPONSE_TYPE.OK
-    columns: Optional[List[Dict]] = None
-    data: Optional[List[Dict]] = None   # resultSet ?
-    status: Optional[int] = None
-    state_track: Optional[List[List]] = None
-    error_code: Optional[int] = None
-    error_message: Optional[str] = None
-    affected_rows: Optional[int] = None
+    result_set: ResultSet | None = None
+    status: int | None = None
+    state_track: List[List] | None = None
+    error_code: int | None = None
+    error_message: str | None = None
+    affected_rows: int | None = None
+    mysql_types: list[MYSQL_DATA_TYPE] | None = None
 
     @property
     def type(self):
         return self.resp_type
 
+    def dump_http_response(self) -> dict:
+        if self.resp_type == RESPONSE_TYPE.OK:
+            return {
+                "type": self.resp_type,
+                "affected_rows": self.affected_rows,
+            }
+        elif self.resp_type in (RESPONSE_TYPE.TABLE, RESPONSE_TYPE.COLUMNS_TABLE):
+            data = self.result_set.to_lists(json_types=True)
+            return {
+                "type": RESPONSE_TYPE.TABLE,
+                "data": data,
+                "column_names": [
+                    column.alias or column.name
+                    for column in self.result_set.columns
+                ],
+            }
+        elif self.resp_type == RESPONSE_TYPE.ERROR:
+            return {
+                "type": RESPONSE_TYPE.ERROR,
+                "error_code": self.error_code or 0,
+                "error_message": self.error_message,
+            }
+        else:
+            raise ValueError(f"Unsupported response type for dump HTTP response: {self.resp_type}")
+
 
 class MysqlProxy(SocketServer.BaseRequestHandler):
     """
@@ -312,14 +332,14 @@ class MysqlProxy(SocketServer.BaseRequestHandler):
         self.session.unregister_stmt(stmt_id)
 
     def send_query_answer(self, answer: SQLAnswer):
-        if answer.type == RESPONSE_TYPE.TABLE:
+        if answer.type in (RESPONSE_TYPE.TABLE, RESPONSE_TYPE.COLUMNS_TABLE):
             packages = []
 
-            if len(answer.data) > 1000:
+            if len(answer.result_set) > 1000:
                 # for big responses leverage pandas map function to convert data to packages
-                self.send_table_packets(columns=answer.columns, data=answer.data)
+                self.send_table_packets(result_set=answer.result_set)
             else:
-                packages += self.get_table_packets(columns=answer.columns, data=answer.data.to_lists())
+                packages += self.get_table_packets(result_set=answer.result_set)
 
             if answer.status is not None:
                 packages.append(self.last_packet(status=answer.status))
@@ -333,7 +353,7 @@ class MysqlProxy(SocketServer.BaseRequestHandler):
                 ErrPacket, err_code=answer.error_code, msg=answer.error_message
             ).send()
 
-    def _get_column_defenition_packets(self, columns, data=None, columns_len=None):
+    def _get_column_defenition_packets(self, columns: dict, data=None):
         if data is None:
             data = []
         packets = []
@@ -348,18 +368,16 @@ class MysqlProxy(SocketServer.BaseRequestHandler):
             column_name = column.get("name", "column_name")
             column_alias = column.get("alias", column_name)
             flags = column.get("flags", 0)
-            if columns_len is not None:
-                length = columns_len[i]
-            else:
-                if len(data) == 0:
-                    length = 0xFFFF
-                else:
-                    length = 1
-                    for row in data:
-                        if isinstance(row, dict):
-                            length = max(len(str(row[column_alias])), length)
-                        else:
-                            length = max(len(str(row[i])), length)
+            if isinstance(flags, list):
+                flags = sum(flags)
+            if column.get('size') is None:
+                length = 1
+                for row in data:
+                    if isinstance(row, dict):
+                        length = max(len(str(row[column_alias])), length)
+                    else:
+                        length = max(len(str(row[i])), length)
+                column['size'] = 1
 
             packets.append(
                 self.packet(
@@ -371,16 +389,19 @@ class MysqlProxy(SocketServer.BaseRequestHandler):
                     column_name=column_name,
                     column_type=column["type"],
                     charset=column.get("charset", CHARSET_NUMBERS["utf8_unicode_ci"]),
-                    max_length=length,
+                    max_length=column["size"],
                     flags=flags,
                 )
             )
         return packets
 
-    def get_table_packets(self, columns, data, status=0):
+    def get_table_packets(self, result_set: ResultSet, status=0):
+        data_frame, columns_dict = dump_result_set_to_mysql(result_set)
+        data = data_frame.to_dict('split')['data']
+
         # TODO remove columns order
-        packets = [self.packet(ColumnCountPacket, count=len(columns))]
-        packets.extend(self._get_column_defenition_packets(columns, data))
+        packets = [self.packet(ColumnCountPacket, count=len(columns_dict))]
+        packets.extend(self._get_column_defenition_packets(columns_dict, data))
 
         if self.client_capabilities.DEPRECATE_EOF is False:
             packets.append(self.packet(EofPacket, status=status))
@@ -388,9 +409,9 @@ class MysqlProxy(SocketServer.BaseRequestHandler):
         packets += [self.packet(ResultsetRowPacket, data=x) for x in data]
         return packets
 
-    def send_table_packets(self, columns, data, status=0):
+    def send_table_packets(self, result_set: ResultSet, status: int = 0):
+        df, columns_dicts = dump_result_set_to_mysql(result_set, infer_column_size=True)
         # text protocol, convert all to string and serialize as packages
-        df = data.get_raw_df()
 
         def apply_f(v):
             if v is None:
@@ -399,23 +420,10 @@ class MysqlProxy(SocketServer.BaseRequestHandler):
                 v = str(v)
             return Datum.serialize_str(v)
 
-        # get column max size
-        # column_len is used by mysql client to determine width of columns, so it is not mandatory
-        # to get exactly max value. We can approximate them by sample.
-        columns_len = None
-        if len(df) > 0:
-            sample = df.head(100)
-            columns_len = []
-            for column in sample.columns:
-                try:
-                    columns_len.append(sample[column].astype(str).str.len().max())
-                except Exception:
-                    columns_len.append(1)
-
         # columns packages
-        packets = [self.packet(ColumnCountPacket, count=len(columns))]
+        packets = [self.packet(ColumnCountPacket, count=len(columns_dicts))]
 
-        packets.extend(self._get_column_defenition_packets(columns, columns_len=columns_len))
+        packets.extend(self._get_column_defenition_packets(columns_dicts))
 
         if self.client_capabilities.DEPRECATE_EOF is False:
             packets.append(self.packet(EofPacket, status=status))
@@ -488,57 +496,12 @@ class MysqlProxy(SocketServer.BaseRequestHandler):
 
         return {"is_cloud": False}
 
-    def to_mysql_columns(self, columns_list):
-        """Converts raw columns data into convinient format(list of lists) for the futher usage.
-        Plus, it is also converts column types into internal ones."""
-
-        result = []
-
-        database = (
-            None if self.session.database == "" else self.session.database.lower()
-        )
-        for column_record in columns_list:
-
-            field_type = column_record.type
-
-            column_type = TYPES.MYSQL_TYPE_VAR_STRING
-            # is already in mysql protocol type?
-            if isinstance(field_type, int):
-                column_type = field_type
-            # pandas checks
-            elif isinstance(field_type, np_dtype):
-                if pd_types.is_integer_dtype(field_type):
-                    column_type = TYPES.MYSQL_TYPE_LONG
-                elif pd_types.is_numeric_dtype(field_type):
-                    column_type = TYPES.MYSQL_TYPE_DOUBLE
-                elif pd_types.is_datetime64_any_dtype(field_type):
-                    column_type = TYPES.MYSQL_TYPE_DATETIME
-            # lightwood checks
-            elif field_type == dtype.date:
-                column_type = TYPES.MYSQL_TYPE_DATE
-            elif field_type == dtype.datetime:
-                column_type = TYPES.MYSQL_TYPE_DATETIME
-            elif field_type == dtype.float:
-                column_type = TYPES.MYSQL_TYPE_DOUBLE
-            elif field_type == dtype.integer:
-                column_type = TYPES.MYSQL_TYPE_LONG
-
-            result.append(
-                {
-                    "database": column_record.database or database,
-                    #  TODO add 'original_table'
-                    "table_name": column_record.table_name,
-                    "name": column_record.name,
-                    "alias": column_record.alias or column_record.name,
-                    # NOTE all work with text-type, but if/when wanted change types to real,
-                    # it will need to check all types casts in BinaryResultsetRowPacket
-                    "type": column_type,
-                }
-            )
-        return result
+    def to_mysql_columns(self, columns_list: list[Column]) -> list[dict[str, str | int]]:
+        database_name = None if self.session.database == "" else self.session.database.lower()
+        return [column_to_mysql_column_dict(column, database_name=database_name) for column in columns_list]
 
     @profiler.profile()
-    def process_query(self, sql):
+    def process_query(self, sql) -> SQLAnswer:
         executor = Executor(session=self.session, sqlserver=self)
         executor.query_execute(sql)
         executor_answer = executor.executor_answer
@@ -553,16 +516,14 @@ class MysqlProxy(SocketServer.BaseRequestHandler):
             resp = SQLAnswer(
                 resp_type=RESPONSE_TYPE.TABLE,
                 state_track=executor_answer.state_track,
-                columns=self.to_mysql_columns(executor_answer.data.columns),
-                data=executor_answer.data,
+                result_set=executor_answer.data,
                 status=executor.server_status,
-                affected_rows=executor_answer.affected_rows
+                affected_rows=executor_answer.affected_rows,
+                mysql_types=executor_answer.data.mysql_types
             )
 
         # Increment the counter and include metadata in attributes
-        metadata = ctx.metadata(query=sql)
-        query_request_counter = get_query_request_counter()
-        query_request_counter.add(1, metadata)
+        increment_otel_query_request_counter(ctx.get_metadata(query=sql))
 
         return resp
 
@@ -613,23 +574,25 @@ class MysqlProxy(SocketServer.BaseRequestHandler):
             return self.send_query_answer(resp)
 
         # TODO prepared_stmt['type'] == 'lock' is not used but it works
-        columns_def = self.to_mysql_columns(executor_answer.data.columns)
-        packages = [self.packet(ColumnCountPacket, count=len(columns_def))]
+        result_set = executor_answer.data
+        data_frame, columns_dict = dump_result_set_to_mysql(result_set)
+        data = data_frame.to_dict('split')['data']
 
-        packages.extend(self._get_column_defenition_packets(columns_def))
+        packages = [self.packet(ColumnCountPacket, count=len(columns_dict))]
+        packages.extend(self._get_column_defenition_packets(columns_dict))
 
         if self.client_capabilities.DEPRECATE_EOF is False:
             packages.append(self.packet(EofPacket, status=0x0062))
 
         # send all
-        for row in executor_answer.data.to_lists():
+        for row in data:
             packages.append(
-                self.packet(BinaryResultsetRowPacket, data=row, columns=columns_def)
+                self.packet(BinaryResultsetRowPacket, data=row, columns=columns_dict)
             )
 
         server_status = executor.server_status or 0x0002
         packages.append(self.last_packet(status=server_status))
-        prepared_stmt["fetched"] += len(executor_answer.data)
+        prepared_stmt["fetched"] += len(data)
 
         return self.send_package_group(packages)
 
@@ -741,7 +704,7 @@ class MysqlProxy(SocketServer.BaseRequestHandler):
             try:
                 if p.type.value == COMMANDS.COM_QUERY:
                     sql = self.decode_utf(p.sql.value)
-                    sql = SqlStatementParser.clear_sql(sql)
+                    sql = clear_sql(sql)
                     logger.debug(f'Incoming query: {sql}')
                     profiler.set_meta(
                         query=sql, api="mysql", environment=config.get("environment")