MindsDB 25.9.2.0a1__py3-none-any.whl → 25.10.0rc1__py3-none-any.whl

mindsdb/integrations/handlers/databricks_handler/databricks_handler.py

@@ -1,24 +1,50 @@
 from typing import Text, Dict, Any, Optional
 
+import pandas as pd
 from databricks.sql import connect, RequestError, ServerOperationError
 from databricks.sql.client import Connection
 from databricks.sqlalchemy import DatabricksDialect
 from mindsdb_sql_parser.ast.base import ASTNode
-from mindsdb.utilities.render.sqlalchemy_render import SqlalchemyRender
-import pandas as pd
 
+from mindsdb.utilities.render.sqlalchemy_render import SqlalchemyRender
 from mindsdb.integrations.libs.base import DatabaseHandler
 from mindsdb.integrations.libs.response import (
     HandlerStatusResponse as StatusResponse,
     HandlerResponse as Response,
     RESPONSE_TYPE,
+    INF_SCHEMA_COLUMNS_NAMES_SET,
 )
 from mindsdb.utilities import log
+from mindsdb.api.mysql.mysql_proxy.libs.constants.mysql import MYSQL_DATA_TYPE
 
 
 logger = log.getLogger(__name__)
 
 
+def _map_type(internal_type_name: str | None) -> MYSQL_DATA_TYPE:
+    """Map MyDatabricks SQL text types names to MySQL types as enum.
+
+    Args:
+        internal_type_name (str): The name of the Databricks type to map.
+
+    Returns:
+        MYSQL_DATA_TYPE: The MySQL type enum that corresponds to the MySQL text type name.
+    """
+    if not isinstance(internal_type_name, str):
+        return MYSQL_DATA_TYPE.TEXT
+    if internal_type_name.upper() == "STRING":
+        return MYSQL_DATA_TYPE.TEXT
+    if internal_type_name.upper() == "LONG":
+        return MYSQL_DATA_TYPE.BIGINT
+    if internal_type_name.upper() == "SHORT":
+        return MYSQL_DATA_TYPE.SMALLINT
+    try:
+        return MYSQL_DATA_TYPE(internal_type_name.upper())
+    except Exception:
+        logger.info(f"Databricks handler: unknown type: {internal_type_name}, use TEXT as fallback.")
+        return MYSQL_DATA_TYPE.TEXT
+
+
 class DatabricksHandler(DatabaseHandler):
     """
     This handler handles the connection and execution of SQL statements on Databricks.
@@ -64,11 +90,8 @@ class DatabricksHandler(DatabaseHandler):
             return self.connection
 
         # Mandatory connection parameters.
-        if not all(
-            key in self.connection_data
-            for key in ["server_hostname", "http_path", "access_token"]
-        ):
-            raise ValueError('Required parameters (server_hostname, http_path, access_token) must be provided.')
+        if not all(key in self.connection_data for key in ["server_hostname", "http_path", "access_token"]):
+            raise ValueError("Required parameters (server_hostname, http_path, access_token) must be provided.")
 
         config = {
             "server_hostname": self.connection_data["server_hostname"],
@@ -88,19 +111,17 @@ class DatabricksHandler(DatabaseHandler):
                 config[parameter] = self.connection_data[parameter]
 
         try:
-            self.connection = connect(
-                **config
-            )
+            self.connection = connect(**config)
             self.is_connected = True
             return self.connection
         except RequestError as request_error:
-            logger.error(f'Request error when connecting to Databricks: {request_error}')
+            logger.error(f"Request error when connecting to Databricks: {request_error}")
             raise
         except RuntimeError as runtime_error:
-            logger.error(f'Runtime error when connecting to Databricks: {runtime_error}')
+            logger.error(f"Runtime error when connecting to Databricks: {runtime_error}")
             raise
         except Exception as unknown_error:
-            logger.error(f'Unknown error when connecting to Databricks: {unknown_error}')
+            logger.error(f"Unknown error when connecting to Databricks: {unknown_error}")
             raise
 
     def disconnect(self):
@@ -129,7 +150,7 @@ class DatabricksHandler(DatabaseHandler):
 
             # Execute a simple query to check the connection.
             query = "SELECT 1 FROM information_schema.schemata"
-            if 'schema' in self.connection_data:
+            if "schema" in self.connection_data:
                 query += f" WHERE schema_name = '{self.connection_data['schema']}'"
 
             with connection.cursor() as cursor:
@@ -138,14 +159,14 @@ class DatabricksHandler(DatabaseHandler):
 
             # If the query does not return a result, the schema does not exist.
             if not result:
-                raise ValueError(f'The schema {self.connection_data["schema"]} does not exist!')
+                raise ValueError(f"The schema {self.connection_data['schema']} does not exist!")
 
             response.success = True
         except (ValueError, RequestError, RuntimeError, ServerOperationError) as known_error:
-            logger.error(f'Connection check to Databricks failed, {known_error}!')
+            logger.error(f"Connection check to Databricks failed, {known_error}!")
             response.error_message = str(known_error)
         except Exception as unknown_error:
-            logger.error(f'Connection check to Databricks failed due to an unknown error, {unknown_error}!')
+            logger.error(f"Connection check to Databricks failed due to an unknown error, {unknown_error}!")
             response.error_message = str(unknown_error)
 
         if response.success and need_to_close:
@@ -176,30 +197,18 @@ class DatabricksHandler(DatabaseHandler):
                 if result:
                     response = Response(
                         RESPONSE_TYPE.TABLE,
-                        data_frame=pd.DataFrame(
-                            result, columns=[x[0] for x in cursor.description]
-                        ),
+                        data_frame=pd.DataFrame(result, columns=[x[0] for x in cursor.description]),
                     )
 
                 else:
                     response = Response(RESPONSE_TYPE.OK)
                     connection.commit()
             except ServerOperationError as server_error:
-                logger.error(
-                    f'Server error running query: {query} on Databricks, {server_error}!'
-                )
-                response = Response(
-                    RESPONSE_TYPE.ERROR,
-                    error_message=str(server_error)
-                )
+                logger.error(f"Server error running query: {query} on Databricks, {server_error}!")
+                response = Response(RESPONSE_TYPE.ERROR, error_message=str(server_error))
             except Exception as unknown_error:
-                logger.error(
-                    f'Unknown error running query: {query} on Databricks, {unknown_error}!'
-                )
-                response = Response(
-                    RESPONSE_TYPE.ERROR,
-                    error_message=str(unknown_error)
-                )
+                logger.error(f"Unknown error running query: {query} on Databricks, {unknown_error}!")
+                response = Response(RESPONSE_TYPE.ERROR, error_message=str(unknown_error))
 
         if need_to_close is True:
             self.disconnect()
@@ -220,29 +229,44 @@ class DatabricksHandler(DatabaseHandler):
         query_str = renderer.get_string(query, with_failback=True)
         return self.native_query(query_str)
 
-    def get_tables(self) -> Response:
+    def get_tables(self, all: bool = False) -> Response:
         """
         Retrieves a list of all non-system tables in the connected schema of the Databricks workspace.
 
+        Args:
+            all (bool): If True - return tables from all schemas.
+
         Returns:
             Response: A response object containing a list of tables in the connected schema.
         """
-        query = """
-            SHOW TABLES;
+        all_filter = "and table_schema = current_schema()"
+        if all is True:
+            all_filter = ""
+        query = f"""
+            SELECT
+                table_schema,
+                table_name,
+                table_type
+            FROM
+                information_schema.tables
+            WHERE
+                table_schema != 'information_schema'
+                {all_filter}
         """
         result = self.native_query(query)
-
-        df = result.data_frame
-        if df is not None:
-            result.data_frame = df.rename(columns={"tableName": "table_name", "database": "schema_name"})
+        if result.resp_type == RESPONSE_TYPE.OK:
+            result = Response(
+                RESPONSE_TYPE.TABLE, data_frame=pd.DataFrame([], columns=list(INF_SCHEMA_COLUMNS_NAMES_SET))
+            )
         return result
 
-    def get_columns(self, table_name: Text) -> Response:
+    def get_columns(self, table_name: str, schema_name: str | None = None) -> Response:
         """
         Retrieves column details for a specified table in the Databricks workspace.
 
         Args:
-            table_name (Text): The name of the table for which to retrieve column information.
+            table_name (str): The name of the table for which to retrieve column information.
+            schema_name (str|None): The name of the schema in which the table is located.
 
         Raises:
             ValueError: If the 'table_name' is not a valid string.
@@ -253,9 +277,37 @@ class DatabricksHandler(DatabaseHandler):
         if not table_name or not isinstance(table_name, str):
             raise ValueError("Invalid table name provided.")
 
-        query = f"DESCRIBE TABLE {table_name};"
+        if isinstance(schema_name, str):
+            schema_name = f"'{schema_name}'"
+        else:
+            schema_name = "current_schema()"
+        query = f"""
+            SELECT
+                COLUMN_NAME,
+                DATA_TYPE,
+                ORDINAL_POSITION,
+                COLUMN_DEFAULT,
+                IS_NULLABLE,
+                CHARACTER_MAXIMUM_LENGTH,
+                CHARACTER_OCTET_LENGTH,
+                NUMERIC_PRECISION,
+                NUMERIC_SCALE,
+                DATETIME_PRECISION,
+                null as CHARACTER_SET_NAME,
+                null as COLLATION_NAME
+            FROM
+                information_schema.columns
+            WHERE
+                table_name = '{table_name}'
+            AND
+                table_schema = {schema_name}
+        """
+
         result = self.native_query(query)
+        if result.resp_type == RESPONSE_TYPE.OK:
+            result = Response(
+                RESPONSE_TYPE.TABLE, data_frame=pd.DataFrame([], columns=list(INF_SCHEMA_COLUMNS_NAMES_SET))
+            )
+        result.to_columns_table_response(map_type_fn=_map_type)
 
-        df = result.data_frame
-        result.data_frame = df.rename(columns={"col_name": "column_name"})
         return result

mindsdb/integrations/handlers/druid_handler/druid_handler.py

@@ -8,13 +8,13 @@ from mindsdb.utilities.render.sqlalchemy_render import SqlalchemyRender
 from mindsdb.integrations.libs.base import DatabaseHandler
 from pydruid.db.sqlalchemy import DruidDialect
 
-from mindsdb_sql_parser.ast.base import ASTNode
+from mindsdb_sql_parser import ASTNode
 
 from mindsdb.utilities import log
 from mindsdb.integrations.libs.response import (
     HandlerStatusResponse as StatusResponse,
     HandlerResponse as Response,
-    RESPONSE_TYPE
+    RESPONSE_TYPE,
 )
 
 logger = log.getLogger(__name__)
@@ -25,7 +25,7 @@ class DruidHandler(DatabaseHandler):
     This handler handles connection and execution of the Apache Druid statements.
     """
 
-    name = 'druid'
+    name = "druid"
 
     def __init__(self, name: str, connection_data: Optional[dict], **kwargs):
         """
@@ -37,18 +37,18 @@ class DruidHandler(DatabaseHandler):
         """
         super().__init__(name)
         self.parser = parse_sql
-        self.dialect = 'druid'
+        self.dialect = "druid"
 
-        optional_parameters = ['user', 'password']
+        optional_parameters = ["user", "password"]
         for parameter in optional_parameters:
             if parameter not in connection_data:
                 connection_data[parameter] = None
 
-        if 'path' not in connection_data:
-            connection_data['path'] = '/druid/v2/sql/'
+        if "path" not in connection_data:
+            connection_data["path"] = "/druid/v2/sql/"
 
-        if 'scheme' not in connection_data:
-            connection_data['scheme'] = 'http'
+        if "scheme" not in connection_data:
+            connection_data["scheme"] = "http"
 
         self.connection_data = connection_data
         self.kwargs = kwargs
@@ -71,12 +71,12 @@ class DruidHandler(DatabaseHandler):
             return self.connection
 
         self.connection = connect(
-            host=self.connection_data['host'],
-            port=self.connection_data['port'],
-            path=self.connection_data['path'],
-            scheme=self.connection_data['scheme'],
-            user=self.connection_data['user'],
-            password=self.connection_data['password']
+            host=self.connection_data["host"],
+            port=self.connection_data["port"],
+            path=self.connection_data["path"],
+            scheme=self.connection_data["scheme"],
+            user=self.connection_data["user"],
+            password=self.connection_data["password"],
         )
         self.is_connected = True
 
@@ -106,11 +106,11 @@ class DruidHandler(DatabaseHandler):
 
         try:
             conn = self.connect()
-            conn.cursor().execute('select 1')  # raise exception if provided wrong credentials
+            conn.cursor().execute("select 1")  # raise exception if provided wrong credentials
 
             response.success = True
         except Exception as e:
-            logger.error(f'Error connecting to Druid, {e}!')
+            logger.error(f"Error connecting to Druid, {e}!")
             response.error_message = str(e)
         finally:
             if response.success is True and need_to_close:
@@ -139,21 +139,14 @@ class DruidHandler(DatabaseHandler):
             result = cursor.fetchall()
             if result:
                 response = Response(
-                    RESPONSE_TYPE.TABLE,
-                    data_frame=pd.DataFrame(
-                        result,
-                        columns=[x[0] for x in cursor.description]
-                    )
+                    RESPONSE_TYPE.TABLE, data_frame=pd.DataFrame(result, columns=[x[0] for x in cursor.description])
                 )
             else:
                 connection.commit()
                 response = Response(RESPONSE_TYPE.OK)
         except Exception as e:
-            logger.error(f'Error running query: {query} on Pinot!')
-            response = Response(
-                RESPONSE_TYPE.ERROR,
-                error_message=str(e)
-            )
+            logger.error(f"Error running query: {query} on Pinot!")
+            response = Response(RESPONSE_TYPE.ERROR, error_message=str(e))
 
         cursor.close()
         if need_to_close is True:
@@ -182,18 +175,18 @@ class DruidHandler(DatabaseHandler):
         """
 
         query = """
-            SELECT *
+            SELECT 
+                TABLE_SCHEMA AS table_schema,
+                TABLE_NAME AS table_name,
+                TABLE_TYPE AS table_type
             FROM INFORMATION_SCHEMA.TABLES
+            WHERE TABLE_SCHEMA not in ('INFORMATION_SCHEMA', 'sys')
         """
         result = self.native_query(query)
-        df = result.data_frame
-
-        df = df[['TABLE_NAME', 'TABLE_TYPE']]
-        result.data_frame = df.rename(columns={'TABLE_NAME': 'table_name', 'TABLE_TYPE': 'table_type'})
 
         return result
 
-    def get_columns(self, table_name: str) -> StatusResponse:
+    def get_columns(self, table_name: str, schema_name: Optional[str] = None) -> StatusResponse:
         """
         Returns a list of entity columns.
         Args:
@@ -201,16 +194,15 @@ class DruidHandler(DatabaseHandler):
         Returns:
             HandlerResponse
         """
-
+        if schema_name is None:
+            schema_name = "druid"
         query = f"""
-            SELECT *
+            SELECT
+                COLUMN_NAME FIELD,
+                DATA_TYPE TYPE
             FROM INFORMATION_SCHEMA.COLUMNS
-            WHERE "TABLE_SCHEMA" = 'druid' AND "TABLE_NAME" = '{table_name}'
+            WHERE "TABLE_SCHEMA" = '{schema_name}' AND "TABLE_NAME" = '{table_name}'
         """
         result = self.native_query(query)
-        df = result.data_frame
-
-        df = df[['COLUMN_NAME', 'DATA_TYPE']]
-        result.data_frame = df.rename(columns={'COLUMN_NAME': 'column_name', 'DATA_TYPE': 'data_type'})
 
         return result

mindsdb/integrations/handlers/file_handler/file_handler.py

@@ -50,6 +50,7 @@ class FileHandler(DatabaseHandler):
         self.chunk_size = connection_data.get("chunk_size", DEFAULT_CHUNK_SIZE)
         self.chunk_overlap = connection_data.get("chunk_overlap", DEFAULT_CHUNK_OVERLAP)
         self.file_controller = file_controller
+        self.thread_safe = True
 
     def connect(self, **kwargs):
         return
@@ -83,6 +84,12 @@ class FileHandler(DatabaseHandler):
                 table_name = table_identifier.parts[-1]
                 try:
                     self.file_controller.delete_file(table_name)
+                except FileNotFoundError as e:
+                    if not query.if_exists:
+                        return Response(
+                            RESPONSE_TYPE.ERROR,
+                            error_message=f"Can't delete table '{table_name}': {e}",
+                        )
                 except Exception as e:
                     return Response(
                         RESPONSE_TYPE.ERROR,

mindsdb/integrations/handlers/gitlab_handler/gitlab_handler.py

@@ -16,7 +16,7 @@ class GitlabHandler(APIHandler):
     """The GitLab handler implementation"""
 
     def __init__(self, name: str, **kwargs):
-        """ constructor
+        """constructor
         Args:
             name (str): the handler name
         """
@@ -36,13 +36,16 @@ class GitlabHandler(APIHandler):
         self._register_table("merge_requests", gitlab_merge_requests_data)
 
     def connect(self) -> StatusResponse:
-        """ Set up the connections required by the handler
+        """Set up the connections required by the handler
         Returns:
             HandlerStatusResponse
         """
 
         connection_kwargs = {}
 
+        if self.connection_data.get("url", None):
+            connection_kwargs["url"] = self.connection_data["url"]
+
         if self.connection_data.get("api_key", None):
             connection_kwargs["private_token"] = self.connection_data["api_key"]