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

mindsdb/api/executor/command_executor.py

@@ -34,6 +34,7 @@ from mindsdb_sql_parser.ast import (
     Use,
     Tuple,
     Function,
+    Variable,
 )
 
 # typed models
@@ -49,6 +50,7 @@ from mindsdb_sql_parser.ast.mindsdb import (
     CreateSkill,
     CreateTrigger,
     CreateView,
+    CreateKnowledgeBaseIndex,
     DropAgent,
     DropChatBot,
     DropDatasource,
@@ -101,6 +103,7 @@ from mindsdb.interfaces.model.functions import (
 )
 from mindsdb.interfaces.query_context.context_controller import query_context_controller
 from mindsdb.interfaces.triggers.triggers_controller import TriggersController
+from mindsdb.interfaces.variables.variables_controller import variables_controller
 from mindsdb.utilities.context import context as ctx
 from mindsdb.utilities.functions import mark_process, resolve_model_identifier, get_handler_install_message
 from mindsdb.utilities.exception import EntityExistsError, EntityNotExistsError
@@ -336,7 +339,7 @@ class ExecuteCommands:
                 df2 = query_df(df, new_statement)
 
                 return ExecuteAnswer(
-                    data=ResultSet().from_df(df2, table_name="session_variables")
+                    data=ResultSet.from_df(df2, table_name="session_variables")
                 )
             elif sql_category == "search_path":
                 return ExecuteAnswer(
@@ -511,28 +514,30 @@ class ExecuteCommands:
             return ExecuteAnswer()
         elif statement_type is Set:
             category = (statement.category or "").lower()
-            if category == "" and isinstance(statement.name, Identifier):
-                param = statement.name.parts[0].lower()
-
-                value = None
-                if isinstance(statement.value, Constant):
-                    value = statement.value.value
-
-                if param == "profiling":
-                    self.session.profiling = value in (1, True)
-                    if self.session.profiling is True:
-                        profiler.enable()
-                    else:
-                        profiler.disable()
-                elif param == "predictor_cache":
-                    self.session.predictor_cache = value in (1, True)
-                elif param == "context":
-                    if value in (0, False, None):
-                        # drop context
-                        query_context_controller.drop_query_context(None)
-                elif param == "show_secrets":
-                    self.session.show_secrets = value in (1, True)
-
+            if category == "":
+                if isinstance(statement.name, Identifier):
+                    param = statement.name.parts[0].lower()
+
+                    value = None
+                    if isinstance(statement.value, Constant):
+                        value = statement.value.value
+
+                    if param == "profiling":
+                        self.session.profiling = value in (1, True)
+                        if self.session.profiling is True:
+                            profiler.enable()
+                        else:
+                            profiler.disable()
+                    elif param == "predictor_cache":
+                        self.session.predictor_cache = value in (1, True)
+                    elif param == "context":
+                        if value in (0, False, None):
+                            # drop context
+                            query_context_controller.drop_query_context(None)
+                    elif param == "show_secrets":
+                        self.session.show_secrets = value in (1, True)
+                elif isinstance(statement.name, Variable):
+                    variables_controller.set_variable(statement.name.value, statement.value)
                 return ExecuteAnswer()
             elif category == "autocommit":
                 return ExecuteAnswer()
@@ -649,6 +654,8 @@ class ExecuteCommands:
         elif statement_type is Evaluate:
             statement.data = parse_sql(statement.query_str)
             return self.answer_evaluate_metric(statement, database_name)
+        elif statement_type is CreateKnowledgeBaseIndex:
+            return self.answer_create_kb_index(statement, database_name)
         else:
             logger.warning(f"Unknown SQL statement: {sql}")
             raise NotSupportedYet(f"Unknown SQL statement: {sql}")
@@ -936,8 +943,18 @@ class ExecuteCommands:
         )
 
         return ExecuteAnswer(
-            data=ResultSet().from_df(df, table_name="")
+            data=ResultSet.from_df(df, table_name="")
+        )
+
+    def answer_create_kb_index(self, statement, database_name):
+        table_name = statement.name.parts[-1]
+        project_name = (
+            statement.name.parts[0]
+            if len(statement.name.parts) > 1
+            else database_name
         )
+        self.session.kb_controller.create_index(table_name=table_name, project_name=project_name)
+        return ExecuteAnswer()
 
     def _get_model_info(self, identifier, except_absent=True, database_name=None):
         if len(identifier.parts) == 1:
@@ -968,6 +985,9 @@ class ExecuteCommands:
         """Checks if there is already a predictor retraining or fine-tuning
         Do not allow to run retrain if there is another model in training process in less that 1h
         """
+        if ctx.company_id is None:
+            # bypass for tests
+            return
         is_cloud = self.session.config.get("cloud", False)
         if is_cloud and ctx.user_class == 0:
             models = get_model_records(active=None)
@@ -1032,7 +1052,7 @@ class ExecuteCommands:
         df = self.session.model_controller.retrain_model(statement, ml_handler)
 
         return ExecuteAnswer(
-            data=ResultSet().from_df(df)
+            data=ResultSet.from_df(df)
         )
 
     @profiler.profile()
@@ -1062,7 +1082,7 @@ class ExecuteCommands:
         df = self.session.model_controller.finetune_model(statement, ml_handler)
 
         return ExecuteAnswer(
-            data=ResultSet().from_df(df)
+            data=ResultSet.from_df(df)
         )
 
     def _create_integration(self, name: str, engine: str, connection_args: dict):
@@ -1572,7 +1592,7 @@ class ExecuteCommands:
         try:
             df = self.session.model_controller.create_model(statement, ml_handler)
 
-            return ExecuteAnswer(data=ResultSet().from_df(df))
+            return ExecuteAnswer(data=ResultSet.from_df(df))
         except EntityExistsError:
             if getattr(statement, "if_not_exists", False) is True:
                 return ExecuteAnswer()

mindsdb/api/executor/datahub/classes/response.py

@@ -1,11 +1,14 @@
 from dataclasses import dataclass, field
-from typing import Optional, List, Dict
+from typing import List, Dict
 
 import pandas as pd
 
+from mindsdb.api.mysql.mysql_proxy.libs.constants.mysql import MYSQL_DATA_TYPE
+
 
 @dataclass
 class DataHubResponse:
     data_frame: pd.DataFrame = field(default_factory=pd.DataFrame)
     columns: List[Dict] = field(default_factory=list)
-    affected_rows: Optional[int] = None
+    affected_rows: int | None = None
+    mysql_types: list[MYSQL_DATA_TYPE] | None = None

mindsdb/api/executor/datahub/datanodes/integration_datanode.py

@@ -1,13 +1,12 @@
 import time
 import inspect
 from dataclasses import astuple
+from typing import Iterable
 
 import numpy as np
-from numpy import dtype as np_dtype
 import pandas as pd
-from pandas.api import types as pd_types
 from sqlalchemy.types import (
-    Integer, Float, Text
+    Integer, Float
 )
 
 from mindsdb_sql_parser.ast.base import ASTNode
@@ -121,9 +120,7 @@ class IntegrationDataNode(DataNode):
             tables=[name],
             if_exists=if_exists
         )
-        result = self._query(drop_ast)
-        if result.type == RESPONSE_TYPE.ERROR:
-            raise Exception(result.error_message)
+        self.query(drop_ast)
 
     def create_table(self, table_name: Identifier, result_set: ResultSet = None, columns=None,
                      is_replace=False, is_create=False, **kwargs) -> DataHubResponse:
@@ -138,27 +135,11 @@ class IntegrationDataNode(DataNode):
 
             df = result_set.get_raw_df()
 
-            for idx, col in enumerate(result_set.columns):
-                dtype = col.type
-                # assume this is pandas type
-                column_type = Text
-                if isinstance(dtype, np_dtype):
-                    if pd_types.is_object_dtype(dtype):
-                        # try to infer
-                        dtype = df[idx].infer_objects().dtype
-
-                    if pd_types.is_integer_dtype(dtype):
-                        column_type = Integer
-                    elif pd_types.is_numeric_dtype(dtype):
-                        column_type = Float
-
-                columns.append(
-                    TableColumn(
-                        name=col.alias,
-                        type=column_type
-                    )
-                )
-                table_columns_meta[col.alias] = column_type
+            columns: list[TableColumn] = result_set.get_ast_columns()
+            table_columns_meta = {
+                column.name: column.type
+                for column in columns
+            }
 
         if is_replace:
             # drop
@@ -166,9 +147,7 @@ class IntegrationDataNode(DataNode):
                 tables=[table_name],
                 if_exists=True
             )
-            result = self._query(drop_ast)
-            if result.type == RESPONSE_TYPE.ERROR:
-                raise Exception(result.error_message)
+            self.query(drop_ast)
             is_create = True
 
         if is_create:
@@ -177,9 +156,7 @@ class IntegrationDataNode(DataNode):
                 columns=columns,
                 is_replace=is_replace
             )
-            result = self._query(create_table_ast)
-            if result.type == RESPONSE_TYPE.ERROR:
-                raise Exception(result.error_message)
+            self.query(create_table_ast)
 
         if result_set is None:
             # it is just a 'create table'
@@ -224,56 +201,44 @@ class IntegrationDataNode(DataNode):
         )
 
         try:
-            result = self._query(insert_ast)
+            result: DataHubResponse = self.query(insert_ast)
         except Exception as e:
             msg = f'[{self.ds_type}/{self.integration_name}]: {str(e)}'
             raise DBHandlerException(msg) from e
 
-        if result.type == RESPONSE_TYPE.ERROR:
-            raise Exception(result.error_message)
-
         return DataHubResponse(affected_rows=result.affected_rows)
 
-    def _query(self, query) -> HandlerResponse:
-        time_before_query = time.perf_counter()
-        result = self.integration_handler.query(query)
-        elapsed_seconds = time.perf_counter() - time_before_query
-        query_time_with_labels = metrics.INTEGRATION_HANDLER_QUERY_TIME.labels(
-            get_class_name(self.integration_handler), result.type)
-        query_time_with_labels.observe(elapsed_seconds)
-
-        num_rows = 0
-        if result.data_frame is not None:
-            num_rows = len(result.data_frame.index)
-        response_size_with_labels = metrics.INTEGRATION_HANDLER_RESPONSE_SIZE.labels(
-            get_class_name(self.integration_handler), result.type)
-        response_size_with_labels.observe(num_rows)
-        return result
-
-    def _native_query(self, native_query) -> HandlerResponse:
-        time_before_query = time.perf_counter()
-        result = self.integration_handler.native_query(native_query)
-        elapsed_seconds = time.perf_counter() - time_before_query
-        query_time_with_labels = metrics.INTEGRATION_HANDLER_QUERY_TIME.labels(
-            get_class_name(self.integration_handler), result.type)
-        query_time_with_labels.observe(elapsed_seconds)
-
-        num_rows = 0
-        if result.data_frame is not None:
-            num_rows = len(result.data_frame.index)
-        response_size_with_labels = metrics.INTEGRATION_HANDLER_RESPONSE_SIZE.labels(
-            get_class_name(self.integration_handler), result.type)
-        response_size_with_labels.observe(num_rows)
-        return result
+    def has_support_stream(self) -> bool:
+        # checks if data handler has query_stream method
+        return hasattr(self.integration_handler, 'query_stream') and callable(self.integration_handler.query_stream)
+
+    @profiler.profile()
+    def query_stream(self, query: ASTNode, fetch_size: int = None) -> Iterable:
+        # returns generator of results from handler (split by chunks)
+        return self.integration_handler.query_stream(query, fetch_size=fetch_size)
 
     @profiler.profile()
     def query(self, query: ASTNode | None = None, native_query: str | None = None, session=None) -> DataHubResponse:
         try:
+            time_before_query = time.perf_counter()
             if query is not None:
-                result: HandlerResponse = self._query(query)
+                result: HandlerResponse = self.integration_handler.query(query)
             else:
                 # try to fetch native query
-                result: HandlerResponse = self._native_query(native_query)
+                result: HandlerResponse = self.integration_handler.native_query(native_query)
+
+            # metrics
+            elapsed_seconds = time.perf_counter() - time_before_query
+            query_time_with_labels = metrics.INTEGRATION_HANDLER_QUERY_TIME.labels(
+                get_class_name(self.integration_handler), result.type)
+            query_time_with_labels.observe(elapsed_seconds)
+
+            num_rows = 0
+            if result.data_frame is not None:
+                num_rows = len(result.data_frame.index)
+            response_size_with_labels = metrics.INTEGRATION_HANDLER_RESPONSE_SIZE.labels(
+                get_class_name(self.integration_handler), result.type)
+            response_size_with_labels.observe(num_rows)
         except Exception as e:
             msg = str(e).strip()
             if msg == '':
@@ -294,7 +259,8 @@ class IntegrationDataNode(DataNode):
 
         try:
             # replace python's Nan, np.NaN, np.nan and pd.NA to None
-            df.replace([np.NaN, pd.NA], None, inplace=True)
+            # TODO keep all NAN to the end of processing, bacause replacing also changes dtypes
+            df.replace([np.NaN, pd.NA, pd.NaT], None, inplace=True)
         except Exception as e:
             logger.error(f"Issue with clearing DF from NaN values: {e}")
         # endregion
@@ -310,5 +276,6 @@ class IntegrationDataNode(DataNode):
         return DataHubResponse(
             data_frame=df,
             columns=columns_info,
-            affected_rows=result.affected_rows
+            affected_rows=result.affected_rows,
+            mysql_types=result.mysql_types
         )

mindsdb/api/executor/planner/query_planner.py

@@ -848,11 +848,12 @@ class QueryPlanner:
         """
 
         # handle fetchdataframe partitioning
+        steps_in = plan.steps
         steps_out = []
 
         step = None
         partition_step = None
-        for step in plan.steps:
+        for step in steps_in:
             if isinstance(step, FetchDataframeStep) and step.params is not None:
                 batch_size = step.params.get('batch_size')
                 if batch_size is not None:
@@ -905,6 +906,14 @@ class QueryPlanner:
 
         if plan.is_resumable and isinstance(step, InsertToTable):
             plan.is_async = True
+        else:
+            # special case: register insert from select (it is the same as mark resumable)
+            if (
+                len(steps_in) == 2
+                and isinstance(steps_in[0], FetchDataframeStep)
+                and isinstance(steps_in[1], InsertToTable)
+            ):
+                plan.is_resumable = True
 
         plan.steps = steps_out
         return plan