MindsDB 25.8.2.0__py3-none-any.whl → 25.9.1.0__py3-none-any.whl

mindsdb/api/mysql/mysql_proxy/mysql_proxy.py

@@ -63,7 +63,7 @@ from mindsdb.api.mysql.mysql_proxy.libs.constants.mysql import (
     NULL_VALUE,
     COMMANDS,
     ERR,
-    getConstName
+    getConstName,
 )
 from mindsdb.api.executor.data_types.answer import ExecuteAnswer
 from mindsdb.api.executor.data_types.response_type import RESPONSE_TYPE
@@ -73,7 +73,7 @@ from mindsdb.api.mysql.mysql_proxy.utilities import (
 )
 from mindsdb.api.executor import exceptions as exec_exc
 
-from mindsdb.api.common.check_auth import check_auth
+from mindsdb.api.common.middleware import check_auth
 from mindsdb.api.mysql.mysql_proxy.libs.constants.mysql import MYSQL_DATA_TYPE
 from mindsdb.api.executor.sql_query.result_set import Column, ResultSet
 from mindsdb.utilities import log
@@ -116,10 +116,7 @@ class SQLAnswer:
             return {
                 "type": RESPONSE_TYPE.TABLE,
                 "data": data,
-                "column_names": [
-                    column.alias or column.name
-                    for column in self.result_set.columns
-                ],
+                "column_names": [column.alias or column.name for column in self.result_set.columns],
             }
         elif self.resp_type == RESPONSE_TYPE.ERROR:
             return {
@@ -149,17 +146,13 @@ class MysqlProxy(SocketServer.BaseRequestHandler):
         super().__init__(request, client_address, server)
 
     def init_session(self):
-        logger.debug(
-            "New connection [{ip}:{port}]".format(
-                ip=self.client_address[0], port=self.client_address[1]
-            )
-        )
+        logger.debug("New connection [{ip}:{port}]".format(ip=self.client_address[0], port=self.client_address[1]))
 
         if self.server.connection_id >= 65025:
             self.server.connection_id = 0
         self.server.connection_id += 1
         self.connection_id = self.server.connection_id
-        self.session = SessionController(api_type='sql')
+        self.session = SessionController(api_type="sql")
 
         if hasattr(self.server, "salt") and isinstance(self.server.salt, str):
             self.salt = self.server.salt
@@ -223,10 +216,9 @@ class MysqlProxy(SocketServer.BaseRequestHandler):
             self.session.is_ssl = True
 
             ssl_context = ssl.SSLContext()
+            ssl_context.minimum_version = ssl.TLSVersion.TLSv1_2
             ssl_context.load_cert_chain(self.server.cert_path)
-            ssl_socket = ssl_context.wrap_socket(
-                self.socket, server_side=True, do_handshake_on_connect=True
-            )
+            ssl_socket = ssl_context.wrap_socket(self.socket, server_side=True, do_handshake_on_connect=True)
 
             self.socket = ssl_socket
             handshake_resp = self.packet(HandshakeResponsePacket)
@@ -245,10 +237,7 @@ class MysqlProxy(SocketServer.BaseRequestHandler):
                     else "mysql_native_password"
                 )
 
-                if (
-                    new_method == "caching_sha2_password"
-                    and self.session.is_ssl is False
-                ):
+                if new_method == "caching_sha2_password" and self.session.is_ssl is False:
                     logger.warning(
                         f"Check auth, user={username}, ssl={self.session.is_ssl}, auth_method={client_auth_plugin}: "
                         "error: cant switch to caching_sha2_password without SSL"
@@ -307,9 +296,7 @@ class MysqlProxy(SocketServer.BaseRequestHandler):
             f"connecting to database {self.session.database}"
         )
 
-        auth_data = self.server.check_auth(
-            username, password, scramble_func, self.salt, ctx.company_id
-        )
+        auth_data = self.server.check_auth(username, password, scramble_func, self.salt, ctx.company_id)
         if auth_data["success"]:
             self.session.username = auth_data["username"]
             self.session.auth = True
@@ -349,9 +336,7 @@ class MysqlProxy(SocketServer.BaseRequestHandler):
         elif answer.type == RESPONSE_TYPE.OK:
             self.packet(OkPacket, state_track=answer.state_track, affected_rows=answer.affected_rows).send()
         elif answer.type == RESPONSE_TYPE.ERROR:
-            self.packet(
-                ErrPacket, err_code=answer.error_code, msg=answer.error_message
-            ).send()
+            self.packet(ErrPacket, err_code=answer.error_code, msg=answer.error_message).send()
 
     def _get_column_defenition_packets(self, columns: dict, data=None):
         if data is None:
@@ -370,14 +355,14 @@ class MysqlProxy(SocketServer.BaseRequestHandler):
             flags = column.get("flags", 0)
             if isinstance(flags, list):
                 flags = sum(flags)
-            if column.get('size') is None:
+            if column.get("size") is None:
                 length = 1
                 for row in data:
                     if isinstance(row, dict):
                         length = max(len(str(row[column_alias])), length)
                     else:
                         length = max(len(str(row[i])), length)
-                column['size'] = 1
+                column["size"] = 1
 
             packets.append(
                 self.packet(
@@ -397,7 +382,7 @@ class MysqlProxy(SocketServer.BaseRequestHandler):
 
     def get_table_packets(self, result_set: ResultSet, status=0):
         data_frame, columns_dict = dump_result_set_to_mysql(result_set)
-        data = data_frame.to_dict('split')['data']
+        data = data_frame.to_dict("split")["data"]
 
         # TODO remove columns order
         packets = [self.packet(ColumnCountPacket, count=len(columns_dict))]
@@ -431,10 +416,12 @@ class MysqlProxy(SocketServer.BaseRequestHandler):
 
         chunk_size = 100
         for start in range(0, len(df), chunk_size):
-            string = b"".join([
-                self.packet(body=body, length=len(body)).accum()
-                for body in df[start:start + chunk_size].applymap(apply_f).values.sum(axis=1)
-            ])
+            string = b"".join(
+                [
+                    self.packet(body=body, length=len(body)).accum()
+                    for body in df[start : start + chunk_size].applymap(apply_f).values.sum(axis=1)
+                ]
+            )
             self.socket.sendall(string)
 
     def decode_utf(self, text):
@@ -510,7 +497,7 @@ class MysqlProxy(SocketServer.BaseRequestHandler):
             resp = SQLAnswer(
                 resp_type=RESPONSE_TYPE.OK,
                 state_track=executor_answer.state_track,
-                affected_rows=executor_answer.affected_rows
+                affected_rows=executor_answer.affected_rows,
             )
         else:
             resp = SQLAnswer(
@@ -519,7 +506,7 @@ class MysqlProxy(SocketServer.BaseRequestHandler):
                 result_set=executor_answer.data,
                 status=executor.server_status,
                 affected_rows=executor_answer.affected_rows,
-                mysql_types=executor_answer.data.mysql_types
+                mysql_types=executor_answer.data.mysql_types,
             )
 
         # Increment the counter and include metadata in attributes
@@ -568,15 +555,13 @@ class MysqlProxy(SocketServer.BaseRequestHandler):
         executor_answer: ExecuteAnswer = executor.executor_answer
 
         if executor_answer.data is None:
-            resp = SQLAnswer(
-                resp_type=RESPONSE_TYPE.OK, state_track=executor_answer.state_track
-            )
+            resp = SQLAnswer(resp_type=RESPONSE_TYPE.OK, state_track=executor_answer.state_track)
             return self.send_query_answer(resp)
 
         # TODO prepared_stmt['type'] == 'lock' is not used but it works
         result_set = executor_answer.data
         data_frame, columns_dict = dump_result_set_to_mysql(result_set)
-        data = data_frame.to_dict('split')['data']
+        data = data_frame.to_dict("split")["data"]
 
         packages = [self.packet(ColumnCountPacket, count=len(columns_dict))]
         packages.extend(self._get_column_defenition_packets(columns_dict))
@@ -586,9 +571,7 @@ class MysqlProxy(SocketServer.BaseRequestHandler):
 
         # send all
         for row in data:
-            packages.append(
-                self.packet(BinaryResultsetRowPacket, data=row, columns=columns_dict)
-            )
+            packages.append(self.packet(BinaryResultsetRowPacket, data=row, columns=columns_dict))
 
         server_status = executor.server_status or 0x0002
         packages.append(self.last_packet(status=server_status))
@@ -603,17 +586,13 @@ class MysqlProxy(SocketServer.BaseRequestHandler):
         executor_answer: ExecuteAnswer = executor.executor_answer
 
         if executor_answer.data is None:
-            resp = SQLAnswer(
-                resp_type=RESPONSE_TYPE.OK, state_track=executor_answer.state_track
-            )
+            resp = SQLAnswer(resp_type=RESPONSE_TYPE.OK, state_track=executor_answer.state_track)
             return self.send_query_answer(resp)
 
         packages = []
         columns = self.to_mysql_columns(executor_answer.data.columns)
         for row in executor_answer.data[fetched:limit].to_lists():
-            packages.append(
-                self.packet(BinaryResultsetRowPacket, data=row, columns=columns)
-            )
+            packages.append(self.packet(BinaryResultsetRowPacket, data=row, columns=columns))
 
         prepared_stmt["fetched"] += len(executor_answer.data[fetched:limit])
 
@@ -656,9 +635,7 @@ class MysqlProxy(SocketServer.BaseRequestHandler):
         else:
             ctx.user_class = cloud_connection["user_class"]
             ctx.email_confirmed = cloud_connection["email_confirmed"]
-            self.client_capabilities = ClentCapabilities(
-                cloud_connection["client_capabilities"]
-            )
+            self.client_capabilities = ClentCapabilities(cloud_connection["client_capabilities"])
             self.session.database = cloud_connection["database"]
             self.session.username = "cloud"
             self.session.auth = True
@@ -678,9 +655,7 @@ class MysqlProxy(SocketServer.BaseRequestHandler):
                 logger.debug("Session closed by client")
                 return
 
-            logger.debug(
-                "Command TYPE: {type}".format(type=getConstName(COMMANDS, p.type.value))
-            )
+            logger.debug("Command TYPE: {type}".format(type=getConstName(COMMANDS, p.type.value)))
 
             command_names = {
                 COMMANDS.COM_QUERY: "COM_QUERY",
@@ -705,10 +680,8 @@ class MysqlProxy(SocketServer.BaseRequestHandler):
                 if p.type.value == COMMANDS.COM_QUERY:
                     sql = self.decode_utf(p.sql.value)
                     sql = clear_sql(sql)
-                    logger.debug(f'Incoming query: {sql}')
-                    profiler.set_meta(
-                        query=sql, api="mysql", environment=config.get("environment")
-                    )
+                    logger.debug(f"Incoming query: {sql}")
+                    profiler.set_meta(query=sql, api="mysql", environment=config.get("environment"))
                     with profiler.Context("mysql_query_processing"):
                         response = self.process_query(sql)
                 elif p.type.value == COMMANDS.COM_STMT_PREPARE:
@@ -791,9 +764,7 @@ class MysqlProxy(SocketServer.BaseRequestHandler):
                 # any other exception
                 error_type = "unexpected"
                 error_traceback = traceback.format_exc()
-                logger.error(
-                    f"ERROR while executing query\n" f"{error_traceback}\n" f"{e}"
-                )
+                logger.error(f"ERROR while executing query\n{error_traceback}\n{e}")
                 error_code = ERR.ER_SYNTAX_ERROR
                 response = SQLAnswer(
                     resp_type=RESPONSE_TYPE.ERROR,
@@ -841,7 +812,7 @@ class MysqlProxy(SocketServer.BaseRequestHandler):
         if "db" in context:
             self.session.database = context["db"]
         else:
-            self.session.database = config.get('default_project')
+            self.session.database = config.get("default_project")
 
         if "profiling" in context:
             self.session.profiling = context["profiling"]
@@ -851,9 +822,7 @@ class MysqlProxy(SocketServer.BaseRequestHandler):
             self.session.show_secrets = context["show_secrets"]
 
     def get_context(self):
-        context = {
-            "show_secrets": self.session.show_secrets
-        }
+        context = {"show_secrets": self.session.show_secrets}
         if self.session.database is not None:
             context["db"] = self.session.database
         if self.session.profiling is True:

mindsdb/api/postgres/postgres_proxy/postgres_proxy.py

@@ -14,17 +14,38 @@ from mindsdb.api.executor.controllers import SessionController
 from mindsdb.api.postgres.postgres_proxy.executor import Executor
 from mindsdb.api.mysql.mysql_proxy.libs.constants.mysql import CHARSET_NUMBERS
 from mindsdb.api.executor.data_types.response_type import RESPONSE_TYPE
-from mindsdb.api.common.check_auth import check_auth
+from mindsdb.api.common.middleware import check_auth
 from mindsdb.api.mysql.mysql_proxy.mysql_proxy import SQLAnswer
 from mindsdb.api.postgres.postgres_proxy.postgres_packets.errors import POSTGRES_SYNTAX_ERROR_CODE
 from mindsdb.api.postgres.postgres_proxy.postgres_packets.postgres_fields import GenericField, PostgresField
-from mindsdb.api.postgres.postgres_proxy.postgres_packets.postgres_message_formats import Terminate, \
-    Query, AuthenticationClearTextPassword, AuthenticationOk, RowDescriptions, DataRow, CommandComplete, \
-    ReadyForQuery, ConnectionFailure, ParameterStatus, Error, Execute, Bind, Parse, Sync, ParseComplete, \
-    InvalidSQLStatementName, BindComplete, Describe, DataException, ParameterDescription
+from mindsdb.api.postgres.postgres_proxy.postgres_packets.postgres_message_formats import (
+    Terminate,
+    Query,
+    AuthenticationClearTextPassword,
+    AuthenticationOk,
+    RowDescriptions,
+    DataRow,
+    CommandComplete,
+    ReadyForQuery,
+    ConnectionFailure,
+    ParameterStatus,
+    Error,
+    Execute,
+    Bind,
+    Parse,
+    Sync,
+    ParseComplete,
+    InvalidSQLStatementName,
+    BindComplete,
+    Describe,
+    DataException,
+    ParameterDescription,
+)
 from mindsdb.api.postgres.postgres_proxy.postgres_packets.postgres_message import PostgresMessage
-from mindsdb.api.postgres.postgres_proxy.postgres_packets.postgres_packets import PostgresPacketReader, \
-    PostgresPacketBuilder
+from mindsdb.api.postgres.postgres_proxy.postgres_packets.postgres_packets import (
+    PostgresPacketReader,
+    PostgresPacketBuilder,
+)
 from mindsdb.api.postgres.postgres_proxy.utilities import strip_null_byte
 from mindsdb.utilities.config import config
 from mindsdb.utilities.context import context as ctx
@@ -38,8 +59,8 @@ class PostgresProxyHandler(socketserver.StreamRequestHandler):
 
     def __init__(self, request, client_address, server):
         self.logger = log.getLogger(__name__)
-        self.charset = 'utf8'
-        self.charset_text_type = CHARSET_NUMBERS['utf8_general_ci']
+        self.charset = "utf8"
+        self.charset_text_type = CHARSET_NUMBERS["utf8_general_ci"]
         self.session = None
         self.client_capabilities = None
         self.user_parameters = None
@@ -48,17 +69,16 @@ class PostgresProxyHandler(socketserver.StreamRequestHandler):
         self.is_cloud = False
         self.named_portals = {}
         self.unnamed_portal = None
-        self.transaction_status = b'I'  # I: Idle, T: Transaction Block, E: Failed Transaction Block
+        self.transaction_status = b"I"  # I: Idle, T: Transaction Block, E: Failed Transaction Block
         super().__init__(request, client_address, server)
 
     def handle(self) -> None:
-
         ctx.set_default()
         self.init_session()
-        self.logger.debug('handle new incoming connection')
+        self.logger.debug("handle new incoming connection")
         cloud_connection = self.is_cloud_connection()
         if cloud_connection["is_cloud"]:
-            ctx.company_id = cloud_connection.get('company_id')
+            ctx.company_id = cloud_connection.get("company_id")
             self.is_cloud = True
 
         self.message_map: Dict[Type[PostgresMessage], Callable[[Any], bool]] = {
@@ -68,7 +88,7 @@ class PostgresProxyHandler(socketserver.StreamRequestHandler):
             Bind: self.bind,
             Execute: self.execute,
             Describe: self.describe,
-            Sync: self.sync
+            Sync: self.sync,
         }
         self.client_buffer = PostgresPacketReader(self.rfile)
         if self.is_cloud:
@@ -83,69 +103,59 @@ class PostgresProxyHandler(socketserver.StreamRequestHandler):
             self.main_loop()
 
     def is_cloud_connection(self):
-        """ Determine source of connection. Must be call before handshake.
-                Idea based on: real mysql connection does not send anything before server handshake, so
-                socket should be in 'out' state. In opposite, clout connection sends '0000' right after
-                connection. '0000' selected because in real mysql connection it should be length of package,
-                and it can not be 0.
-
-                Copied from mysql_proxy
-                TODO: Extract method into common
-            """
-        is_cloud = config.get('cloud', False)
-
-        if sys.platform != 'linux' or is_cloud is False:
-            return {
-                'is_cloud': False
-            }
+        """Determine source of connection. Must be call before handshake.
+        Idea based on: real mysql connection does not send anything before server handshake, so
+        socket should be in 'out' state. In opposite, clout connection sends '0000' right after
+        connection. '0000' selected because in real mysql connection it should be length of package,
+        and it can not be 0.
+
+        Copied from mysql_proxy
+        TODO: Extract method into common
+        """
+        is_cloud = config.get("cloud", False)
+
+        if sys.platform != "linux" or is_cloud is False:
+            return {"is_cloud": False}
 
         read_poller = select.poll()
         read_poller.register(self.request, select.POLLIN)
         events = read_poller.poll(30)
 
         if len(events) == 0:
-            return {
-                'is_cloud': False
-            }
+            return {"is_cloud": False}
 
         first_byte = self.request.recv(4, socket.MSG_PEEK)
-        if first_byte == b'\x00\x00\x00\x00':
+        if first_byte == b"\x00\x00\x00\x00":
             self.request.recv(4)
             client_capabilities = self.request.recv(8)
-            client_capabilities = struct.unpack('L', client_capabilities)[0]
+            client_capabilities = struct.unpack("L", client_capabilities)[0]
 
             company_id = self.request.recv(4)
-            company_id = struct.unpack('I', company_id)[0]
+            company_id = struct.unpack("I", company_id)[0]
 
             user_class = self.request.recv(1)
-            user_class = struct.unpack('B', user_class)[0]
+            user_class = struct.unpack("B", user_class)[0]
 
             database_name_len = self.request.recv(2)
-            database_name_len = struct.unpack('H', database_name_len)[0]
+            database_name_len = struct.unpack("H", database_name_len)[0]
 
-            database_name = ''
+            database_name = ""
             if database_name_len > 0:
                 database_name = self.request.recv(database_name_len).decode()
 
             return {
-                'is_cloud': True,
-                'client_capabilities': client_capabilities,
-                'company_id': company_id,
-                'user_class': user_class,
-                'database': database_name
+                "is_cloud": True,
+                "client_capabilities": client_capabilities,
+                "company_id": company_id,
+                "user_class": user_class,
+                "database": database_name,
             }
 
-        return {
-            'is_cloud': False
-        }
+        return {"is_cloud": False}
 
     def parse(self, message: Parse):
         self.logger.info("Postgres_Proxy: Parsing")
-        executor = Executor(
-            session=self.session,
-            proxy_server=self,
-            charset=self.charset
-        )
+        executor = Executor(session=self.session, proxy_server=self, charset=self.charset)
         # TODO: Remove comment if unneeded ot use session since we're storing in this proxy class per session anyway
         # stmt_id = self.session.register_stmt(executor)
         executor.stmt_prepare(sql=message.query)
@@ -179,7 +189,7 @@ class PostgresProxyHandler(socketserver.StreamRequestHandler):
 
     def describe(self, message: Describe):
         self.logger.info("Postgres_Proxy: Describing")
-        if message.describe_type == b'P':
+        if message.describe_type == b"P":
             if message.name:
                 describing = self.named_portals[message.name]
             elif self.unnamed_statement:
@@ -187,7 +197,7 @@ class PostgresProxyHandler(socketserver.StreamRequestHandler):
             else:
                 self.send(InvalidSQLStatementName("Portal Does not Exist"))
                 return True
-        elif message.describe_type == b'S':
+        elif message.describe_type == b"S":
             if message.name:
                 describing = self.named_statements[message.name]
             elif self.unnamed_statement:
@@ -227,8 +237,7 @@ class PostgresProxyHandler(socketserver.StreamRequestHandler):
         return True
 
     def init_session(self):
-        self.logger.info('New connection [{ip}:{port}]'.format(
-            ip=self.client_address[0], port=self.client_address[1]))
+        self.logger.info("New connection [{ip}:{port}]".format(ip=self.client_address[0], port=self.client_address[1]))
         self.logger.debug(self.__dict__)
 
         if self.server.connection_id >= 65025:
@@ -236,9 +245,9 @@ class PostgresProxyHandler(socketserver.StreamRequestHandler):
         self.server.connection_id += 1
         self.connection_id = self.server.connection_id
         self.session = SessionController()
-        self.session.database = config.get('default_project')
+        self.session.database = config.get("default_project")
 
-        if hasattr(self.server, 'salt') and isinstance(self.server.salt, str):
+        if hasattr(self.server, "salt") and isinstance(self.server.salt, str):
             self.salt = self.server.salt
         else:
             self.salt = base64.b64encode(os.urandom(15)).decode()
@@ -247,14 +256,10 @@ class PostgresProxyHandler(socketserver.StreamRequestHandler):
 
         self.current_transaction = None
 
-        self.logger.debug('session salt: {salt}'.format(salt=self.salt))
+        self.logger.debug("session salt: {salt}".format(salt=self.salt))
 
     def process_query(self, sql):
-        executor = Executor(
-            session=self.session,
-            proxy_server=self,
-            charset=self.charset
-        )
+        executor = Executor(session=self.session, proxy_server=self, charset=self.charset)
         self.logger.debug("processing query\n%s", sql)
         try:
             executor.query_execute(sql)
@@ -262,7 +267,7 @@ class PostgresProxyHandler(socketserver.StreamRequestHandler):
             return SQLAnswer(
                 resp_type=RESPONSE_TYPE.ERROR,
                 error_message=str(e).encode(self.get_encoding()),
-                error_code=POSTGRES_SYNTAX_ERROR_CODE.encode(self.get_encoding())
+                error_code=POSTGRES_SYNTAX_ERROR_CODE.encode(self.get_encoding()),
             )
         return self.return_executor_data(executor)
 
@@ -277,7 +282,7 @@ class PostgresProxyHandler(socketserver.StreamRequestHandler):
                 resp_type=RESPONSE_TYPE.TABLE,
                 state_track=executor.state_track,
                 result_set=executor.data,
-                status=executor.server_status
+                status=executor.server_status,
             )
         return resp
 
@@ -295,7 +300,7 @@ class PostgresProxyHandler(socketserver.StreamRequestHandler):
         self.client_buffer.read_verify_ssl_request()
         # self.send(NoticeResponse()) -- Should Probably not send. Looks in protocol manual to be sent for warning
         self.logger.debug("Sending No to SSL Request")
-        PostgresPacketBuilder().write_char(b'N', self.wfile)
+        PostgresPacketBuilder().write_char(b"N", self.wfile)
         self.user_parameters = self.client_buffer.read_startup_message()
 
     def authenticate(self, ask_for_password=False):
@@ -303,12 +308,12 @@ class PostgresProxyHandler(socketserver.StreamRequestHandler):
             self.send(AuthenticationClearTextPassword())
             password = self.client_buffer.read_authentication(encoding=self.charset)
         else:
-            password = ''
-        username = self.user_parameters[b'user'].decode(encoding=self.charset)
+            password = ""
+        username = self.user_parameters[b"user"].decode(encoding=self.charset)
         auth_data = self.server.check_auth(username, password, scramble_func, self.salt, ctx.company_id)
-        if auth_data['success']:
+        if auth_data["success"]:
             self.logger.debug("Authentication succeeded")
-            self.session.username = auth_data['username']
+            self.session.username = auth_data["username"]
             self.session.auth = True
             self.send(AuthenticationOk())
             return True
@@ -329,9 +334,9 @@ class PostgresProxyHandler(socketserver.StreamRequestHandler):
     def return_ok(self, sql, rows: int = 0):
         command = self.get_command(sql)
         if command == "BEGIN":
-            self.transaction_status = b'T'
+            self.transaction_status = b"T"
         if command == "COMMIT":
-            self.transaction_status = b'I'
+            self.transaction_status = b"I"
         if command == "CREATE":
             command = "SELECT"
         if command in ("INSERT", "DELETE", "UPDATE", "SELECT", "MOVE", "FETCH", "COPY"):
@@ -347,7 +352,7 @@ class PostgresProxyHandler(socketserver.StreamRequestHandler):
         if type(sql) == bytes:
             encoding = self.get_encoding()
             sql: str = sql.decode(encoding)
-        return strip_null_byte(sql).strip(';')
+        return strip_null_byte(sql).strip(";")
 
     def return_table(self, sql_answer: SQLAnswer, row_descs=True):
         fields = self.to_postgres_fields(sql_answer.result_set.columns)
@@ -356,7 +361,7 @@ class PostgresProxyHandler(socketserver.StreamRequestHandler):
             self.send(RowDescriptions(fields=fields))
         self.send(DataRow(rows=rows))
         encoding = self.get_encoding()
-        tag = ('SELECT %s' % str(len(rows))).encode(encoding)
+        tag = ("SELECT %s" % str(len(rows))).encode(encoding)
         self.send(CommandComplete(tag=tag))
         return True
 
@@ -389,11 +394,7 @@ class PostgresProxyHandler(socketserver.StreamRequestHandler):
         fields = []
         i = 0
         for column in columns:
-            fields.append(GenericField(
-                name=column['name'],
-                object_id=column['type'].value,
-                column_id=i
-            ))
+            fields.append(GenericField(name=column["name"], object_id=column["type"].value, column_id=i))
             i += 1
         return fields
 
@@ -410,12 +411,12 @@ class PostgresProxyHandler(socketserver.StreamRequestHandler):
                     column = json.dumps(column)
                 if isinstance(column, datetime.date) or isinstance(column, datetime.datetime):
                     try:
-                        column = datetime.datetime.strftime(column, '%Y-%m-%d')
+                        column = datetime.datetime.strftime(column, "%Y-%m-%d")
                     except ValueError:
                         try:
-                            column = datetime.datetime.strftime(column, '%Y-%m-%dT%H:%M:%S')
+                            column = datetime.datetime.strftime(column, "%Y-%m-%dT%H:%M:%S")
                         except ValueError:
-                            column = datetime.datetime.strptime(column, '%Y-%m-%dT%H:%M:%S.%f')
+                            column = datetime.datetime.strptime(column, "%Y-%m-%dT%H:%M:%S.%f")
                 if isinstance(column, bool):
                     if column:
                         column = "true"
@@ -427,7 +428,7 @@ class PostgresProxyHandler(socketserver.StreamRequestHandler):
 
     def send_initial_data(self):
         server_encoding = self.charset.encode(self.charset)
-        client_encoding = self.user_parameters.get(b'client_encoding', server_encoding)
+        client_encoding = self.user_parameters.get(b"client_encoding", server_encoding)
         # TODO: Send BackendKeyData Here (55.2.1)
         self.send(ParameterStatus(name=b"server_version", value=b"14.6"))
         self.send(ParameterStatus(name=b"server_encoding", value=server_encoding))
@@ -459,8 +460,8 @@ class PostgresProxyHandler(socketserver.StreamRequestHandler):
 
     @staticmethod
     def startProxy():
-        host = config['api']['postgres']['host']
-        port = int(config['api']['postgres']['port'])
+        host = config["api"]["postgres"]["host"]
+        port = int(config["api"]["postgres"]["port"])
         server = TcpServer((host, port), PostgresProxyHandler)
         server.connection_id = 0
         server.mindsdb_config = config

mindsdb/integrations/handlers/crate_handler/crate_handler.py

@@ -56,9 +56,7 @@ class CrateHandler(DatabaseHandler):
         if self.is_connected:
             return self.connection
 
-        is_local = (
-            self.host.startswith("localhost") or self.host == "127.0.0.1"
-        )
+        is_local = self.host.startswith("localhost") or self.host == "127.0.0.1"
 
         try:
             # Build URL based on connection type
@@ -121,7 +119,7 @@ class CrateHandler(DatabaseHandler):
         """Receive raw query and act upon it somehow.
         Args:
             query (Any): query in native format (str for sql databases,
-                dict for mongo, etc)
+                etc)
         Returns:
             HandlerResponse
         """
@@ -136,9 +134,7 @@ class CrateHandler(DatabaseHandler):
                 result = cur.fetchall()
                 response = Response(
                     RESPONSE_TYPE.TABLE,
-                    data_frame=pd.DataFrame(
-                        result, columns=[x[0] for x in cur.description]
-                    ),
+                    data_frame=pd.DataFrame(result, columns=[x[0] for x in cur.description]),
                 )
             else:
                 response = Response(RESPONSE_TYPE.OK)