MindsDB 25.9.3rc1__py3-none-any.whl → 25.10.0__py3-none-any.whl

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/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"]