MindsDB 25.4.2.0__py3-none-any.whl → 25.4.2.1__py3-none-any.whl

mindsdb/__about__.py

@@ -1,6 +1,6 @@
 __title__ = 'MindsDB'
 __package_name__ = 'mindsdb'
-__version__ = '25.4.2.0'
+__version__ = '25.4.2.1'
 __description__ = "MindsDB's AI SQL Server enables developers to build AI tools that need access to real-time data to perform their tasks"
 __email__ = "jorge@mindsdb.com"
 __author__ = 'MindsDB Inc'

mindsdb/api/executor/command_executor.py

@@ -34,6 +34,7 @@ from mindsdb_sql_parser.ast import (
     Update,
     Use,
     Tuple,
+    Function,
 )
 
 # typed models
@@ -599,6 +600,9 @@ class ExecuteCommands:
         ):
             return ExecuteAnswer()
         elif statement_type is Select:
+            ret = self.exec_service_function(statement, database_name)
+            if ret is not None:
+                return ret
             query = SQLQuery(statement, session=self.session, database=database_name)
             return self.answer_select(query)
         elif statement_type is Union:
@@ -648,6 +652,31 @@ class ExecuteCommands:
             logger.warning(f"Unknown SQL statement: {sql}")
             raise NotSupportedYet(f"Unknown SQL statement: {sql}")
 
+    def exec_service_function(self, statement: Select, database_name: str) -> Optional[ExecuteAnswer]:
+        """
+        If input query is a single line select without FROM
+          and has function in targets that matches with one of the mindsdb service functions:
+          - execute this function and return response
+        Otherwise, return None to allow to continue execution query outside
+        """
+
+        if statement.from_table is not None or len(statement.targets) != 1:
+            return
+
+        target = statement.targets[0]
+        if not isinstance(target, Function):
+            return
+
+        command = target.op.lower()
+        args = [arg.value for arg in target.args if isinstance(arg, Constant)]
+        if command == 'query_resume':
+            ret = SQLQuery(None, session=self.session, database=database_name, query_id=args[0])
+            return self.answer_select(ret)
+
+        elif command == 'query_cancel':
+            query_context_controller.cancel_query(*args)
+            return ExecuteAnswer()
+
     def answer_create_trigger(self, statement, database_name):
         triggers_controller = TriggersController()
 

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

@@ -17,7 +17,7 @@ from .system_tables import (
     PluginsTable, EnginesTable, KeyColumnUsageTable, StatisticsTable,
     CharacterSetsTable, CollationsTable)
 from .mindsdb_tables import (
-    ModelsTable, DatabasesTable, MLEnginesTable, HandlersTable, JobsTable,
+    ModelsTable, DatabasesTable, MLEnginesTable, HandlersTable, JobsTable, QueriesTable,
     ChatbotsTable, KBTable, SkillsTable, AgentsTable, ViewsTable, TriggersTable)
 
 
@@ -32,7 +32,8 @@ class InformationSchemaDataNode(DataNode):
         PluginsTable, EnginesTable, KeyColumnUsageTable, StatisticsTable,
         CharacterSetsTable, CollationsTable,
         ModelsTable, DatabasesTable, MLEnginesTable, HandlersTable, JobsTable,
-        ChatbotsTable, KBTable, SkillsTable, AgentsTable, ViewsTable, TriggersTable
+        ChatbotsTable, KBTable, SkillsTable, AgentsTable, ViewsTable, TriggersTable,
+        QueriesTable
     ]
 
     def __init__(self, session):

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

@@ -9,6 +9,7 @@ from mindsdb.interfaces.jobs.jobs_controller import JobsController
 from mindsdb.interfaces.skills.skills_controller import SkillsController
 from mindsdb.interfaces.database.views import ViewController
 from mindsdb.interfaces.database.projects import ProjectController
+from mindsdb.interfaces.query_context.context_controller import query_context_controller
 
 from mindsdb.api.executor.datahub.datanodes.system_tables import Table
 
@@ -326,7 +327,8 @@ class ChatbotsTable(MdbTable):
 
 class KBTable(MdbTable):
     name = 'KNOWLEDGE_BASES'
-    columns = ["NAME", "PROJECT", "MODEL", "STORAGE", "PARAMS"]
+    columns = ["NAME", "PROJECT", "MODEL", "STORAGE", "PARAMS",
+               "INSERT_STARTED_AT", "INSERT_FINISHED_AT", "PROCESSED_ROWS", "ERROR", "QUERY_ID"]
 
     @classmethod
     def get_data(cls, query: ASTNode = None, inf_schema=None, **kwargs):
@@ -336,17 +338,36 @@ class KBTable(MdbTable):
         controller = KnowledgeBaseController(inf_schema.session)
         kb_list = controller.list(project_name)
 
+        # shouldn't be a lot of queries, we can fetch them all
+        queries_data = {
+            item['id']: item
+            for item in query_context_controller.list_queries()
+        }
+
         data = []
 
         for kb in kb_list:
             vector_database_name = kb['vector_database'] or ''
 
+            query_item = {}
+            query_id = kb['query_id']
+            if query_id is not None:
+                if query_id in queries_data:
+                    query_item = queries_data.get(query_id)
+                else:
+                    query_id = None
+
             data.append((
                 kb['name'],
                 kb['project_name'],
                 kb['embedding_model'],
                 vector_database_name + '.' + kb['vector_database_table'],
                 to_json(kb['params']),
+                query_item.get('started_at'),
+                query_item.get('finished_at'),
+                query_item.get('processed_rows'),
+                query_item.get('error'),
+                query_id,
             ))
 
         return pd.DataFrame(data, columns=cls.columns)
@@ -426,3 +447,24 @@ class ViewsTable(MdbTable):
         data = [[row[k] for k in columns_lower] for row in data]
 
         return pd.DataFrame(data, columns=cls.columns)
+
+
+class QueriesTable(MdbTable):
+    name = 'QUERIES'
+    columns = ["ID", "STARTED_AT", "FINISHED_AT", "PROCESSED_ROWS", "ERROR", "SQL", "PARAMETERS", "CONTEXT", "UPDATED_AT"]
+
+    @classmethod
+    def get_data(cls, **kwargs):
+        """
+        Returns all queries in progres or recently completed
+        Only queries marked as is_resumable by planner are stored in this table
+        :param kwargs:
+        :return:
+        """
+
+        data = query_context_controller.list_queries()
+        columns_lower = [col.lower() for col in cls.columns]
+
+        data = [[row[k] for k in columns_lower] for row in data]
+
+        return pd.DataFrame(data, columns=cls.columns)

mindsdb/api/executor/planner/plan_join.py

@@ -423,7 +423,7 @@ class PlanJoinTablesQuery:
             else:
                 query2.where = cond
 
-        step = self.planner.get_integration_select_step(query2)
+        step = self.planner.get_integration_select_step(query2, params=query_in.using)
         self.tables_fetch_step[item.index] = step
 
         self.add_plan_step(step)

mindsdb/api/executor/planner/query_plan.py

@@ -2,6 +2,7 @@
 class QueryPlan:
     def __init__(self, steps=None, **kwargs):
         self.steps = []
+        self.is_resumable = False
 
         if steps:
             for step in steps:

mindsdb/api/executor/planner/query_planner.py

@@ -12,14 +12,13 @@ from mindsdb.api.executor.planner.exceptions import PlanningException
 from mindsdb.api.executor.planner import utils
 from mindsdb.api.executor.planner.query_plan import QueryPlan
 from mindsdb.api.executor.planner.steps import (
-    FetchDataframeStep, ProjectStep, ApplyPredictorStep,
+    PlanStep, FetchDataframeStep, ProjectStep, ApplyPredictorStep,
     ApplyPredictorRowStep, UnionStep, GetPredictorColumns, SaveToTable,
-    InsertToTable, UpdateToTable, SubSelectStep, QueryStep,
-    DeleteStep, DataStep, CreateTableStep
+    InsertToTable, UpdateToTable, SubSelectStep, QueryStep, JoinStep,
+    DeleteStep, DataStep, CreateTableStep, FetchDataframeStepPartition
 )
 from mindsdb.api.executor.planner.utils import (
     disambiguate_predictor_column_identifier,
-    get_deepest_select,
     recursively_extract_column_values,
     query_traversal, filters_to_bin_op
 )
@@ -166,7 +165,11 @@ class QueryPlanner:
 
         query_traversal(query, _prepare_integration_select)
 
-    def get_integration_select_step(self, select):
+    def get_integration_select_step(self, select: Select, params: dict = None) -> PlanStep:
+        """
+        Generate planner step to execute query over integration or over results of previous step (if it is CTE)
+        """
+
         if isinstance(select.from_table, NativeQuery):
             integration_name = select.from_table.integration.parts[-1]
         else:
@@ -188,12 +191,22 @@ class QueryPlanner:
         if fetch_df_select.using is not None:
             fetch_df_select.using = None
 
-        return FetchDataframeStep(integration=integration_name, query=fetch_df_select)
+        if params:
+            fetch_params = params.copy()
+            # remove partition parameters
+            for key in ('batch_size', 'track_column'):
+                if key in params:
+                    del params[key]
+            if 'track_column' in fetch_params and isinstance(fetch_params['track_column'], Identifier):
+                fetch_params['track_column'] = fetch_params['track_column'].parts[-1]
+        else:
+            fetch_params = None
+        return FetchDataframeStep(integration=integration_name, query=fetch_df_select, params=fetch_params)
 
     def plan_integration_select(self, select):
         """Plan for a select query that can be fully executed in an integration"""
 
-        return self.plan.add_step(self.get_integration_select_step(select))
+        return self.plan.add_step(self.get_integration_select_step(select, params=select.using))
 
     def resolve_database_table(self, node: Identifier):
         # resolves integration name and table name
@@ -414,12 +427,6 @@ class QueryPlanner:
 
         return self.plan_mdb_nested_select(select)
 
-    def plan_integration_nested_select(self, select, integration_name):
-        fetch_df_select = copy.deepcopy(select)
-        deepest_select = get_deepest_select(fetch_df_select)
-        self.prepare_integration_select(integration_name, deepest_select)
-        return self.plan.add_step(FetchDataframeStep(integration=integration_name, query=fetch_df_select))
-
     def plan_mdb_nested_select(self, select):
         # plan nested select
 
@@ -818,7 +825,72 @@ class QueryPlanner:
         else:
             raise PlanningException(f'Unsupported query type {type(query)}')
 
-        return self.plan
+        plan = self.handle_partitioning(self.plan)
+
+        return plan
+
+    def handle_partitioning(self, plan: QueryPlan) -> QueryPlan:
+        """
+        If plan has fetching in partitions:
+          try to rebuild plan to send fetched chunk of data through the following steps, if it is possible
+        """
+
+        # handle fetchdataframe partitioning
+        steps_out = []
+
+        partition_step = None
+        for step in plan.steps:
+            if isinstance(step, FetchDataframeStep) and step.params is not None:
+                batch_size = step.params.get('batch_size')
+                if batch_size is not None:
+                    # found batched fetch
+                    partition_step = FetchDataframeStepPartition(
+                        step_num=step.step_num,
+                        integration=step.integration,
+                        query=step.query,
+                        raw_query=step.raw_query,
+                        params=step.params
+                    )
+                    steps_out.append(partition_step)
+                    # mark plan
+                    plan.is_resumable = True
+                    continue
+                else:
+                    step.params = None
+
+            if partition_step is not None:
+                # check and add step into partition
+
+                can_be_partitioned = False
+                if isinstance(step, (JoinStep, ApplyPredictorStep, InsertToTable)):
+                    can_be_partitioned = True
+                elif isinstance(step, QueryStep):
+                    query = step.query
+                    if (
+                        query.group_by is None and query.order_by is None and query.distinct is False
+                        and query.limit is None and query.offset is None
+                    ):
+                        no_identifiers = [
+                            target
+                            for target in step.query.targets
+                            if not isinstance(target, (Star, Identifier))
+                        ]
+                        if len(no_identifiers) == 0:
+                            can_be_partitioned = True
+
+                if not can_be_partitioned:
+                    if len(partition_step.steps) == 0:
+                        # Nothing can be partitioned, failback to old plan
+                        plan.is_resumable = False
+                        return plan
+                    partition_step = None
+                else:
+                    partition_step.steps.append(step)
+                    continue
+
+            steps_out.append(step)
+        plan.steps = steps_out
+        return plan
 
     def prepare_steps(self, query):
         statement_planner = PreparedStatementPlanner(self)

mindsdb/api/executor/planner/steps.py

@@ -104,11 +104,19 @@ class LimitOffsetStep(PlanStep):
 
 class FetchDataframeStep(PlanStep):
     """Fetches a dataframe from external integration"""
-    def __init__(self, integration, query=None, raw_query=None, *args, **kwargs):
+    def __init__(self, integration, query=None, raw_query=None, params=None, *args, **kwargs):
         super().__init__(*args, **kwargs)
         self.integration = integration
         self.query = query
         self.raw_query = raw_query
+        self.params = params
+
+
+class FetchDataframeStepPartition(FetchDataframeStep):
+    """Fetches a dataframe from external integration in partitions"""
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self.steps = []
 
 
 class ApplyPredictorStep(PlanStep):

mindsdb/api/executor/sql_query/sql_query.py

@@ -10,9 +10,9 @@
 """
 import inspect
 from textwrap import dedent
-from typing import Dict
+from typing import Union, Dict
 
-from mindsdb_sql_parser import parse_sql
+from mindsdb_sql_parser import parse_sql, ASTNode
 from mindsdb.api.executor.planner.steps import (
     ApplyTimeseriesPredictorStep,
     ApplyPredictorRowStep,
@@ -33,6 +33,9 @@ from mindsdb.api.executor.exceptions import (
 import mindsdb.utilities.profiler as profiler
 from mindsdb.utilities.fs import create_process_mark, delete_process_mark
 from mindsdb.utilities.exception import EntityNotExistsError
+from mindsdb.interfaces.query_context.context_controller import query_context_controller
+from mindsdb.utilities.context import context as ctx
+
 
 from . import steps
 from .result_set import ResultSet, Column
@@ -43,7 +46,8 @@ class SQLQuery:
 
     step_handlers = {}
 
-    def __init__(self, sql, session, execute=True, database=None):
+    def __init__(self, sql: Union[ASTNode, str], session, execute: bool = True,
+                 database: str = None, query_id: int = None):
         self.session = session
 
         if database is not None:
@@ -63,6 +67,15 @@ class SQLQuery:
         self.parameters = []
         self.fetched_data: ResultSet = None
 
+        self.outer_query = None
+        self.run_query = None
+        self.query_id = query_id
+        if query_id is not None:
+            # resume query
+            run_query = query_context_controller.get_query(self.query_id)
+            run_query.clear_error()
+            sql = run_query.sql
+
         if isinstance(sql, str):
             self.query = parse_sql(sql)
             self.context['query_str'] = sql
@@ -217,10 +230,22 @@ class SQLQuery:
             # no need to execute
             return
 
+        try:
+            steps = list(self.planner.execute_steps())
+        except PlanningException as e:
+            raise LogicError(e)
+
+        if self.planner.plan.is_resumable:
+            # create query
+            if self.query_id is not None:
+                self.run_query = query_context_controller.get_query(self.query_id)
+            else:
+                self.run_query = query_context_controller.create_query(self.context['query_str'])
+            ctx.run_query_id = self.run_query.record.id
+
         step_result = None
         process_mark = None
         try:
-            steps = list(self.planner.execute_steps())
             steps_classes = (x.__class__ for x in steps)
             predict_steps = (ApplyPredictorRowStep, ApplyPredictorStep, ApplyTimeseriesPredictorStep)
             if any(s in predict_steps for s in steps_classes):
@@ -229,10 +254,16 @@ class SQLQuery:
                 with profiler.Context(f'step: {step.__class__.__name__}'):
                     step_result = self.execute_step(step)
                 self.steps_data[step.step_num] = step_result
-        except PlanningException as e:
-            raise LogicError(e)
         except Exception as e:
+            if self.run_query is not None:
+                # set error and place where it stopped
+                self.run_query.on_error(e, step.step_num, self.steps_data)
             raise e
+        else:
+            # mark running query as completed
+            if self.run_query is not None:
+                self.run_query.finish()
+                ctx.run_query_id = None
         finally:
             if process_mark is not None:
                 delete_process_mark('predict', process_mark)

mindsdb/api/executor/sql_query/steps/__init__.py

@@ -1,6 +1,7 @@
 from .apply_predictor_step import ApplyPredictorStepCall, ApplyPredictorRowStepCall, ApplyTimeseriesPredictorStepCall
 from .delete_step import DeleteStepCall
 from .fetch_dataframe import FetchDataframeStepCall
+from .fetch_dataframe_partition import FetchDataframePartitionCall
 from .insert_step import InsertToTableCall, SaveToTableCall, CreateTableCall
 from .join_step import JoinStepCall
 from .map_reduce_step import MapReduceStepCall