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

mindsdb/interfaces/query_context/context_controller.py

@@ -1,11 +1,17 @@
 from typing import List
+import pickle
+import datetime as dt
 
+from sqlalchemy.orm.attributes import flag_modified
 import pandas as pd
 
+from mindsdb_sql_parser import Select, Star, OrderBy
+
 from mindsdb_sql_parser.ast import (
     Identifier, BinaryOperation, Last, Constant, ASTNode
 )
 from mindsdb.integrations.utilities.query_traversal import query_traversal
+from mindsdb.utilities.cache import get_cache
 
 from mindsdb.interfaces.storage import db
 from mindsdb.utilities.context import context as ctx
@@ -13,6 +19,147 @@ from mindsdb.utilities.context import context as ctx
 from .last_query import LastQuery
 
 
+class RunningQuery:
+    """
+      Query in progres
+    """
+
+    def __init__(self, record: db.Queries):
+        self.record = record
+        self.sql = record.sql
+
+    def get_partition_query(self, step_num: int, query: Select) -> Select:
+        """
+           Generate query for fetching the next partition
+           It wraps query to
+              select * from ({query})
+              where {track_column} > {previous_value}
+              order by track_column
+              limit size {batch_size}
+           And fill track_column, previous_value, batch_size
+        """
+
+        track_column = self.record.parameters['track_column']
+
+        query = Select(
+            targets=[Star()],
+            from_table=query,
+            order_by=[OrderBy(Identifier(track_column))],
+            limit=Constant(self.batch_size)
+        )
+
+        track_value = self.record.context.get('track_value')
+        # is it different step?
+        cur_step_num = self.record.context.get('step_num')
+        if cur_step_num is not None and cur_step_num != step_num:
+            # reset track_value
+            track_value = None
+            self.record.context['track_value'] = None
+            self.record.context['step_num'] = step_num
+            flag_modified(self.record, 'context')
+            db.session.commit()
+
+        if track_value is not None:
+            query.where = BinaryOperation(
+                op='>',
+                args=[Identifier(track_column), Constant(track_value)],
+            )
+
+        return query
+
+    def set_params(self, params: dict):
+        """
+            Store parameters of the step which is about to be split into partitions
+        """
+
+        if 'track_column' not in params:
+            raise ValueError('Track column is not defined')
+        if 'batch_size' not in params:
+            params['batch_size'] = 1000
+
+        self.record.parameters = params
+        self.batch_size = self.record.parameters['batch_size']
+        db.session.commit()
+
+    def get_max_track_value(self, df: pd.DataFrame) -> pd.DataFrame:
+        """
+            return max value to use in `set_progress`.
+            this function is called before execution substeps,
+             `set_progress` function - after
+        """
+
+        track_column = self.record.parameters['track_column']
+        return df[track_column].max()
+
+    def set_progress(self, df: pd.DataFrame, max_track_value: int):
+        """
+           Store progres of the query, it is called after processing of batch
+        """
+
+        if len(df) == 0:
+            return
+
+        self.record.processed_rows = self.record.processed_rows + len(df)
+
+        cur_value = self.record.context.get('track_value')
+        new_value = max_track_value
+        if new_value is not None:
+            if cur_value is None or new_value > cur_value:
+                self.record.context['track_value'] = new_value
+                flag_modified(self.record, 'context')
+
+        db.session.commit()
+
+    def on_error(self, error: Exception, step_num: int, steps_data: dict):
+        """
+            Saves error of the query in database
+            Also saves step data and current step num to be able to resume query
+        """
+        self.record.error = str(error)
+        self.record.context['step_num'] = step_num
+        flag_modified(self.record, 'context')
+
+        # save steps_data
+        cache = get_cache('steps_data')
+        data = pickle.dumps(steps_data, protocol=5)
+        cache.set(str(self.record.id), data)
+
+        db.session.commit()
+
+    def clear_error(self):
+        """
+            Reset error of the query in database
+        """
+
+        if self.record.error is not None:
+            self.record.error = None
+            db.session.commit()
+
+    def get_state(self) -> dict:
+        """
+            Returns stored state for resuming the query
+        """
+        cache = get_cache('steps_data')
+        key = self.record.id
+        data = cache.get(key)
+        cache.delete(key)
+
+        steps_data = pickle.loads(data)
+
+        return {
+            'step_num': self.record.context.get('step_num'),
+            'steps_data': steps_data,
+        }
+
+    def finish(self):
+        """
+            Mark query as finished
+        """
+
+        self.record.finished_at = dt.datetime.now()
+        db.session.commit()
+
+
 class QueryContextController:
     IGNORE_CONTEXT = '<IGNORE>'
 
@@ -156,10 +303,12 @@ class QueryContextController:
         last_values = {}
         for query, info in l_query.get_init_queries():
 
-            data, columns_info = dn.query(
+            response = dn.query(
                 query=query,
                 session=session
             )
+            data = response.data_frame
+            columns_info = response.columns
 
             if len(data) == 0:
                 value = None
@@ -285,5 +434,79 @@ class QueryContextController:
         rec.values = values
         db.session.commit()
 
+    def get_query(self, query_id: int) -> RunningQuery:
+        """
+           Get running query by id
+        """
+
+        rec = db.Queries.query.filter(
+            db.Queries.id == query_id,
+            db.Queries.company_id == ctx.company_id
+        ).first()
+
+        if rec is None:
+            raise RuntimeError(f'Query not found: {query_id}')
+        return RunningQuery(rec)
+
+    def create_query(self, query: ASTNode) -> RunningQuery:
+        """
+           Create a new running query from AST query
+        """
+
+        # remove old queries
+        remove_query = db.session.query(db.Queries).filter(
+            db.Queries.company_id == ctx.company_id,
+            db.Queries.finished_at < (dt.datetime.now() - dt.timedelta(days=1))
+        )
+        for rec in remove_query.all():
+            db.session.delete(rec)
+
+        rec = db.Queries(
+            sql=str(query),
+            company_id=ctx.company_id,
+        )
+
+        db.session.add(rec)
+        db.session.commit()
+        return RunningQuery(rec)
+
+    def list_queries(self) -> List[dict]:
+        """
+           Get list of all running queries with metadata
+        """
+
+        query = db.session.query(db.Queries).filter(
+            db.Queries.company_id == ctx.company_id
+        )
+        return [
+            {
+                'id': record.id,
+                'sql': record.sql,
+                'started_at': record.started_at,
+                'finished_at': record.finished_at,
+                'parameters': record.parameters,
+                'context': record.context,
+                'processed_rows': record.processed_rows,
+                'error': record.error,
+                'updated_at': record.updated_at,
+            }
+            for record in query
+        ]
+
+    def cancel_query(self, query_id: int):
+        """
+           Cancels running query by id
+        """
+        rec = db.Queries.query.filter(
+            db.Queries.id == query_id,
+            db.Queries.company_id == ctx.company_id
+        ).first()
+        if rec is None:
+            raise RuntimeError(f'Query not found: {query_id}')
+
+        # the query in progress will fail when it tries to update status
+        db.session.delete(rec)
+        db.session.commit()
+
 
 query_context_controller = QueryContextController()

mindsdb/interfaces/storage/db.py

@@ -523,6 +523,7 @@ class KnowledgeBase(Base):
     embedding_model = relationship(
         "Predictor", foreign_keys=[embedding_model_id], doc="embedding model"
     )
+    query_id = Column(Integer, nullable=True)
 
     created_at = Column(DateTime, default=datetime.datetime.now)
     updated_at = Column(
@@ -564,6 +565,28 @@ class QueryContext(Base):
     created_at: datetime.datetime = Column(DateTime, default=datetime.datetime.now)
 
 
+class Queries(Base):
+    __tablename__ = "queries"
+    id: int = Column(Integer, primary_key=True)
+    company_id: int = Column(Integer, nullable=True)
+
+    sql: str = Column(String, nullable=False)
+    # step_data: JSON = Column(JSON, nullable=True)
+
+    started_at: datetime.datetime = Column(DateTime, default=datetime.datetime.now)
+    finished_at: datetime.datetime = Column(DateTime)
+
+    parameters = Column(JSON, default={})
+    context = Column(JSON, default={})
+    processed_rows = Column(Integer, default=0)
+    error: str = Column(String, nullable=True)
+
+    updated_at: datetime.datetime = Column(
+        DateTime, default=datetime.datetime.now, onupdate=datetime.datetime.now
+    )
+    created_at: datetime.datetime = Column(DateTime, default=datetime.datetime.now)
+
+
 class LLMLog(Base):
     __tablename__ = "llm_log"
     id: int = Column(Integer, primary_key=True)

mindsdb/migrations/versions/2025-03-21_fda503400e43_queries.py

@@ -0,0 +1,45 @@
+"""queries
+
+Revision ID: fda503400e43
+Revises: 11347c213b36
+Create Date: 2025-03-21 18:50:20.795930
+
+"""
+from alembic import op
+import sqlalchemy as sa
+import mindsdb.interfaces.storage.db  # noqa
+
+
+# revision identifiers, used by Alembic.
+revision = 'fda503400e43'
+down_revision = '11347c213b36'
+branch_labels = None
+depends_on = None
+
+
+def upgrade():
+    op.create_table(
+        'queries',
+        sa.Column('id', sa.Integer(), nullable=False),
+        sa.Column('company_id', sa.Integer(), nullable=True),
+        sa.Column('sql', sa.String(), nullable=False),
+        sa.Column('started_at', sa.DateTime(), nullable=True),
+        sa.Column('finished_at', sa.DateTime(), nullable=True),
+        sa.Column('parameters', sa.JSON(), nullable=True),
+        sa.Column('context', sa.JSON(), nullable=True),
+        sa.Column('processed_rows', sa.Integer(), nullable=True),
+        sa.Column('error', sa.String(), nullable=True),
+        sa.Column('updated_at', sa.DateTime(), nullable=True),
+        sa.Column('created_at', sa.DateTime(), nullable=True),
+        sa.PrimaryKeyConstraint('id')
+    )
+
+    with op.batch_alter_table('knowledge_base', schema=None) as batch_op:
+        batch_op.add_column(sa.Column('query_id', sa.INTEGER(), nullable=True))
+
+
+def downgrade():
+    with op.batch_alter_table('knowledge_base', schema=None) as batch_op:
+        batch_op.drop_column('query_id')
+
+    op.drop_table('queries')

mindsdb/utilities/context_executor.py

@@ -43,7 +43,7 @@ def execute_in_threads(func, tasks, thread_count=3, queue_size_k=1.5):
         for i in range(queue_size):
             try:
                 args = next(tasks)
-                futures.append(executor.submit(func, *args))
+                futures.append(executor.submit(func, args))
             except StopIteration:
                 break
 

mindsdb/utilities/partitioning.py

@@ -6,6 +6,35 @@ from mindsdb.utilities.config import Config
 from mindsdb.utilities.context_executor import execute_in_threads
 
 
+def get_max_thread_count() -> int:
+    """
+        Calculate the maximum number of threads allowed for the system.
+    """
+    # workers count
+    is_cloud = Config().is_cloud
+    if is_cloud:
+        max_threads = int(os.getenv('MINDSDB_MAX_PARTITIONING_THREADS', 10))
+    else:
+        max_threads = os.cpu_count() - 3
+
+    if max_threads < 1:
+        max_threads = 1
+
+    return max_threads
+
+
+def split_data_frame(df: pd.DataFrame, partition_size: int) -> Iterable[pd.DataFrame]:
+    """
+    Split data frame into chunks with partition_size and yield them out
+    """
+    chunk = 0
+    while chunk * partition_size < len(df):
+        # create results with partition
+        df1 = df.iloc[chunk * partition_size: (chunk + 1) * partition_size]
+        chunk += 1
+        yield df1
+
+
 def process_dataframe_in_partitions(df: pd.DataFrame, callback: Callable, partition_size: int) -> Iterable:
     """
     Splits dataframe into partitions and apply callback on each partition
@@ -17,35 +46,21 @@ def process_dataframe_in_partitions(df: pd.DataFrame, callback: Callable, partit
     """
 
     # tasks
-    def split_data_f(df):
-        chunk = 0
-        while chunk * partition_size < len(df):
-            # create results with partition
-            df1 = df.iloc[chunk * partition_size: (chunk + 1) * partition_size]
-            chunk += 1
-            yield [df1]
 
-    tasks = split_data_f(df)
+    tasks = split_data_frame(df, partition_size)
 
-    # workers count
-    is_cloud = Config().is_cloud
-    if is_cloud:
-        max_threads = int(os.getenv('MINDSDB_MAX_PARTITIONING_THREADS', 10))
-    else:
-        max_threads = os.cpu_count() - 2
+    max_threads = get_max_thread_count()
 
-    # don't exceed chunk_count
     chunk_count = int(len(df) / partition_size)
-    max_threads = min(max_threads, chunk_count)
-
-    if max_threads < 1:
-        max_threads = 1
+    # don't exceed chunk_count
+    if chunk_count > 0:
+        max_threads = min(max_threads, chunk_count)
 
     if max_threads == 1:
         # don't spawn threads
 
         for task in tasks:
-            yield callback(*task)
+            yield callback(task)
 
     else:
         for result in execute_in_threads(callback, tasks, thread_count=max_threads):