MindsDB 25.5.4.1__py3-none-any.whl → 25.6.2.0__py3-none-any.whl

mindsdb/interfaces/database/integrations.py

@@ -32,19 +32,19 @@ from mindsdb.utilities import log
 from mindsdb.integrations.libs.ml_exec_base import BaseMLEngineExec
 from mindsdb.integrations.libs.base import BaseHandler
 import mindsdb.utilities.profiler as profiler
+from mindsdb.interfaces.data_catalog.data_catalog_loader import DataCatalogLoader
 
 logger = log.getLogger(__name__)
 
 
 class HandlersCache:
-    """ Cache for data handlers that keep connections opened during ttl time from handler last use
-    """
+    """Cache for data handlers that keep connections opened during ttl time from handler last use"""
 
     def __init__(self, ttl: int = 60):
-        """ init cache
+        """init cache
 
-            Args:
-                ttl (int): time to live (in seconds) for record in cache
+        Args:
+            ttl (int): time to live (in seconds) for record in cache
         """
         self.ttl = ttl
         self.handlers = {}
@@ -56,50 +56,46 @@ class HandlersCache:
         self._stop_clean()
 
     def _start_clean(self) -> None:
-        """ start worker that close connections after ttl expired
-        """
-        if (
-            isinstance(self.cleaner_thread, threading.Thread)
-            and self.cleaner_thread.is_alive()
-        ):
+        """start worker that close connections after ttl expired"""
+        if isinstance(self.cleaner_thread, threading.Thread) and self.cleaner_thread.is_alive():
             return
         self._stop_event.clear()
-        self.cleaner_thread = threading.Thread(target=self._clean, name='HandlersCache.clean')
+        self.cleaner_thread = threading.Thread(target=self._clean, name="HandlersCache.clean")
         self.cleaner_thread.daemon = True
         self.cleaner_thread.start()
 
     def _stop_clean(self) -> None:
-        """ stop clean worker
-        """
+        """stop clean worker"""
         self._stop_event.set()
 
     def set(self, handler: DatabaseHandler):
-        """ add (or replace) handler in cache
+        """add (or replace) handler in cache
 
-            Args:
-                handler (DatabaseHandler)
+        Args:
+            handler (DatabaseHandler)
         """
         with self._lock:
             try:
                 # If the handler is defined to be thread safe, set 0 as the last element of the key, otherwise set the thrad ID.
-                key = (handler.name, ctx.company_id, 0 if getattr(handler, 'thread_safe', False) else threading.get_native_id())
+                key = (
+                    handler.name,
+                    ctx.company_id,
+                    0 if getattr(handler, "thread_safe", False) else threading.get_native_id(),
+                )
                 handler.connect()
-                self.handlers[key] = {
-                    'handler': handler,
-                    'expired_at': time.time() + self.ttl
-                }
+                self.handlers[key] = {"handler": handler, "expired_at": time.time() + self.ttl}
             except Exception:
                 pass
             self._start_clean()
 
     def get(self, name: str) -> Optional[DatabaseHandler]:
-        """ get handler from cache by name
+        """get handler from cache by name
 
-            Args:
-                name (str): handler name
+        Args:
+            name (str): handler name
 
-            Returns:
-                DatabaseHandler
+        Returns:
+            DatabaseHandler
         """
         with self._lock:
             # If the handler is not thread safe, the thread ID will be assigned to the last element of the key.
@@ -107,19 +103,16 @@ class HandlersCache:
             if key not in self.handlers:
                 # If the handler is thread safe, a 0 will be assigned to the last element of the key.
                 key = (name, ctx.company_id, 0)
-            if (
-                key not in self.handlers
-                or self.handlers[key]['expired_at'] < time.time()
-            ):
+            if key not in self.handlers or self.handlers[key]["expired_at"] < time.time():
                 return None
-            self.handlers[key]['expired_at'] = time.time() + self.ttl
-            return self.handlers[key]['handler']
+            self.handlers[key]["expired_at"] = time.time() + self.ttl
+            return self.handlers[key]["handler"]
 
     def delete(self, name: str) -> None:
-        """ delete handler from cache
+        """delete handler from cache
 
-            Args:
-                name (str): handler name
+        Args:
+            name (str): handler name
         """
         with self._lock:
             key = (name, ctx.company_id, threading.get_native_id())
@@ -133,14 +126,13 @@ class HandlersCache:
                 self._stop_clean()
 
     def _clean(self) -> None:
-        """ worker that delete from cache handlers that was not in use for ttl
-        """
+        """worker that delete from cache handlers that was not in use for ttl"""
         while self._stop_event.wait(timeout=3) is False:
             with self._lock:
                 for key in list(self.handlers.keys()):
                     if (
-                        self.handlers[key]['expired_at'] < time.time()
-                        and sys.getrefcount(self.handlers[key]) == 2    # returned ref count is always 1 higher
+                        self.handlers[key]["expired_at"] < time.time()
+                        and sys.getrefcount(self.handlers[key]) == 2  # returned ref count is always 1 higher
                     ):
                         try:
                             self.handlers[key].disconnect()
@@ -163,50 +155,43 @@ class IntegrationController:
 
     def _add_integration_record(self, name, engine, connection_args):
         integration_record = db.Integration(
-            name=name,
-            engine=engine,
-            data=connection_args or {},
-            company_id=ctx.company_id
+            name=name, engine=engine, data=connection_args or {}, company_id=ctx.company_id
         )
         db.session.add(integration_record)
         db.session.commit()
         return integration_record.id
 
     def add(self, name, engine, connection_args):
-
         logger.debug(
             "%s: add method calling name=%s, engine=%s, connection_args=%s, company_id=%s",
-            self.__class__.__name__, name, engine, connection_args, ctx.company_id
+            self.__class__.__name__,
+            name,
+            engine,
+            connection_args,
+            ctx.company_id,
         )
         handler_meta = self.get_handler_meta(engine)
 
-        accept_connection_args = handler_meta.get('connection_args')
+        accept_connection_args = handler_meta.get("connection_args")
         logger.debug("%s: accept_connection_args - %s", self.__class__.__name__, accept_connection_args)
 
         files_dir = None
         if accept_connection_args is not None and connection_args is not None:
             for arg_name, arg_value in connection_args.items():
-                if (
-                    arg_name in accept_connection_args
-                    and accept_connection_args[arg_name]['type'] == ARG_TYPE.PATH
-                ):
+                if arg_name in accept_connection_args and accept_connection_args[arg_name]["type"] == ARG_TYPE.PATH:
                     if files_dir is None:
-                        files_dir = tempfile.mkdtemp(prefix='mindsdb_files_')
+                        files_dir = tempfile.mkdtemp(prefix="mindsdb_files_")
                     shutil.copy(arg_value, files_dir)
                     connection_args[arg_name] = Path(arg_value).name
 
         integration_id = self._add_integration_record(name, engine, connection_args)
 
         if files_dir is not None:
-            store = FileStorage(
-                resource_group=RESOURCE_GROUP.INTEGRATION,
-                resource_id=integration_id,
-                sync=False
-            )
-            store.add(files_dir, '')
+            store = FileStorage(resource_group=RESOURCE_GROUP.INTEGRATION, resource_id=integration_id, sync=False)
+            store.add(files_dir, "")
             store.push()
 
-        if handler_meta.get('type') == HANDLER_TYPE.ML:
+        if handler_meta.get("type") == HANDLER_TYPE.ML:
             ml_handler = self.get_ml_handler(name)
             ml_handler.create_engine(connection_args, integration_id)
 
@@ -215,7 +200,7 @@ class IntegrationController:
     def modify(self, name, data):
         self.handlers_cache.delete(name)
         integration_record = self._get_integration_record(name)
-        if isinstance(integration_record.data, dict) and integration_record.data.get('is_demo') is True:
+        if isinstance(integration_record.data, dict) and integration_record.data.get("is_demo") is True:
             raise ValueError("It is forbidden to change properties of the demo object")
         old_data = deepcopy(integration_record.data)
         for k in old_data:
@@ -226,8 +211,8 @@ class IntegrationController:
         db.session.commit()
 
     def delete(self, name):
-        if name in ('files', 'lightwood'):
-            raise Exception('Unable to drop: is system database')
+        if name in ("files", "lightwood"):
+            raise Exception("Unable to drop: is system database")
 
         self.handlers_cache.delete(name)
 
@@ -235,32 +220,32 @@ class IntegrationController:
         if name in self.handler_modules:
             handler = self.handler_modules[name]
 
-            if getattr(handler, 'permanent', False) is True:
-                raise Exception('Unable to drop permanent integration')
+            if getattr(handler, "permanent", False) is True:
+                raise Exception("Unable to drop permanent integration")
 
         integration_record = self._get_integration_record(name)
-        if isinstance(integration_record.data, dict) and integration_record.data.get('is_demo') is True:
-            raise Exception('Unable to drop demo object')
+        if isinstance(integration_record.data, dict) and integration_record.data.get("is_demo") is True:
+            raise Exception("Unable to drop demo object")
 
         # if this is ml engine
         engine_models = get_model_records(ml_handler_name=name, deleted_at=None)
         active_models = [m.name for m in engine_models if m.deleted_at is None]
         if len(active_models) > 0:
-            raise Exception(f'Unable to drop ml engine with active models: {active_models}')
+            raise Exception(f"Unable to drop ml engine with active models: {active_models}")
 
         # check linked KBs
         kb = db.KnowledgeBase.query.filter_by(vector_database_id=integration_record.id).first()
         if kb is not None:
-            raise Exception(f'Unable to drop, integration is used by knowledge base: {kb.name}')
+            raise Exception(f"Unable to drop, integration is used by knowledge base: {kb.name}")
 
         # check linked predictors
         models = get_model_records()
         for model in models:
             if (
                 model.data_integration_ref is not None
-                and model.data_integration_ref.get('type') == 'integration'
-                and isinstance(model.data_integration_ref.get('id'), int)
-                and model.data_integration_ref['id'] == integration_record.id
+                and model.data_integration_ref.get("type") == "integration"
+                and isinstance(model.data_integration_ref.get("id"), int)
+                and model.data_integration_ref["id"] == integration_record.id
             ):
                 model.data_integration_ref = None
 
@@ -269,6 +254,12 @@ class IntegrationController:
             if model.deleted_at is not None:
                 model.integration_id = None
 
+        # Remove the integration metadata from the data catalog (if enabled).
+        # TODO: Can this be handled via cascading delete in the database?
+        if Config().get("data_catalog", {}).get("enabled", False):
+            data_catalog_reader = DataCatalogLoader(database_name=name)
+            data_catalog_reader.unload_metadata()
+
         db.session.delete(integration_record)
         db.session.commit()
 
@@ -281,82 +272,77 @@ class IntegrationController:
             return None
         data = deepcopy(integration_record.data)
 
-        bundle_path = data.get('secure_connect_bundle')
-        mysql_ssl_ca = data.get('ssl_ca')
-        mysql_ssl_cert = data.get('ssl_cert')
-        mysql_ssl_key = data.get('ssl_key')
+        bundle_path = data.get("secure_connect_bundle")
+        mysql_ssl_ca = data.get("ssl_ca")
+        mysql_ssl_cert = data.get("ssl_cert")
+        mysql_ssl_key = data.get("ssl_key")
         if (
-            data.get('type') in ('mysql', 'mariadb')
+            data.get("type") in ("mysql", "mariadb")
             and (
                 self._is_not_empty_str(mysql_ssl_ca)
                 or self._is_not_empty_str(mysql_ssl_cert)
                 or self._is_not_empty_str(mysql_ssl_key)
             )
-            or data.get('type') in ('cassandra', 'scylla')
+            or data.get("type") in ("cassandra", "scylla")
             and bundle_path is not None
         ):
             fs_store = FsStore()
-            integrations_dir = Config()['paths']['integrations']
-            folder_name = f'integration_files_{integration_record.company_id}_{integration_record.id}'
-            fs_store.get(
-                folder_name,
-                base_dir=integrations_dir
-            )
+            integrations_dir = Config()["paths"]["integrations"]
+            folder_name = f"integration_files_{integration_record.company_id}_{integration_record.id}"
+            fs_store.get(folder_name, base_dir=integrations_dir)
 
         handler_meta = self.get_handler_metadata(integration_record.engine)
         integration_type = None
         if isinstance(handler_meta, dict):
             # in other cases, the handler directory is likely not exist.
-            integration_type = handler_meta.get('type')
+            integration_type = handler_meta.get("type")
 
         if show_secrets is False and handler_meta is not None:
-            connection_args = handler_meta.get('connection_args', None)
+            connection_args = handler_meta.get("connection_args", None)
             if isinstance(connection_args, dict):
                 if integration_type == HANDLER_TYPE.DATA:
                     for key, value in connection_args.items():
-                        if key in data and value.get('secret', False) is True:
-                            data[key] = '******'
+                        if key in data and value.get("secret", False) is True:
+                            data[key] = "******"
                 elif integration_type == HANDLER_TYPE.ML:
-                    creation_args = connection_args.get('creation_args')
+                    creation_args = connection_args.get("creation_args")
                     if isinstance(creation_args, dict):
                         for key, value in creation_args.items():
-                            if key in data and value.get('secret', False) is True:
-                                data[key] = '******'
+                            if key in data and value.get("secret", False) is True:
+                                data[key] = "******"
                 else:
-                    raise ValueError(f'Unexpected handler type: {integration_type}')
+                    raise ValueError(f"Unexpected handler type: {integration_type}")
             else:
                 # region obsolete, del in future
-                if 'password' in data:
-                    data['password'] = None
+                if "password" in data:
+                    data["password"] = None
                 if (
-                    data.get('type') == 'redis'
-                    and isinstance(data.get('connection'), dict)
-                    and 'password' in data['connection']
+                    data.get("type") == "redis"
+                    and isinstance(data.get("connection"), dict)
+                    and "password" in data["connection"]
                 ):
-                    data['connection'] = None
+                    data["connection"] = None
                 # endregion
 
         class_type, permanent = None, False
         if handler_meta is not None:
-            class_type = handler_meta.get('class_type')
-            permanent = handler_meta.get('permanent', False)
+            class_type = handler_meta.get("class_type")
+            permanent = handler_meta.get("permanent", False)
 
         return {
-            'id': integration_record.id,
-            'name': integration_record.name,
-            'type': integration_type,
-            'class_type': class_type,
-            'engine': integration_record.engine,
-            'permanent': permanent,
-            'date_last_update': deepcopy(integration_record.updated_at),  # to del ?
-            'connection_data': data
+            "id": integration_record.id,
+            "name": integration_record.name,
+            "type": integration_type,
+            "class_type": class_type,
+            "engine": integration_record.engine,
+            "permanent": permanent,
+            "date_last_update": deepcopy(integration_record.updated_at),  # to del ?
+            "connection_data": data,
         }
 
     def get_by_id(self, integration_id, show_secrets=True):
         integration_record = (
-            db.session.query(db.Integration)
-            .filter_by(company_id=ctx.company_id, id=integration_id)
-            .first()
+            db.session.query(db.Integration).filter_by(company_id=ctx.company_id, id=integration_id).first()
         )
         return self._get_integration_record_data(integration_record, show_secrets)
 
@@ -379,20 +365,21 @@ class IntegrationController:
             db.Integration
         """
         if case_sensitive:
-            integration_records = db.session.query(db.Integration).filter_by(
-                company_id=ctx.company_id,
-                name=name
-            ).all()
+            integration_records = db.session.query(db.Integration).filter_by(company_id=ctx.company_id, name=name).all()
             if len(integration_records) > 1:
                 raise Exception(f"There is {len(integration_records)} integrations with name '{name}'")
             if len(integration_records) == 0:
                 raise EntityNotExistsError(f"There is no integration with name '{name}'")
             integration_record = integration_records[0]
         else:
-            integration_record = db.session.query(db.Integration).filter(
-                (db.Integration.company_id == ctx.company_id)
-                & (func.lower(db.Integration.name) == func.lower(name))
-            ).first()
+            integration_record = (
+                db.session.query(db.Integration)
+                .filter(
+                    (db.Integration.company_id == ctx.company_id)
+                    & (func.lower(db.Integration.name) == func.lower(name))
+                )
+                .first()
+            )
             if integration_record is None:
                 raise EntityNotExistsError(f"There is no integration with name '{name}'")
 
@@ -407,21 +394,28 @@ class IntegrationController:
             integration_dict[record.name] = self._get_integration_record_data(record, show_secrets)
         return integration_dict
 
-    def _make_handler_args(self, name: str, handler_type: str, connection_data: dict, integration_id: int = None,
-                           file_storage: FileStorage = None, handler_storage: HandlerStorage = None):
+    def _make_handler_args(
+        self,
+        name: str,
+        handler_type: str,
+        connection_data: dict,
+        integration_id: int = None,
+        file_storage: FileStorage = None,
+        handler_storage: HandlerStorage = None,
+    ):
         handler_args = dict(
             name=name,
             integration_id=integration_id,
             connection_data=connection_data,
             file_storage=file_storage,
-            handler_storage=handler_storage
+            handler_storage=handler_storage,
         )
 
-        if handler_type == 'files':
-            handler_args['file_controller'] = FileController()
+        if handler_type == "files":
+            handler_args["file_controller"] = FileController()
         elif self.handler_modules.get(handler_type, False).type == HANDLER_TYPE.ML:
-            handler_args['handler_controller'] = self
-            handler_args['company_id'] = ctx.company_id
+            handler_args["handler_controller"] = self
+            handler_args["company_id"] = ctx.company_id
 
         return handler_args
 
@@ -439,12 +433,9 @@ class IntegrationController:
         integration_id = int(time.time() * 10000)
 
         file_storage = FileStorage(
-            resource_group=RESOURCE_GROUP.INTEGRATION,
-            resource_id=integration_id,
-            root_dir='tmp',
-            sync=False
+            resource_group=RESOURCE_GROUP.INTEGRATION, resource_id=integration_id, root_dir="tmp", sync=False
         )
-        handler_storage = HandlerStorage(integration_id, root_dir='tmp', is_temporal=True)
+        handler_storage = HandlerStorage(integration_id, root_dir="tmp", is_temporal=True)
 
         handler_meta = self.get_handler_meta(engine)
         if handler_meta is None:
@@ -466,7 +457,7 @@ class IntegrationController:
 
     def copy_integration_storage(self, integration_id_from, integration_id_to):
         storage_from = HandlerStorage(integration_id_from)
-        root_path = ''
+        root_path = ""
 
         if storage_from.is_empty():
             return None
@@ -494,7 +485,7 @@ class IntegrationController:
         if integration_meta is None:
             raise Exception(f"Handler '{name}' does not exists")
 
-        if integration_meta.get('type') != HANDLER_TYPE.ML:
+        if integration_meta.get("type") != HANDLER_TYPE.ML:
             raise Exception(f"Handler '{name}' must be ML type")
 
         logger.info(
@@ -504,7 +495,7 @@ class IntegrationController:
         handler = BaseMLEngineExec(
             name=integration_record.name,
             integration_id=integration_record.id,
-            handler_module=self.handler_modules[integration_engine]
+            handler_module=self.handler_modules[integration_engine],
         )
 
         return handler
@@ -532,35 +523,36 @@ class IntegrationController:
         if integration_meta is None:
             raise Exception(f"Handler '{name}' does not exist")
 
-        if integration_meta.get('type') != HANDLER_TYPE.DATA:
+        if integration_meta.get("type") != HANDLER_TYPE.DATA:
             raise Exception(f"Handler '{name}' must be DATA type")
 
         integration_data = self._get_integration_record_data(integration_record, True)
         if integration_data is None:
             raise Exception(f"Can't find integration_record for handler '{name}'")
-        connection_data = integration_data.get('connection_data', {})
+        connection_data = integration_data.get("connection_data", {})
         logger.debug(
             "%s.get_handler: connection_data=%s, engine=%s",
             self.__class__.__name__,
-            connection_data, integration_engine
+            connection_data,
+            integration_engine,
         )
 
         if integration_meta["import"]["success"] is False:
-            msg = dedent(f'''\
+            msg = dedent(f"""\
                 Handler '{integration_engine}' cannot be used. Reason is:
-                    {integration_meta['import']['error_message']}
-            ''')
-            is_cloud = Config().get('cloud', False)
+                    {integration_meta["import"]["error_message"]}
+            """)
+            is_cloud = Config().get("cloud", False)
             if is_cloud is False:
-                msg += dedent(f'''
+                msg += dedent(f"""
 
                 If error is related to missing dependencies, then try to run command in shell and restart mindsdb:
                     pip install mindsdb[{integration_engine}]
-                ''')
+                """)
             logger.debug(msg)
             raise Exception(msg)
 
-        connection_args = integration_meta.get('connection_args')
+        connection_args = integration_meta.get("connection_args")
         logger.debug("%s.get_handler: connection args - %s", self.__class__.__name__, connection_args)
 
         file_storage = FileStorage(
@@ -572,11 +564,11 @@ class IntegrationController:
 
         if isinstance(connection_args, (dict, OrderedDict)):
             files_to_get = {
-                arg_name: arg_value for arg_name, arg_value in connection_data.items()
-                if arg_name in connection_args and connection_args.get(arg_name)['type'] == ARG_TYPE.PATH
+                arg_name: arg_value
+                for arg_name, arg_value in connection_data.items()
+                if arg_name in connection_args and connection_args.get(arg_name)["type"] == ARG_TYPE.PATH
             }
             if len(files_to_get) > 0:
-
                 for file_name, file_path in files_to_get.items():
                     connection_data[file_name] = file_storage.get_path(file_path)
 
@@ -584,9 +576,9 @@ class IntegrationController:
             name=name,
             handler_type=integration_engine,
             connection_data=connection_data,
-            integration_id=integration_data['id'],
+            integration_id=integration_data["id"],
             file_storage=file_storage,
-            handler_storage=handler_storage
+            handler_storage=handler_storage,
         )
 
         HandlerClass = self.handler_modules[integration_engine].Handler
@@ -602,82 +594,76 @@ class IntegrationController:
             handler_meta = self._get_handler_meta(handler_name)
         except Exception as e:
             handler_meta = self.handlers_import_status[handler_name]
-            handler_meta['import']['success'] = False
-            handler_meta['import']['error_message'] = str(e)
+            handler_meta["import"]["success"] = False
+            handler_meta["import"]["error_message"] = str(e)
 
         self.handlers_import_status[handler_name] = handler_meta
 
     def _read_dependencies(self, path):
         dependencies = []
-        requirements_txt = Path(path).joinpath('requirements.txt')
+        requirements_txt = Path(path).joinpath("requirements.txt")
         if requirements_txt.is_file():
-            with open(str(requirements_txt), 'rt') as f:
-                dependencies = [x.strip(' \t\n') for x in f.readlines()]
+            with open(str(requirements_txt), "rt") as f:
+                dependencies = [x.strip(" \t\n") for x in f.readlines()]
                 dependencies = [x for x in dependencies if len(x) > 0]
         return dependencies
 
     def _get_handler_meta(self, handler_name):
-
         module = self.handler_modules[handler_name]
 
         handler_dir = Path(module.__path__[0])
         handler_folder_name = handler_dir.name
 
-        import_error = getattr(module, 'import_error', None)
+        import_error = getattr(module, "import_error", None)
         handler_meta = self.handlers_import_status[handler_name]
-        handler_meta['import']['success'] = import_error is None
-        handler_meta['version'] = module.version
-        handler_meta['thread_safe'] = getattr(module, 'thread_safe', False)
+        handler_meta["import"]["success"] = import_error is None
+        handler_meta["version"] = module.version
+        handler_meta["thread_safe"] = getattr(module, "thread_safe", False)
 
         if import_error is not None:
-            handler_meta['import']['error_message'] = str(import_error)
+            handler_meta["import"]["error_message"] = str(import_error)
 
-        handler_type = getattr(module, 'type', None)
+        handler_type = getattr(module, "type", None)
         handler_class = None
-        if hasattr(module, 'Handler') and inspect.isclass(module.Handler):
+        if hasattr(module, "Handler") and inspect.isclass(module.Handler):
             handler_class = module.Handler
             if issubclass(handler_class, BaseMLEngine):
-                handler_meta['class_type'] = 'ml'
+                handler_meta["class_type"] = "ml"
             elif issubclass(handler_class, DatabaseHandler):
-                handler_meta['class_type'] = 'sql'
+                handler_meta["class_type"] = "sql"
             if issubclass(handler_class, APIHandler):
-                handler_meta['class_type'] = 'api'
+                handler_meta["class_type"] = "api"
 
         if handler_type == HANDLER_TYPE.ML:
             # for ml engines, patch the connection_args from the argument probing
             if handler_class:
                 try:
                     prediction_args = handler_class.prediction_args()
-                    creation_args = getattr(module, 'creation_args', handler_class.creation_args())
-                    connection_args = {
-                        "prediction": prediction_args,
-                        "creation_args": creation_args
-                    }
-                    setattr(module, 'connection_args', connection_args)
+                    creation_args = getattr(module, "creation_args", handler_class.creation_args())
+                    connection_args = {"prediction": prediction_args, "creation_args": creation_args}
+                    setattr(module, "connection_args", connection_args)
                     logger.debug("Patched connection_args for %s", handler_folder_name)
                 except Exception as e:
                     # do nothing
                     logger.debug("Failed to patch connection_args for %s, reason: %s", handler_folder_name, str(e))
 
-        module_attrs = [attr for attr in [
-            'connection_args_example',
-            'connection_args',
-            'description',
-            'type',
-            'title'
-        ] if hasattr(module, attr)]
+        module_attrs = [
+            attr
+            for attr in ["connection_args_example", "connection_args", "description", "type", "title"]
+            if hasattr(module, attr)
+        ]
 
         for attr in module_attrs:
             handler_meta[attr] = getattr(module, attr)
 
         # endregion
-        if hasattr(module, 'permanent'):
-            handler_meta['permanent'] = module.permanent
+        if hasattr(module, "permanent"):
+            handler_meta["permanent"] = module.permanent
         else:
-            if handler_meta.get('name') in ('files', 'views', 'lightwood'):
-                handler_meta['permanent'] = True
+            if handler_meta.get("name") in ("files", "views", "lightwood"):
+                handler_meta["permanent"] = True
             else:
-                handler_meta['permanent'] = False
+                handler_meta["permanent"] = False
 
         return handler_meta
 
@@ -685,60 +671,60 @@ class IntegrationController:
         icon = {}
         try:
             icon_path = handler_dir.joinpath(icon_path)
-            icon_type = icon_path.name[icon_path.name.rfind('.') + 1:].lower()
+            icon_type = icon_path.name[icon_path.name.rfind(".") + 1 :].lower()
 
-            if icon_type == 'svg':
-                with open(str(icon_path), 'rt') as f:
-                    icon['data'] = f.read()
+            if icon_type == "svg":
+                with open(str(icon_path), "rt") as f:
+                    icon["data"] = f.read()
             else:
-                with open(str(icon_path), 'rb') as f:
-                    icon['data'] = base64.b64encode(f.read()).decode('utf-8')
+                with open(str(icon_path), "rb") as f:
+                    icon["data"] = base64.b64encode(f.read()).decode("utf-8")
 
-            icon['name'] = icon_path.name
-            icon['type'] = icon_type
+            icon["name"] = icon_path.name
+            icon["type"] = icon_type
 
         except Exception as e:
-            logger.error(f'Error reading icon for {handler_dir}, {e}!')
+            logger.error(f"Error reading icon for {handler_dir}, {e}!")
         return icon
 
     def _load_handler_modules(self):
-        mindsdb_path = Path(importlib.util.find_spec('mindsdb').origin).parent
-        handlers_path = mindsdb_path.joinpath('integrations/handlers')
+        mindsdb_path = Path(importlib.util.find_spec("mindsdb").origin).parent
+        handlers_path = mindsdb_path.joinpath("integrations/handlers")
 
         # edge case: running from tests directory, find_spec finds the base folder instead of actual package
         if not os.path.isdir(handlers_path):
-            mindsdb_path = Path(importlib.util.find_spec('mindsdb').origin).parent.joinpath('mindsdb')
-            handlers_path = mindsdb_path.joinpath('integrations/handlers')
+            mindsdb_path = Path(importlib.util.find_spec("mindsdb").origin).parent.joinpath("mindsdb")
+            handlers_path = mindsdb_path.joinpath("integrations/handlers")
 
         self.handler_modules = {}
         self.handlers_import_status = {}
         for handler_dir in handlers_path.iterdir():
-            if handler_dir.is_dir() is False or handler_dir.name.startswith('__'):
+            if handler_dir.is_dir() is False or handler_dir.name.startswith("__"):
                 continue
 
             handler_info = self._get_handler_info(handler_dir)
-            if 'name' not in handler_info:
+            if "name" not in handler_info:
                 continue
-            handler_name = handler_info['name']
+            handler_name = handler_info["name"]
             dependencies = self._read_dependencies(handler_dir)
             handler_meta = {
-                'path': handler_dir,
-                'import': {
-                    'success': None,
-                    'error_message': None,
-                    'folder': handler_dir.name,
-                    'dependencies': dependencies,
+                "path": handler_dir,
+                "import": {
+                    "success": None,
+                    "error_message": None,
+                    "folder": handler_dir.name,
+                    "dependencies": dependencies,
                 },
-                'name': handler_name,
-                'permanent': handler_info.get('permanent', False),
-                'connection_args': handler_info.get('connection_args', None),
-                'class_type': handler_info.get('class_type', None),
-                'type': handler_info.get('type')
+                "name": handler_name,
+                "permanent": handler_info.get("permanent", False),
+                "connection_args": handler_info.get("connection_args", None),
+                "class_type": handler_info.get("class_type", None),
+                "type": handler_info.get("type"),
             }
-            if 'icon_path' in handler_info:
-                icon = self._get_handler_icon(handler_dir, handler_info['icon_path'])
+            if "icon_path" in handler_info:
+                icon = self._get_handler_icon(handler_dir, handler_info["icon_path"])
                 if icon:
-                    handler_meta['icon'] = icon
+                    handler_meta["icon"] = icon
             self.handlers_import_status[handler_name] = handler_meta
 
     def _get_connection_args(self, args_file: Path, param_name: str) -> dict:
@@ -758,7 +744,7 @@ class IntegrationController:
                 continue
             if not item.targets[0].id == param_name:
                 continue
-            if hasattr(item.value, 'keywords'):
+            if hasattr(item.value, "keywords"):
                 for keyword in item.value.keywords:
                     name = keyword.arg
                     params = keyword.value
@@ -802,7 +788,7 @@ class IntegrationController:
         if module_file is None:
             return
 
-        path = handler_dir / f'{module_file}.py'
+        path = handler_dir / f"{module_file}.py"
 
         if not path.exists():
             return
@@ -812,9 +798,9 @@ class IntegrationController:
         for item in code.body:
             if isinstance(item, ast.ClassDef):
                 bases = [base.id for base in item.bases]
-                if 'APIHandler' in bases:
-                    return 'api'
-        return 'sql'
+                if "APIHandler" in bases:
+                    return "api"
+        return "sql"
 
     def _get_handler_info(self, handler_dir: Path) -> dict:
         """
@@ -825,7 +811,7 @@ class IntegrationController:
           - connection arguments
         """
 
-        init_file = handler_dir / '__init__.py'
+        init_file = handler_dir / "__init__.py"
         if not init_file.exists():
             return {}
         code = ast.parse(init_file.read_text())
@@ -838,47 +824,47 @@ class IntegrationController:
                 name = item.targets[0].id
                 if isinstance(item.value, ast.Constant):
                     info[name] = item.value.value
-                if isinstance(item.value, ast.Attribute) and name == 'type':
-                    if item.value.attr == 'ML':
+                if isinstance(item.value, ast.Attribute) and name == "type":
+                    if item.value.attr == "ML":
                         info[name] = HANDLER_TYPE.ML
-                        info['class_type'] = 'ml'
+                        info["class_type"] = "ml"
                     else:
                         info[name] = HANDLER_TYPE.DATA
-                        info['class_type'] = self._get_base_class_type(code, handler_dir) or 'sql'
+                        info["class_type"] = self._get_base_class_type(code, handler_dir) or "sql"
 
         # connection args
-        if info['type'] == HANDLER_TYPE.ML:
-            args_file = handler_dir / 'creation_args.py'
+        if info["type"] == HANDLER_TYPE.ML:
+            args_file = handler_dir / "creation_args.py"
             if args_file.exists():
-                info['connection_args'] = {
+                info["connection_args"] = {
                     "prediction": {},
-                    "creation_args": self._get_connection_args(args_file, 'creation_args')
+                    "creation_args": self._get_connection_args(args_file, "creation_args"),
                 }
         else:
-            args_file = handler_dir / 'connection_args.py'
+            args_file = handler_dir / "connection_args.py"
             if args_file.exists():
-                info['connection_args'] = self._get_connection_args(args_file, 'connection_args')
+                info["connection_args"] = self._get_connection_args(args_file, "connection_args")
 
         return info
 
     def import_handler(self, handler_name: str, base_import: str = None):
         with self._import_lock:
             handler_meta = self.handlers_import_status[handler_name]
-            handler_dir = handler_meta['path']
+            handler_dir = handler_meta["path"]
 
             handler_folder_name = str(handler_dir.name)
             if base_import is None:
-                base_import = 'mindsdb.integrations.handlers.'
+                base_import = "mindsdb.integrations.handlers."
 
             try:
-                handler_module = importlib.import_module(f'{base_import}{handler_folder_name}')
+                handler_module = importlib.import_module(f"{base_import}{handler_folder_name}")
                 self.handler_modules[handler_name] = handler_module
                 handler_meta = self._get_handler_meta(handler_name)
             except Exception as e:
-                handler_meta['import']['success'] = False
-                handler_meta['import']['error_message'] = str(e)
+                handler_meta["import"]["success"] = False
+                handler_meta["import"]["error_message"] = str(e)
 
-            self.handlers_import_status[handler_meta['name']] = handler_meta
+            self.handlers_import_status[handler_meta["name"]] = handler_meta
             return handler_meta
 
     def get_handlers_import_status(self):