MindsDB 25.1.2.1__py3-none-any.whl → 25.1.4.0__py3-none-any.whl

mindsdb/utilities/langfuse.py

@@ -0,0 +1,264 @@
+import os
+import typing
+
+from mindsdb.utilities import log
+from langfuse import Langfuse
+from langfuse.client import StatefulSpanClient
+from langfuse.callback import CallbackHandler
+from langfuse.api.resources.commons.errors.not_found_error import NotFoundError as TraceNotFoundError
+
+logger = log.getLogger(__name__)
+
+# Define Langfuse public key.
+LANGFUSE_PUBLIC_KEY = os.getenv("LANGFUSE_PUBLIC_KEY", "langfuse_public_key")
+
+# Define Langfuse secret key.
+LANGFUSE_SECRET_KEY = os.getenv("LANGFUSE_SECRET_KEY", "langfuse_secret_key")
+
+# Define Langfuse host.
+LANGFUSE_HOST = os.getenv("LANGFUSE_HOST", "http://localhost:3000")
+
+# Define Langfuse environment.
+LANGFUSE_ENVIRONMENT = os.getenv("LANGFUSE_ENVIRONMENT", "local")
+
+# Define Langfuse release.
+LANGFUSE_RELEASE = os.getenv("LANGFUSE_RELEASE", "local")
+
+# Define Langfuse debug mode.
+LANGFUSE_DEBUG = os.getenv("LANGFUSE_DEBUG", "false").lower() == "true"
+
+# Define Langfuse timeout.
+LANGFUSE_TIMEOUT = int(os.getenv("LANGFUSE_TIMEOUT", 10))
+
+# Define Langfuse sample rate.
+LANGFUSE_SAMPLE_RATE = float(os.getenv("LANGFUSE_SAMPLE_RATE", 1.0))
+
+# Define if Langfuse is disabled.
+LANGFUSE_DISABLED = os.getenv("LANGFUSE_DISABLED", "false").lower() == "true" or LANGFUSE_ENVIRONMENT == "local"
+LANGFUSE_FORCE_RUN = os.getenv("LANGFUSE_FORCE_RUN", "false").lower() == "true"
+
+
+class LangfuseClientWrapper:
+    """
+    Langfuse client wrapper. Defines Langfuse client configuration and initializes Langfuse client.
+    """
+
+    def __init__(self,
+                 public_key: str = LANGFUSE_PUBLIC_KEY,
+                 secret_key: str = LANGFUSE_SECRET_KEY,
+                 host: str = LANGFUSE_HOST,
+                 environment: str = LANGFUSE_ENVIRONMENT,
+                 release: str = LANGFUSE_RELEASE,
+                 debug: bool = LANGFUSE_DEBUG,
+                 timeout: int = LANGFUSE_TIMEOUT,
+                 sample_rate: float = LANGFUSE_SAMPLE_RATE,
+                 disable: bool = LANGFUSE_DISABLED,
+                 force_run: bool = LANGFUSE_FORCE_RUN) -> None:
+        """
+        Initialize Langfuse client.
+
+        Args:
+            public_key (str): Langfuse public key.
+            secret_key (str): Langfuse secret key.
+            host (str): Langfuse host.
+            release (str): Langfuse release.
+            timeout (int): Langfuse timeout.
+            sample_rate (float): Langfuse sample rate.
+        """
+
+        self.metadata = None
+        self.public_key = public_key
+        self.secret_key = secret_key
+        self.host = host
+        self.environment = environment
+        self.release = release
+        self.debug = debug
+        self.timeout = timeout
+        self.sample_rate = sample_rate
+        self.disable = disable
+        self.force_run = force_run
+
+        self.client = None
+        self.trace = None
+        self.metadata = None
+        self.tags = None
+
+        # Check if Langfuse is disabled.
+        if LANGFUSE_DISABLED and not LANGFUSE_FORCE_RUN:
+            logger.info("Langfuse is disabled.")
+            return
+
+        logger.info("Langfuse enabled")
+        logger.debug(f"LANGFUSE_PUBLIC_KEY: {LANGFUSE_PUBLIC_KEY}")
+        logger.debug(f"LANGFUSE_SECRET_KEY: {'*' * len(LANGFUSE_SECRET_KEY)}")
+        logger.debug(f"LANGFUSE_HOST: {LANGFUSE_HOST}")
+        logger.debug(f"LANGFUSE_ENVIRONMENT: {LANGFUSE_ENVIRONMENT}")
+        logger.debug(f"LANGFUSE_RELEASE: {LANGFUSE_RELEASE}")
+        logger.debug(f"LANGFUSE_DEBUG: {LANGFUSE_DEBUG}")
+        logger.debug(f"LANGFUSE_TIMEOUT: {LANGFUSE_TIMEOUT}")
+        logger.debug(f"LANGFUSE_SAMPLE_RATE: {LANGFUSE_SAMPLE_RATE * 100}%")
+
+        self.client = Langfuse(
+            public_key=public_key,
+            secret_key=secret_key,
+            host=host,
+            release=release,
+            debug=debug,
+            timeout=timeout,
+            sample_rate=sample_rate
+        )
+
+    def setup_trace(self,
+                    name: str,
+                    input: typing.Optional[typing.Any] = None,
+                    tags: typing.Optional[typing.List] = None,
+                    metadata: typing.Optional[typing.Dict] = None,
+                    user_id: str = None,
+                    session_id: str = None) -> None:
+        """
+        Setup trace. If Langfuse is disabled, nothing will be done.
+        Args:
+            name (str): Trace name.
+            input (dict): Trace input.
+            tags (dict): Trace tags.
+            metadata (dict): Trace metadata.
+            user_id (str): User ID.
+            session_id (str): Session ID.
+        """
+
+        if self.client is None:
+            logger.debug("Langfuse is disabled.")
+            return
+
+        self.set_metadata(metadata)
+        self.set_tags(tags)
+
+        try:
+            self.trace = self.client.trace(
+                name=name,
+                input=input,
+                metadata=self.metadata,
+                tags=self.tags,
+                user_id=user_id,
+                session_id=session_id
+            )
+        except Exception as e:
+            logger.error(f'Something went wrong while processing Langfuse trace {self.trace.id}: {str(e)}')
+
+        logger.info(f"Langfuse trace configured with ID: {self.trace.id}")
+
+    def start_span(self,
+                   name: str,
+                   input: typing.Optional[typing.Any] = None) -> typing.Optional[StatefulSpanClient]:
+        """
+        Create span. If Langfuse is disabled, nothing will be done.
+
+        Args:
+            name (str): Span name.
+            input (dict): Span input.
+        """
+
+        if self.client is None:
+            logger.debug("Langfuse is disabled.")
+            return None
+
+        return self.trace.span(name=name, input=input)
+
+    def end_span_stream(self,
+                        span: typing.Optional[StatefulSpanClient] = None) -> None:
+        """
+        End span. If Langfuse is disabled, nothing will happen.
+        Args:
+            span (Any): Span object.
+        """
+
+        if self.client is None:
+            logger.debug("Langfuse is disabled.")
+            return
+
+        span.end()
+        self.trace.update()
+
+    def end_span(self,
+                 span: typing.Optional[StatefulSpanClient] = None,
+                 output: typing.Optional[typing.Any] = None) -> None:
+        """
+        End trace. If Langfuse is disabled, nothing will be done.
+
+        Args:
+            span (Any): Span object.
+            output (Any): Span output.
+        """
+
+        if self.client is None:
+            logger.debug("Langfuse is disabled.")
+            return
+
+        if span is None:
+            logger.debug("Langfuse span is not created.")
+            return
+
+        span.end(output=output)
+        self.trace.update(output=output)
+
+        metadata = self.metadata or {}
+
+        try:
+            # Ensure all batched traces are sent before fetching.
+            self.client.flush()
+            metadata['tool_usage'] = self._get_tool_usage()
+            self.trace.update(metadata=metadata)
+
+        except Exception as e:
+            logger.error(f'Something went wrong while processing Langfuse trace {self.trace.id}: {str(e)}')
+
+    def get_langchain_handler(self) -> typing.Optional[CallbackHandler]:
+        """
+        Get Langchain handler. If Langfuse is disabled, returns None.
+        """
+
+        if self.client is None:
+            logger.debug("Langfuse is disabled.")
+            return None
+
+        return self.trace.get_langchain_handler()
+
+    def set_metadata(self, custom_metadata: dict = None) -> None:
+        """
+        Get default metadata.
+        """
+        self.metadata = custom_metadata or {}
+
+        self.metadata["environment"] = self.environment
+        self.metadata["release"] = self.release
+
+    def set_tags(self, custom_tags: typing.Optional[typing.List] = None) -> None:
+        """
+        Get default tags.
+        """
+        self.tags = custom_tags or []
+
+        self.tags.append(self.environment)
+        self.tags.append(self.release)
+
+    def _get_tool_usage(self) -> typing.Dict:
+        """ Retrieves tool usage information from a langfuse trace.
+        Note: assumes trace marks an action with string `AgentAction` """
+
+        tool_usage = {}
+
+        try:
+            fetched_trace = self.client.get_trace(self.trace.id)
+            steps = [s.name for s in fetched_trace.observations]
+            for step in steps:
+                if 'AgentAction' in step:
+                    tool_name = step.split('-')[1]
+                    if tool_name not in tool_usage:
+                        tool_usage[tool_name] = 0
+                    tool_usage[tool_name] += 1
+        except TraceNotFoundError:
+            logger.warning(f'Langfuse trace {self.trace.id} not found')
+        except Exception as e:
+            logger.error(f'Something went wrong while processing Langfuse trace {self.trace.id}: {str(e)}')
+
+        return tool_usage

mindsdb/utilities/log.py

@@ -29,6 +29,23 @@ class ColorFormatter(logging.Formatter):
         return log_fmt.format(record)
 
 
+def get_console_handler_config_level() -> int:
+    console_handler_config = app_config['logging']['handlers']['console']
+    return getattr(logging, console_handler_config["level"])
+
+
+def get_file_handler_config_level() -> int:
+    file_handler_config = app_config['logging']['handlers']['file']
+    return getattr(logging, file_handler_config["level"])
+
+
+def get_mindsdb_log_level() -> int:
+    console_handler_config_level = get_console_handler_config_level()
+    file_handler_config_level = get_file_handler_config_level()
+
+    return min(console_handler_config_level, file_handler_config_level)
+
+
 def configure_logging():
     handlers_config = {}
     console_handler_config = app_config['logging']['handlers']['console']
@@ -39,6 +56,7 @@ def configure_logging():
             "formatter": "f",
             "level": console_handler_config_level
         }
+
     file_handler_config = app_config['logging']['handlers']['file']
     file_handler_config_level = getattr(logging, file_handler_config["level"])
     if file_handler_config['enabled'] is True:
@@ -51,7 +69,7 @@ def configure_logging():
             "backupCount": file_handler_config["backupCount"]
         }
 
-    mindsdb_log_level = min(console_handler_config_level, file_handler_config_level)
+    mindsdb_log_level = get_mindsdb_log_level()
 
     logging_config = dict(
         version=1,
@@ -65,7 +83,7 @@ def configure_logging():
         loggers={
             "": {  # root logger
                 "handlers": list(handlers_config.keys()),
-                "level": logging.WARNING,
+                "level": mindsdb_log_level,
             },
             "__main__": {
                 "level": mindsdb_log_level,

mindsdb/utilities/otel/__init__.py

@@ -0,0 +1,206 @@
+import os
+import typing
+
+from opentelemetry import trace  # noqa: F401
+from opentelemetry import metrics  # noqa: F401
+from opentelemetry.exporter.otlp.proto.grpc._log_exporter import OTLPLogExporter as OTLPLogExporterGRPC
+from opentelemetry.exporter.otlp.proto.http._log_exporter import OTLPLogExporter as OTLPLogExporterHTTP
+from opentelemetry.sdk._logs._internal.export import LogExporter
+from opentelemetry.exporter.otlp.proto.grpc.metric_exporter import OTLPMetricExporter as OTLPMetricExporterGRPC
+from opentelemetry.exporter.otlp.proto.http.metric_exporter import OTLPMetricExporter as OTLPMetricExporterHTTP
+from opentelemetry.sdk.metrics.export import MetricExporter, ConsoleMetricExporter
+from opentelemetry.exporter.otlp.proto.grpc.trace_exporter import OTLPSpanExporter as OTLPSpanExporterGRPC
+from opentelemetry.exporter.otlp.proto.http.trace_exporter import OTLPSpanExporter as OTLPSpanExporterHTTP
+from opentelemetry.sdk.trace.export import SpanExporter, ConsoleSpanExporter
+from opentelemetry.sdk.resources import Resource
+from opentelemetry.sdk.trace.sampling import TraceIdRatioBased
+
+from mindsdb.utilities.otel.logger import setup_logger
+from mindsdb.utilities.otel.meter import setup_meter
+from mindsdb.utilities.otel.tracer import setup_tracer
+from mindsdb.utilities.utils import parse_csv_attributes
+from mindsdb.utilities import log
+
+logger = log.getLogger(__name__)
+
+# Check OpenTelemetry exporter type
+OTEL_EXPORTER_TYPE = os.getenv("OTEL_EXPORTER_TYPE", "console")  # console or otlp
+
+# Define OpenTelemetry exporter protocol
+OTEL_EXPORTER_PROTOCOL = os.getenv("OTEL_EXPORTER_PROTOCOL", "grpc")  # grpc or http
+
+# Define OTLP endpoint. If not set, the default OTLP endpoint will be used
+OTEL_OTLP_ENDPOINT = os.getenv("OTEL_OTLP_ENDPOINT", "http://localhost:4317")
+
+# Define OTLP logging endpoint. If not set, the default OTLP logging endpoint will be used
+OTEL_OTLP_LOGGING_ENDPOINT = os.getenv("OTEL_OTLP_LOGGING_ENDPOINT", OTEL_OTLP_ENDPOINT)
+
+# Define OTLP tracing endpoint. If not set, the default OTLP tracing endpoint will be used
+OTEL_OTLP_TRACING_ENDPOINT = os.getenv("OTEL_OTLP_TRACING_ENDPOINT", OTEL_OTLP_ENDPOINT)
+
+# Define OTLP metrics endpoint. If not set, the default OTLP metrics endpoint will be used
+OTEL_OTLP_METRICS_ENDPOINT = os.getenv("OTEL_OTLP_METRICS_ENDPOINT", OTEL_OTLP_ENDPOINT)
+
+# Define service name
+OTEL_SERVICE_NAME = os.getenv("OTEL_SERVICE_NAME", "mindsdb")
+
+# Define service instace ID
+OTEL_SERVICE_INSTANCE_ID = os.getenv("OTEL_SERVICE_INSTANCE_ID", "mindsdb-instance")
+
+# The name of the environment we"re on, by default local for development, this is set differently per-env in our Helm
+# chart values files
+OTEL_SERVICE_ENVIRONMENT = os.getenv("OTEL_SERVICE_ENVIRONMENT", "local").lower()
+
+# Define service release
+OTEL_SERVICE_RELEASE = os.getenv("OTEL_SERVICE_RELEASE", "local").lower()
+
+# Define how often to capture traces
+OTEL_TRACE_SAMPLE_RATE = float(os.getenv("OTEL_TRACE_SAMPLE_RATE", "1.0"))
+
+# Define extra attributes
+OTEL_EXTRA_ATTRIBUTES = os.getenv("OTEL_EXTRA_ATTRIBUTES", "")
+
+# By default, we have Open Telemetry SDK enabled on all envs, except for local which is disabled by default.
+OTEL_SDK_DISABLED = (os.getenv("OTEL_SDK_DISABLED", "false").lower() == "true"
+                     or os.getenv("OTEL_SERVICE_ENVIRONMENT", "local").lower() == "local")
+
+# Define if OpenTelemetry logging is disabled. By default, it is disabled.
+OTEL_LOGGING_DISABLED = os.getenv("OTEL_LOGGING_DISABLED", "true").lower() == "true"
+
+# Define if OpenTelemetry tracing is disabled. By default, it is enabled.
+OTEL_TRACING_DISABLED = os.getenv("OTEL_TRACING_DISABLED", "false").lower() == "true"
+
+# Define if OpenTelemetry metrics is disabled. By default, it is disabled.
+OTEL_METRICS_DISABLED = os.getenv("OTEL_METRICS_DISABLED", "true").lower() == "true"
+
+# If you want to enable Open Telemetry on local for some reason please set OTEL_SDK_FORCE_RUN to true
+OTEL_SDK_FORCE_RUN = os.getenv("OTEL_SDK_FORCE_RUN", "false").lower() == "true"
+
+
+def get_otel_attributes() -> dict:
+    """
+    Get OpenTelemetry attributes
+
+    Returns:
+        dict: OpenTelemetry attributes
+    """
+
+    base_attributes = {
+        "service.name": OTEL_SERVICE_NAME,
+        "service.instance.id": OTEL_SERVICE_INSTANCE_ID,
+        "environment": OTEL_SERVICE_ENVIRONMENT,
+        "release": OTEL_SERVICE_RELEASE,
+    }
+
+    extra_attributes = {}
+    try:
+        extra_attributes = parse_csv_attributes(OTEL_EXTRA_ATTRIBUTES)
+    except Exception as e:
+        logger.error(f"Failed to parse OTEL_EXTRA_ATTRIBUTES: {e}")
+
+    attributes = {**extra_attributes, **base_attributes}  # Base attributes take precedence over extra attributes
+
+    return attributes
+
+
+def get_logging_exporter() -> typing.Optional[LogExporter]:
+    """
+    Get OpenTelemetry logging exporter.
+
+    Returns:
+        OTLPLogExporter: OpenTelemetry logging exporter
+    """
+
+    if OTEL_EXPORTER_TYPE == "otlp":
+
+        if OTEL_EXPORTER_PROTOCOL == "grpc":
+            return OTLPLogExporterGRPC(
+                endpoint=OTEL_OTLP_LOGGING_ENDPOINT,
+                insecure=True
+            )
+
+        elif OTEL_EXPORTER_PROTOCOL == "http":
+            return OTLPLogExporterHTTP(
+                endpoint=OTEL_OTLP_LOGGING_ENDPOINT
+            )
+
+    return None
+
+
+def get_span_exporter() -> SpanExporter:
+    """
+    Get OpenTelemetry span exporter
+
+    Returns:
+        OTLPSpanExporter: OpenTelemetry span exporter
+    """
+
+    if OTEL_EXPORTER_TYPE == "otlp":
+
+        if OTEL_EXPORTER_PROTOCOL == "grpc":
+            return OTLPSpanExporterGRPC(
+                endpoint=OTEL_OTLP_TRACING_ENDPOINT,
+                insecure=True
+            )
+
+        elif OTEL_EXPORTER_PROTOCOL == "http":
+            return OTLPSpanExporterHTTP(
+                endpoint=OTEL_OTLP_TRACING_ENDPOINT
+            )
+
+    return ConsoleSpanExporter()
+
+
+def get_metrics_exporter() -> typing.Optional[MetricExporter]:
+    """
+    Get OpenTelemetry metrics exporter
+
+    Returns:
+        OTLPLogExporter: OpenTelemetry metrics exporter
+    """
+
+    if OTEL_EXPORTER_TYPE == "otlp":
+
+        if OTEL_EXPORTER_PROTOCOL == "grpc":
+            return OTLPMetricExporterGRPC(
+                endpoint=OTEL_OTLP_METRICS_ENDPOINT,
+                insecure=True
+            )
+
+        elif OTEL_EXPORTER_PROTOCOL == "http":
+            return OTLPMetricExporterHTTP(
+                endpoint=OTEL_OTLP_METRICS_ENDPOINT
+            )
+
+    return ConsoleMetricExporter()
+
+
+if not OTEL_SDK_DISABLED or OTEL_SDK_FORCE_RUN:
+    logger.info("OpenTelemetry enabled")
+    logger.info(f"OpenTelemetry exporter type: {OTEL_EXPORTER_TYPE}")
+    logger.info(f"OpenTelemetry service name: {OTEL_SERVICE_NAME}")
+    logger.info(f"OpenTelemetry service environment: {OTEL_SERVICE_ENVIRONMENT}")
+    logger.info(f"OpenTelemetry service release: {OTEL_SERVICE_RELEASE}")
+    logger.info(f"OpenTelemetry trace sample rate: {OTEL_TRACE_SAMPLE_RATE}")
+    logger.info(f"OpenTelemetry extra attributes: {OTEL_EXTRA_ATTRIBUTES}")
+
+    # Define OpenTelemetry resources (e.g., service name)
+    attributes = get_otel_attributes()
+
+    # Define OpenTelemetry sampler
+    sampler = TraceIdRatioBased(OTEL_TRACE_SAMPLE_RATE)
+
+    # Define OpenTelemetry resources (e.g., service name)
+    resource = Resource(attributes=attributes)
+
+    if not OTEL_LOGGING_DISABLED:
+        logger.info("OpenTelemetry Logging is enabled")
+        setup_logger(resource, get_logging_exporter())
+
+    if not OTEL_TRACING_DISABLED:
+        logger.info("OpenTelemetry Tracing is enabled")
+        setup_tracer(resource, sampler, get_span_exporter())
+
+    if not OTEL_METRICS_DISABLED:
+        logger.info("OpenTelemetry Metrics is enabled")
+        setup_meter(resource, get_metrics_exporter())

mindsdb/utilities/otel/logger.py

@@ -0,0 +1,25 @@
+import logging
+
+from opentelemetry._logs import set_logger_provider
+from opentelemetry.sdk._logs._internal.export import LogExporter
+from opentelemetry.sdk._logs import LoggerProvider, LoggingHandler
+from opentelemetry.sdk._logs.export import BatchLogRecordProcessor
+from opentelemetry.sdk.resources import Resource
+
+from mindsdb.utilities.log import get_mindsdb_log_level
+
+
+def setup_logger(resource: Resource, exporter: LogExporter) -> None:
+    """
+    Setup OpenTelemetry logging
+    """
+    mindsdb_log_level = get_mindsdb_log_level()
+
+    logger_provider = LoggerProvider(resource=resource)
+    set_logger_provider(logger_provider)
+
+    logger_provider.add_log_record_processor(BatchLogRecordProcessor(exporter))
+    handler = LoggingHandler(level=mindsdb_log_level, logger_provider=logger_provider)
+
+    # Attach OTLP handler to root logger
+    logging.getLogger().addHandler(handler)

mindsdb/utilities/otel/meter.py

@@ -0,0 +1,19 @@
+from opentelemetry import metrics
+from opentelemetry.sdk.metrics import MeterProvider
+from opentelemetry.sdk.resources import Resource
+from opentelemetry.sdk.metrics.export import (
+    MetricExporter,
+    PeriodicExportingMetricReader,
+)
+
+
+def setup_meter(resource: Resource, exporter: MetricExporter) -> None:
+    """
+    Setup OpenTelemetry metrics
+    """
+
+    metric_reader = PeriodicExportingMetricReader(exporter=exporter)
+    provider = MeterProvider(resource=resource, metric_readers=[metric_reader])
+
+    # Sets the global default meter provider
+    metrics.set_meter_provider(provider)

mindsdb/utilities/otel/metric_handlers/__init__.py

@@ -0,0 +1,25 @@
+from mindsdb.utilities.otel import metrics, OTEL_SERVICE_NAME
+
+_query_request_counter = None
+
+
+def get_query_request_counter():
+    """
+    Get the query request counter
+    """
+    global _query_request_counter
+
+    # Create the counter if it does not exist
+    if _query_request_counter is None:
+        meter_name = f"{OTEL_SERVICE_NAME}.query_service_meter"
+        
+        # Get the meter from the main metrics object
+        meter = metrics.get_meter(meter_name)
+
+        _query_request_counter = meter.create_counter(
+            name="query_request_count",
+            description="Counts the number of times the SQL query is called",
+            unit="1",
+        )
+
+    return _query_request_counter

mindsdb/utilities/otel/tracer.py

@@ -0,0 +1,16 @@
+from opentelemetry import trace
+from opentelemetry.sdk.resources import Resource
+from opentelemetry.sdk.trace import TracerProvider
+from opentelemetry.sdk.trace.export import BatchSpanProcessor, SpanExporter
+from opentelemetry.sdk.trace.sampling import Sampler
+
+
+def setup_tracer(resource: Resource, sampler: Sampler, exporter: SpanExporter) -> None:
+    """
+    Setup OpenTelemetry tracing
+    """
+    # Set the tracer provider with the custom resource
+    trace.set_tracer_provider(TracerProvider(resource=resource, sampler=sampler))
+
+    # Replace the default span processor with the custom one
+    trace.get_tracer_provider().add_span_processor(BatchSpanProcessor(exporter))

mindsdb/utilities/partitioning.py

@@ -0,0 +1,52 @@
+import os
+from typing import Iterable, Callable
+import pandas as pd
+
+from mindsdb.utilities.config import Config
+from mindsdb.utilities.context_executor import execute_in_threads
+
+
+def process_dataframe_in_partitions(df: pd.DataFrame, callback: Callable, partition_size: int) -> Iterable:
+    """
+    Splits dataframe into partitions and apply callback on each partition
+
+    :param df: input dataframe
+    :param callback: function to apply on each partition
+    :param partition_size: size of each partition
+    :return: yield results
+    """
+
+    # 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)
+
+    # 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
+
+    # 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
+
+    if max_threads == 1:
+        # don't spawn threads
+
+        for task in tasks:
+            yield callback(*task)
+
+    else:
+        for result in execute_in_threads(callback, tasks, thread_count=max_threads):
+            yield result

mindsdb/utilities/render/sqlalchemy_render.py

@@ -14,6 +14,10 @@ from sqlalchemy.sql import functions as sa_fnc
 from mindsdb_sql_parser import ast
 
 
+RESERVED_WORDS = {
+    "collation"
+}
+
 sa_type_names = [
     key for key, val in sa.types.__dict__.items() if hasattr(val, '__module__')
     and val.__module__ in ('sqlalchemy.sql.sqltypes', 'sqlalchemy.sql.type_api')
@@ -98,7 +102,7 @@ class SqlalchemyRender:
                     #   in that case use origin string
 
                     part_lower = str(sa.column(i.lower()).compile(dialect=self.dialect))
-                    if part.lower() != part_lower:
+                    if part.lower() != part_lower and i.lower() not in RESERVED_WORDS:
                         part = i
 
             parts2.append(part)
@@ -506,6 +510,8 @@ class SqlalchemyRender:
                             condition = self.to_expression(item['condition'])
 
                         join_type = item['join_type']
+                        if 'ASOF' in join_type:
+                            raise NotImplementedError(f'Unsupported join type: {join_type}')
                         method = 'join'
                         is_full = False
                         if join_type == 'LEFT JOIN':