MindsDB 25.5.4.2__py3-none-any.whl → 25.6.3.0__py3-none-any.whl

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

@@ -1,13 +1,11 @@
 import time
 import inspect
 from dataclasses import astuple
-from typing import Iterable
+from typing import Iterable, List
 
 import numpy as np
 import pandas as pd
-from sqlalchemy.types import (
-    Integer, Float
-)
+from sqlalchemy.types import Integer, Float
 
 from mindsdb_sql_parser.ast.base import ASTNode
 from mindsdb_sql_parser.ast import Insert, Identifier, CreateTable, TableColumn, DropTables
@@ -32,7 +30,7 @@ class DBHandlerException(Exception):
 
 
 class IntegrationDataNode(DataNode):
-    type = 'integration'
+    type = "integration"
 
     def __init__(self, integration_name, ds_type, integration_controller):
         self.integration_name = integration_name
@@ -46,15 +44,17 @@ class IntegrationDataNode(DataNode):
     def get_tables(self):
         response = self.integration_handler.get_tables()
         if response.type == RESPONSE_TYPE.TABLE:
-            result_dict = response.data_frame.to_dict(orient='records')
+            result_dict = response.data_frame.to_dict(orient="records")
             result = []
             for row in result_dict:
-
                 result.append(TablesRow.from_dict(row))
             return result
         else:
             raise Exception(f"Can't get tables: {response.error_message}")
 
+        result_dict = response.data_frame.to_dict(orient="records")
+        return [TablesRow.from_dict(row) for row in result_dict]
+
     def get_table_columns_df(self, table_name: str, schema_name: str | None = None) -> pd.DataFrame:
         """Get a DataFrame containing representation of information_schema.columns for the specified table.
 
@@ -66,7 +66,7 @@ class IntegrationDataNode(DataNode):
             pd.DataFrame: A DataFrame containing representation of information_schema.columns for the specified table.
                           The DataFrame has list of columns as in the integrations.libs.response.INF_SCHEMA_COLUMNS_NAMES.
         """
-        if 'schema_name' in inspect.signature(self.integration_handler.get_columns).parameters:
+        if "schema_name" in inspect.signature(self.integration_handler.get_columns).parameters:
             response = self.integration_handler.get_columns(table_name, schema_name)
         else:
             response = self.integration_handler.get_columns(table_name)
@@ -81,18 +81,18 @@ class IntegrationDataNode(DataNode):
         # region fallback for old handlers
         df = response.data_frame
         df.columns = [name.upper() for name in df.columns]
-        if 'FIELD' not in df.columns or 'TYPE' not in df.columns:
+        if "FIELD" not in df.columns or "TYPE" not in df.columns:
             logger.warning(
                 f"Response from the handler's `get_columns` call does not contain required columns: f{df.columns}"
             )
             return pd.DataFrame([], columns=astuple(INF_SCHEMA_COLUMNS_NAMES))
 
-        new_df = df[['FIELD', 'TYPE']]
-        new_df.columns = ['COLUMN_NAME', 'DATA_TYPE']
+        new_df = df[["FIELD", "TYPE"]]
+        new_df.columns = ["COLUMN_NAME", "DATA_TYPE"]
 
-        new_df[INF_SCHEMA_COLUMNS_NAMES.MYSQL_DATA_TYPE] = new_df[
-            INF_SCHEMA_COLUMNS_NAMES.DATA_TYPE
-        ].apply(lambda x: infer_mysql_type(x).value)
+        new_df[INF_SCHEMA_COLUMNS_NAMES.MYSQL_DATA_TYPE] = new_df[INF_SCHEMA_COLUMNS_NAMES.DATA_TYPE].apply(
+            lambda x: infer_mysql_type(x).value
+        )
 
         for column_name in astuple(INF_SCHEMA_COLUMNS_NAMES):
             if column_name in new_df.columns:
@@ -116,54 +116,50 @@ class IntegrationDataNode(DataNode):
         return df[INF_SCHEMA_COLUMNS_NAMES.COLUMN_NAME].to_list()
 
     def drop_table(self, name: Identifier, if_exists=False):
-        drop_ast = DropTables(
-            tables=[name],
-            if_exists=if_exists
-        )
+        drop_ast = DropTables(tables=[name], if_exists=if_exists)
         self.query(drop_ast)
 
-    def create_table(self, table_name: Identifier, result_set: ResultSet = None, columns=None,
-                     is_replace=False, is_create=False, **kwargs) -> DataHubResponse:
+    def create_table(
+        self,
+        table_name: Identifier,
+        result_set: ResultSet = None,
+        columns: List[TableColumn] = None,
+        is_replace: bool = False,
+        is_create: bool = False,
+        raise_if_exists: bool = True,
+        **kwargs,
+    ) -> DataHubResponse:
         # is_create - create table
+        #   if !raise_if_exists: error will be skipped
         # is_replace - drop table if exists
         # is_create==False and is_replace==False: just insert
 
         table_columns_meta = {}
 
         if columns is None:
-            columns = []
-
-            df = result_set.get_raw_df()
-
             columns: list[TableColumn] = result_set.get_ast_columns()
-            table_columns_meta = {
-                column.name: column.type
-                for column in columns
-            }
+            table_columns_meta = {column.name: column.type for column in columns}
 
         if is_replace:
             # drop
-            drop_ast = DropTables(
-                tables=[table_name],
-                if_exists=True
-            )
+            drop_ast = DropTables(tables=[table_name], if_exists=True)
             self.query(drop_ast)
             is_create = True
 
         if is_create:
-            create_table_ast = CreateTable(
-                name=table_name,
-                columns=columns,
-                is_replace=is_replace
-            )
-            self.query(create_table_ast)
+            create_table_ast = CreateTable(name=table_name, columns=columns, is_replace=is_replace)
+            try:
+                self.query(create_table_ast)
+            except Exception as e:
+                if raise_if_exists:
+                    raise e
 
         if result_set is None:
             # it is just a 'create table'
             return DataHubResponse()
 
         # native insert
-        if hasattr(self.integration_handler, 'insert'):
+        if hasattr(self.integration_handler, "insert"):
             df = result_set.to_df()
 
             result: HandlerResponse = self.integration_handler.insert(table_name.parts[-1], df)
@@ -176,9 +172,9 @@ class IntegrationDataNode(DataNode):
             column_type = table_columns_meta[col.alias]
 
             if column_type == Integer:
-                type_name = 'int'
+                type_name = "int"
             elif column_type == Float:
-                type_name = 'float'
+                type_name = "float"
             else:
                 continue
 
@@ -193,24 +189,19 @@ class IntegrationDataNode(DataNode):
             # not need to insert
             return DataHubResponse()
 
-        insert_ast = Insert(
-            table=table_name,
-            columns=insert_columns,
-            values=values,
-            is_plain=True
-        )
+        insert_ast = Insert(table=table_name, columns=insert_columns, values=values, is_plain=True)
 
         try:
             result: DataHubResponse = self.query(insert_ast)
         except Exception as e:
-            msg = f'[{self.ds_type}/{self.integration_name}]: {str(e)}'
+            msg = f"[{self.ds_type}/{self.integration_name}]: {str(e)}"
             raise DBHandlerException(msg) from e
 
         return DataHubResponse(affected_rows=result.affected_rows)
 
     def has_support_stream(self) -> bool:
         # checks if data handler has query_stream method
-        return hasattr(self.integration_handler, 'query_stream') and callable(self.integration_handler.query_stream)
+        return hasattr(self.integration_handler, "query_stream") and callable(self.integration_handler.query_stream)
 
     @profiler.profile()
     def query_stream(self, query: ASTNode, fetch_size: int = None) -> Iterable:
@@ -230,24 +221,26 @@ class IntegrationDataNode(DataNode):
             # metrics
             elapsed_seconds = time.perf_counter() - time_before_query
             query_time_with_labels = metrics.INTEGRATION_HANDLER_QUERY_TIME.labels(
-                get_class_name(self.integration_handler), result.type)
+                get_class_name(self.integration_handler), result.type
+            )
             query_time_with_labels.observe(elapsed_seconds)
 
             num_rows = 0
             if result.data_frame is not None:
                 num_rows = len(result.data_frame.index)
             response_size_with_labels = metrics.INTEGRATION_HANDLER_RESPONSE_SIZE.labels(
-                get_class_name(self.integration_handler), result.type)
+                get_class_name(self.integration_handler), result.type
+            )
             response_size_with_labels.observe(num_rows)
         except Exception as e:
             msg = str(e).strip()
-            if msg == '':
+            if msg == "":
                 msg = e.__class__.__name__
-            msg = f'[{self.ds_type}/{self.integration_name}]: {msg}'
+            msg = f"[{self.ds_type}/{self.integration_name}]: {msg}"
             raise DBHandlerException(msg) from e
 
         if result.type == RESPONSE_TYPE.ERROR:
-            raise Exception(f'Error in {self.integration_name}: {result.error_message}')
+            raise Exception(f"Error in {self.integration_name}: {result.error_message}")
         if result.type == RESPONSE_TYPE.OK:
             return DataHubResponse(affected_rows=result.affected_rows)
 
@@ -265,17 +258,8 @@ class IntegrationDataNode(DataNode):
             logger.error(f"Issue with clearing DF from NaN values: {e}")
         # endregion
 
-        columns_info = [
-            {
-                'name': k,
-                'type': v
-            }
-            for k, v in df.dtypes.items()
-        ]
+        columns_info = [{"name": k, "type": v} for k, v in df.dtypes.items()]
 
         return DataHubResponse(
-            data_frame=df,
-            columns=columns_info,
-            affected_rows=result.affected_rows,
-            mysql_types=result.mysql_types
+            data_frame=df, columns=columns_info, affected_rows=result.affected_rows, mysql_types=result.mysql_types
         )

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

@@ -26,18 +26,17 @@ def to_json(obj):
 def get_project_name(query: ASTNode = None):
     project_name = None
     if (
-            isinstance(query, Select)
-            and type(query.where) is BinaryOperation
-            and query.where.op == '='
-            and query.where.args[0].parts == ['project']
-            and isinstance(query.where.args[1], Constant)
+        isinstance(query, Select)
+        and type(query.where) is BinaryOperation
+        and query.where.op == "="
+        and query.where.args[0].parts == ["project"]
+        and isinstance(query.where.args[1], Constant)
     ):
         project_name = query.where.args[1].value
     return project_name
 
 
 class MdbTable(Table):
-
     visible: bool = True
 
 
@@ -75,27 +74,29 @@ class ModelsTable(MdbTable):
                 table_name = row["name"]
                 table_meta = row["metadata"]
 
-                data.append([
-                    table_name,
-                    table_meta["engine"],
-                    project_name,
-                    table_meta["active"],
-                    table_meta["version"],
-                    table_meta["status"],
-                    table_meta["accuracy"],
-                    table_meta["predict"],
-                    table_meta["update_status"],
-                    table_meta["mindsdb_version"],
-                    table_meta["error"],
-                    table_meta["select_data_query"],
-                    to_json(table_meta["training_options"]),
-                    table_meta["current_training_phase"],
-                    table_meta["total_training_phases"],
-                    table_meta["training_phase_name"],
-                    table_meta["label"],
-                    row["created_at"],
-                    table_meta["training_time"],
-                ])
+                data.append(
+                    [
+                        table_name,
+                        table_meta["engine"],
+                        project_name,
+                        table_meta["active"],
+                        table_meta["version"],
+                        table_meta["status"],
+                        table_meta["accuracy"],
+                        table_meta["predict"],
+                        table_meta["update_status"],
+                        table_meta["mindsdb_version"],
+                        table_meta["error"],
+                        table_meta["select_data_query"],
+                        to_json(table_meta["training_options"]),
+                        table_meta["current_training_phase"],
+                        table_meta["total_training_phases"],
+                        table_meta["training_phase_name"],
+                        table_meta["label"],
+                        row["created_at"],
+                        table_meta["training_time"],
+                    ]
+                )
             # TODO optimise here
             # if target_table is not None and target_table != project_name:
             #     continue
@@ -110,12 +111,8 @@ class DatabasesTable(MdbTable):
 
     @classmethod
     def get_data(cls, session, inf_schema, **kwargs):
-
         project = inf_schema.database_controller.get_list(with_secrets=session.show_secrets)
-        data = [
-            [x["name"], x["type"], x["engine"], to_json(x.get("connection_data"))]
-            for x in project
-        ]
+        data = [[x["name"], x["type"], x["engine"], to_json(x.get("connection_data"))] for x in project]
 
         df = pd.DataFrame(data, columns=cls.columns)
         return df
@@ -123,17 +120,12 @@ class DatabasesTable(MdbTable):
 
 class MLEnginesTable(MdbTable):
     name = "ML_ENGINES"
-    columns = [
-        "NAME", "HANDLER", "CONNECTION_DATA"
-    ]
+    columns = ["NAME", "HANDLER", "CONNECTION_DATA"]
 
     @classmethod
     def get_data(cls, session, inf_schema, **kwargs):
-
         integrations = inf_schema.integration_controller.get_all(show_secrets=session.show_secrets)
-        ml_integrations = {
-            key: val for key, val in integrations.items() if val["type"] == "ml"
-        }
+        ml_integrations = {key: val for key, val in integrations.items() if val["type"] == "ml"}
 
         data = []
         for _key, val in ml_integrations.items():
@@ -158,7 +150,6 @@ class HandlersTable(MdbTable):
 
     @classmethod
     def get_data(cls, inf_schema, **kwargs):
-
         handlers = inf_schema.integration_controller.get_handlers_import_status()
 
         data = []
@@ -272,7 +263,7 @@ class TriggersTable(MdbTable):
         data = triggers_controller.get_list(project_name)
 
         columns = cls.mindsdb_columns
-        if inf_schema.session.api_type == 'sql':
+        if inf_schema.session.api_type == "sql":
             columns = columns + cls.columns
         columns_lower = [col.lower() for col in columns]
 
@@ -283,7 +274,7 @@ class TriggersTable(MdbTable):
 
 
 class ChatbotsTable(MdbTable):
-    name = 'CHATBOTS'
+    name = "CHATBOTS"
     columns = [
         "NAME",
         "PROJECT",
@@ -319,62 +310,78 @@ class ChatbotsTable(MdbTable):
         # to list of lists
         data = []
         for row in chatbot_data:
-            row['params'] = to_json(row['params'])
+            row["params"] = to_json(row["params"])
             data.append([row[k] for k in columns_lower])
 
         return pd.DataFrame(data, columns=columns)
 
 
 class KBTable(MdbTable):
-    name = 'KNOWLEDGE_BASES'
-    columns = ["NAME", "PROJECT", "MODEL", "STORAGE", "PARAMS",
-               "INSERT_STARTED_AT", "INSERT_FINISHED_AT", "PROCESSED_ROWS", "ERROR", "QUERY_ID"]
+    name = "KNOWLEDGE_BASES"
+    columns = [
+        "NAME",
+        "PROJECT",
+        "EMBEDDING_MODEL",
+        "RERANKING_MODEL",
+        "STORAGE",
+        "METADATA_COLUMNS",
+        "CONTENT_COLUMNS",
+        "ID_COLUMN",
+        "PARAMS",
+        "INSERT_STARTED_AT",
+        "INSERT_FINISHED_AT",
+        "PROCESSED_ROWS",
+        "ERROR",
+        "QUERY_ID",
+    ]
 
     @classmethod
     def get_data(cls, query: ASTNode = None, inf_schema=None, **kwargs):
         project_name = get_project_name(query)
 
         from mindsdb.interfaces.knowledge_base.controller import KnowledgeBaseController
+
         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()
-        }
+        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']
+            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,
-            ))
+            data.append(
+                (
+                    kb["name"],
+                    kb["project_name"],
+                    to_json(kb["embedding_model"]),
+                    to_json(kb["reranking_model"]),
+                    kb["vector_database"] + "." + kb["vector_database_table"],
+                    to_json(kb["metadata_columns"]),
+                    to_json(kb["content_columns"]),
+                    kb["id_column"],
+                    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)
 
 
 class SkillsTable(MdbTable):
-    name = 'SKILLS'
+    name = "SKILLS"
     columns = ["NAME", "PROJECT", "TYPE", "PARAMS"]
 
     @classmethod
@@ -394,14 +401,8 @@ class SkillsTable(MdbTable):
 
 
 class AgentsTable(MdbTable):
-    name = 'AGENTS'
-    columns = [
-        "NAME",
-        "PROJECT",
-        "MODEL_NAME",
-        "SKILLS",
-        "PARAMS"
-    ]
+    name = "AGENTS"
+    columns = ["NAME", "PROJECT", "MODEL_NAME", "SKILLS", "PARAMS"]
 
     @classmethod
     def get_data(cls, query: ASTNode = None, inf_schema=None, **kwargs):
@@ -411,10 +412,7 @@ class AgentsTable(MdbTable):
         all_agents = agents_controller.get_agents(project_name)
 
         project_controller = ProjectController()
-        project_names = {
-            i.id: i.name
-            for i in project_controller.get_list()
-        }
+        project_names = {i.id: i.name for i in project_controller.get_list()}
 
         # NAME, PROJECT, MODEL, SKILLS, PARAMS
         data = [
@@ -423,7 +421,7 @@ class AgentsTable(MdbTable):
                 project_names[a.project_id],
                 a.model_name,
                 [rel.skill.name for rel in a.skills_relationships],
-                to_json(a.params)
+                to_json(a.params),
             )
             for a in all_agents
         ]
@@ -431,12 +429,11 @@ class AgentsTable(MdbTable):
 
 
 class ViewsTable(MdbTable):
-    name = 'VIEWS'
+    name = "VIEWS"
     columns = ["NAME", "PROJECT", "QUERY"]
 
     @classmethod
     def get_data(cls, query: ASTNode = None, **kwargs):
-
         project_name = get_project_name(query)
 
         data = ViewController().list(project_name)
@@ -450,9 +447,19 @@ class ViewsTable(MdbTable):
 
 
 class QueriesTable(MdbTable):
-    name = 'QUERIES'
-    columns = ["ID", "STARTED_AT", "FINISHED_AT", "PROCESSED_ROWS", "ERROR", "SQL", "DATABASE",
-               "PARAMETERS", "CONTEXT", "UPDATED_AT"]
+    name = "QUERIES"
+    columns = [
+        "ID",
+        "STARTED_AT",
+        "FINISHED_AT",
+        "PROCESSED_ROWS",
+        "ERROR",
+        "SQL",
+        "DATABASE",
+        "PARAMETERS",
+        "CONTEXT",
+        "UPDATED_AT",
+    ]
 
     @classmethod
     def get_data(cls, **kwargs):