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

mindsdb/api/postgres/postgres_proxy/postgres_proxy.py

@@ -1,477 +0,0 @@
-import base64
-import datetime
-import os
-import json
-import select
-import socketserver
-import struct
-import sys
-from functools import partial
-import socket
-from typing import Callable, Dict, Type, Any, Iterable, Sequence
-
-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.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 import PostgresMessage
-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
-from mindsdb.utilities import log
-from mindsdb.api.mysql.mysql_proxy.external_libs.mysql_scramble import scramble as scramble_func
-
-
-class PostgresProxyHandler(socketserver.StreamRequestHandler):
-    client_buffer: PostgresPacketReader
-    user_parameters: Dict[bytes, bytes]
-
-    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.session = None
-        self.client_capabilities = None
-        self.user_parameters = None
-        self.named_statements = {}
-        self.unnamed_statement = None
-        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
-        super().__init__(request, client_address, server)
-
-    def handle(self) -> None:
-        ctx.set_default()
-        self.init_session()
-        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")
-            self.is_cloud = True
-
-        self.message_map: Dict[Type[PostgresMessage], Callable[[Any], bool]] = {
-            Terminate: self.terminate,
-            Query: self.query,
-            Parse: self.parse,
-            Bind: self.bind,
-            Execute: self.execute,
-            Describe: self.describe,
-            Sync: self.sync,
-        }
-        self.client_buffer = PostgresPacketReader(self.rfile)
-        if self.is_cloud:
-            # We already have a connection started through the gateway.
-            started = True
-            self.handshake()
-        else:
-            started = self.start_connection()
-        if started:
-            self.logger.debug("connection started")
-            self.send_initial_data()
-            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}
-
-        read_poller = select.poll()
-        read_poller.register(self.request, select.POLLIN)
-        events = read_poller.poll(30)
-
-        if len(events) == 0:
-            return {"is_cloud": False}
-
-        first_byte = self.request.recv(4, socket.MSG_PEEK)
-        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]
-
-            company_id = self.request.recv(4)
-            company_id = struct.unpack("I", company_id)[0]
-
-            user_class = self.request.recv(1)
-            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 = ""
-            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,
-            }
-
-        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)
-        # 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)
-        statement = {"executor": executor, "parse": message}
-        if message.name:
-            self.named_statements[message.name] = statement
-        else:
-            self.unnamed_statement = statement
-        self.send(ParseComplete())
-        return True
-
-    def bind(self, message: Bind):
-        self.logger.info(f"Postgres_Proxy: Binding {message.name} with params {message.parameters}")
-        if message.statement_name:
-            statement = self.named_statements[message.statement_name]
-        elif self.unnamed_statement:
-            statement = self.unnamed_statement
-        else:
-            self.send(InvalidSQLStatementName())
-            return True
-
-        # TODO Should check validity of statement here and not at parse stage
-        portal = statement.copy()
-        portal["bind"] = message
-        if message.name:
-            self.named_portals[message.name] = portal
-        else:
-            self.unnamed_portal = portal
-        self.send(BindComplete())
-        return True
-
-    def describe(self, message: Describe):
-        self.logger.info("Postgres_Proxy: Describing")
-        if message.describe_type == b"P":
-            if message.name:
-                describing = self.named_portals[message.name]
-            elif self.unnamed_statement:
-                describing = self.unnamed_portal
-            else:
-                self.send(InvalidSQLStatementName("Portal Does not Exist"))
-                return True
-        elif message.describe_type == b"S":
-            if message.name:
-                describing = self.named_statements[message.name]
-            elif self.unnamed_statement:
-                describing = self.unnamed_statement
-            else:
-                self.send(InvalidSQLStatementName())
-                return True
-            self.send(ParameterDescription(parameters=describing["parse"]["parameters"]))
-        else:
-            self.send(DataException(message="Describe did not have correct type. Can be 'P' or 'S'"))
-            return True
-
-        fields = self.to_postgres_fields(describing["executor"].columns)
-        self.send(RowDescriptions(fields=fields))
-        return True
-
-    def execute(self, message: Execute):
-        self.logger.info("Postgres_Proxy: Executing")
-        if message.name:
-            portal = self.named_portals[message.name]
-        elif self.unnamed_portal:
-            portal = self.unnamed_portal
-        else:
-            self.send(InvalidSQLStatementName("Portal does not exist"))
-
-        executor = portal["executor"]
-        params = portal["bind"].parameters
-        executor.stmt_execute(param_values=params)
-        sql_answer = self.return_executor_data(executor)
-        self.respond_from_sql_answer(sql=executor.sql, sql_answer=sql_answer, row_descs=False)
-        return True
-
-    def sync(self, message: Sync):
-        self.logger.info("Postgres_Proxy: Syncing")
-        # TODO: Close/commit transaction if outside of a block. Maybe no collaries since Proxy
-        self.send_ready()
-        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.debug(self.__dict__)
-
-        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()
-        self.session.database = config.get("default_project")
-
-        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()
-
-        self.socket = self.request
-
-        self.current_transaction = None
-
-        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)
-        self.logger.debug("processing query\n%s", sql)
-        try:
-            executor.query_execute(sql)
-        except Exception as e:
-            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()),
-            )
-        return self.return_executor_data(executor)
-
-    def return_executor_data(self, executor):
-        if executor.data is None:
-            resp = SQLAnswer(
-                resp_type=RESPONSE_TYPE.OK,
-                state_track=executor.state_track,
-            )
-        else:
-            resp = SQLAnswer(
-                resp_type=RESPONSE_TYPE.TABLE,
-                state_track=executor.state_track,
-                result_set=executor.data,
-                status=executor.server_status,
-            )
-        return resp
-
-    def start_connection(self):
-        self.logger.debug("starting handshake")
-        self.handshake()
-        self.logger.debug("handshake complete, checking authentication")
-        return self.authenticate()
-
-    def send(self, message: PostgresMessage):
-        self.logger.debug("Sending message of type %s" % message.__class__.__name__)
-        message.send(self.wfile)
-
-    def handshake(self):
-        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)
-        self.user_parameters = self.client_buffer.read_startup_message()
-
-    def authenticate(self, ask_for_password=False):
-        if ask_for_password:
-            self.send(AuthenticationClearTextPassword())
-            password = self.client_buffer.read_authentication(encoding=self.charset)
-        else:
-            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"]:
-            self.logger.debug("Authentication succeeded")
-            self.session.username = auth_data["username"]
-            self.session.auth = True
-            self.send(AuthenticationOk())
-            return True
-        else:
-            if not ask_for_password:  # try asking for password
-                return self.authenticate(ask_for_password=True)
-            self.logger.debug("Authentication failed")
-            self.send(ConnectionFailure(message="Authentication failed."))
-            return False
-
-    def terminate(self, message: Terminate) -> bool:
-        self.logger.info("Postgres_Proxy: Terminating")
-        return False
-
-    def get_encoding(self) -> str:
-        return self.user_parameters.get(b"user_encoding", None) or self.charset
-
-    def return_ok(self, sql, rows: int = 0):
-        command = self.get_command(sql)
-        if command == "BEGIN":
-            self.transaction_status = b"T"
-        if command == "COMMIT":
-            self.transaction_status = b"I"
-        if command == "CREATE":
-            command = "SELECT"
-        if command in ("INSERT", "DELETE", "UPDATE", "SELECT", "MOVE", "FETCH", "COPY"):
-            command = f"{command} {rows}"
-        self.send(CommandComplete(tag=command.encode(encoding=self.get_encoding())))
-        return True
-
-    def get_command(self, sql):
-        sql = self.stringify_sql(sql)
-        return sql.split(" ", 1)[0]
-
-    def stringify_sql(self, sql):
-        if type(sql) == bytes:
-            encoding = self.get_encoding()
-            sql: str = sql.decode(encoding)
-        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)
-        rows = self.to_postgres_rows(sql_answer.result_set)
-        if row_descs:
-            self.send(RowDescriptions(fields=fields))
-        self.send(DataRow(rows=rows))
-        encoding = self.get_encoding()
-        tag = ("SELECT %s" % str(len(rows))).encode(encoding)
-        self.send(CommandComplete(tag=tag))
-        return True
-
-    def return_error(self, sql_answer: SQLAnswer):
-        self.send(Error.from_answer(error_code=sql_answer.error_code, error_message=sql_answer.error_message))
-        return True
-
-    def query(self, message: Query) -> bool:
-        self.logger.debug("Postgres Proxy: Got query of:\n%s" % message.sql)
-        sql = message.get_parsed_sql()
-        sql_answer = self.process_query(sql)
-        self.respond_from_sql_answer(sql=sql, sql_answer=sql_answer)
-        self.send_ready()
-        return True
-
-    def respond_from_sql_answer(self, sql, sql_answer: SQLAnswer, row_descs=True) -> bool:
-        # TODO Add command complete passthrough for Complex Queries that exceed row limit in one go
-        rows = 0
-        if sql_answer.result_set:
-            rows = len(sql_answer.result_set)
-        if RESPONSE_TYPE.OK == sql_answer.type:
-            return self.return_ok(sql, rows=rows)
-        elif RESPONSE_TYPE.TABLE == sql_answer.type:
-            return self.return_table(sql_answer, row_descs=row_descs)
-        elif RESPONSE_TYPE.ERROR == sql_answer.type:
-            return self.return_error(sql_answer)
-
-    @staticmethod
-    def to_postgres_fields(columns: Iterable[Dict[str, Any]]) -> Sequence[PostgresField]:
-        fields = []
-        i = 0
-        for column in columns:
-            fields.append(GenericField(name=column["name"], object_id=column["type"].value, column_id=i))
-            i += 1
-        return fields
-
-    def to_postgres_rows(self, rows: Iterable[Iterable[Any]]) -> Sequence[Sequence[bytes]]:
-        p_rows = []
-        for row in rows:
-            p_row = []
-            for column in row:
-                if column is None:
-                    column = ""
-                elif type(column) == int or type(column) == float:
-                    column = str(column)
-                elif type(column) == list or type(column) == dict:
-                    column = json.dumps(column)
-                if isinstance(column, datetime.date) or isinstance(column, datetime.datetime):
-                    try:
-                        column = datetime.datetime.strftime(column, "%Y-%m-%d")
-                    except ValueError:
-                        try:
-                            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")
-                if isinstance(column, bool):
-                    if column:
-                        column = "true"
-                    else:
-                        column = "false"
-                p_row.append(column.encode(encoding=self.charset))
-            p_rows.append(p_row)
-        return p_rows
-
-    def send_initial_data(self):
-        server_encoding = self.charset.encode(self.charset)
-        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))
-        self.send(ParameterStatus(name=b"client_encoding", value=client_encoding))
-        # TODO Send Parameters to complete set on 55.2.7 Asynchronous Operations E.G. - At present there is a
-        #  hard-wired set of parameters for which ParameterStatus will be generated: they are server_version,
-        #  server_encoding, client_encoding, application_name, default_transaction_read_only, in_hot_standby,
-        #  is_superuser, session_authorization, DateStyle, IntervalStyle, TimeZone, integer_datetimes,
-        #  and standard_conforming_strings
-        return
-
-    def send_ready(self):
-        self.logger.debug("Ready for Query")
-        self.send(ReadyForQuery(transaction_status=self.transaction_status))
-
-    def main_loop(self):
-        self.send_ready()
-        while True:
-            message: PostgresMessage = self.client_buffer.read_message()
-            if message is None:  # Empty Data, Buffer done
-                break
-            tof = type(message)
-            if tof in self.message_map:
-                res = self.message_map[tof](message)
-                if not res:
-                    break
-            else:
-                self.logger.warning("Ignoring unsupported message type %s" % tof)
-
-    @staticmethod
-    def startProxy():
-        host = config["api"]["postgres"]["host"]
-        port = int(config["api"]["postgres"]["port"])
-        server = TcpServer((host, port), PostgresProxyHandler)
-        server.connection_id = 0
-        server.mindsdb_config = config
-        server.check_auth = partial(check_auth, config=config)
-        server.serve_forever()
-
-
-class TcpServer(socketserver.ThreadingTCPServer):
-    allow_reuse_address = True
-
-
-if __name__ == "__main__":
-    PostgresProxyHandler.startProxy()

mindsdb/api/postgres/postgres_proxy/utilities/__init__.py

@@ -1,10 +0,0 @@
-import re
-from typing import Union
-
-
-def strip_null_byte(x: Union[str, bytes], encoding=None):
-    if type(x) == bytes:
-        if encoding is None:
-            encoding = "UTF-8"
-        x = x.decode(encoding=encoding)
-    return re.sub(r'[\s\x00]+$', '', x)

mindsdb/api/postgres/start.py

@@ -1,11 +0,0 @@
-import mindsdb.interfaces.storage.db as db
-from mindsdb.api.postgres.postgres_proxy.postgres_proxy import PostgresProxyHandler
-from mindsdb.utilities import log
-
-
-def start(verbose=False):
-    logger = log.getLogger(__name__)
-    logger.info("Postgres API is starting..")
-    db.init()
-
-    PostgresProxyHandler.startProxy()

mindsdb/integrations/handlers/mssql_handler/tests/test_mssql_handler.py

@@ -1,169 +0,0 @@
-import os
-import pytest
-import pymssql
-import pandas as pd
-from mindsdb.integrations.handlers.mssql_handler.mssql_handler import SqlServerHandler
-from mindsdb.integrations.libs.response import (
-    HandlerResponse as Response,
-    RESPONSE_TYPE,
-)
-
-HANDLER_KWARGS = {
-    "connection_data": {
-        "host": os.environ.get("MDB_TEST_MSSQL_HOST", "localhost"),
-        "port": os.environ.get("MDB_TEST_MSSQL_PORT", "1433"),
-        "user": os.environ.get("MDB_TEST_MSSQL_USER", "sa"),
-        "password": os.environ.get("MDB_TEST_MSSQL_PASSWORD", "admin5678@"),
-        "database": os.environ.get("MDB_TEST_MSSQL_DATABASE", "mdb_db_handler_test"),
-    }
-}
-
-
-@pytest.fixture(scope="class")
-def sql_server_handler():
-    seed_db()
-    handler = SqlServerHandler("test_sqlserver_handler", **HANDLER_KWARGS)
-    yield handler
-    handler.disconnect()
-
-
-def seed_db():
-    """Seed the test DB with some data"""
-
-    # Connect to 'master' while we create our test DB
-    conn_info = HANDLER_KWARGS["connection_data"].copy()
-    conn_info["database"] = "master"
-    db = pymssql.connect(**conn_info, autocommit=True)
-    cursor = db.cursor()
-
-    with open("mindsdb/integrations/handlers/mssql_handler/tests/seed.sql", "r") as f:
-        for line in f.readlines():
-            cursor.execute(line)
-    cursor.close()
-    db.close()
-
-
-def check_valid_response(res):
-    if res.resp_type == RESPONSE_TYPE.TABLE:
-        assert res.data_frame is not None, "expected to have some data, but got None"
-    assert (
-        res.error_code == 0
-    ), f"expected to have zero error_code, but got {res.error_code}"
-    assert (
-        res.error_message is None
-    ), f"expected to have None in error message, but got {res.error_message}"
-
-
-def get_table_names(sql_server_handler):
-    res = sql_server_handler.get_tables()
-    tables = res.data_frame
-    assert tables is not None, "expected to have some tables in the db, but got None"
-    assert (
-        "table_name" in tables
-    ), f"expected to get 'table_name' column in the response:\n{tables}"
-    return list(tables["table_name"])
-
-
-@pytest.mark.usefixtures("sql_server_handler")
-class TestMSSQLHandlerConnect:
-    def test_connect(self, sql_server_handler):
-        sql_server_handler.connect()
-        assert sql_server_handler.is_connected, "the handler has failed to connect"
-
-    def test_check_connection(self, sql_server_handler):
-        res = sql_server_handler.check_connection()
-        assert res.success, res.error_message
-
-
-class TestMSSQLHandlerDisconnect:
-    def test_disconnect(self, sql_server_handler):
-        sql_server_handler.disconnect()
-        assert sql_server_handler.is_connected is False, "failed to disconnect"
-
-    def test_check_connection(self, sql_server_handler):
-        res = sql_server_handler.check_connection()
-        assert res.success, res.error_message
-
-
-@pytest.mark.usefixtures("sql_server_handler")
-class TestMSSQLHandlerTables:
-    table_for_creation = "test_mdb"
-
-    def test_get_tables(self, sql_server_handler):
-        res = sql_server_handler.get_tables()
-        tables = res.data_frame
-        assert (
-            tables is not None
-        ), "expected to have some tables in the db, but got None"
-        assert (
-            "table_name" in tables
-        ), f"expected to get 'table_name' in the response but got: {tables}"
-        assert (
-            "test" in tables["table_name"].values
-        ), "expected to have 'test' in the response."
-
-    def test_get_columns(self, sql_server_handler):
-        response = sql_server_handler.get_columns("test")
-        assert response.type == RESPONSE_TYPE.TABLE, "expected a TABLE"
-        assert len(response.data_frame) > 0, "expected > O columns"
-
-        expected_columns = {
-            "Field": ["col_one", "col_two", "col_three", "col_four"],
-            "Type": ["int", "int", "float", "varchar"],
-        }
-        expected_df = pd.DataFrame(expected_columns)
-        assert response.data_frame.equals(
-            expected_df
-        ), "response does not contain the expected columns"
-
-    def test_create_table(self, sql_server_handler):
-        query = f"""
-        IF NOT EXISTS (SELECT * FROM sys.tables WHERE name = '{self.table_for_creation}')
-        BEGIN
-            CREATE TABLE {self.table_for_creation} (test_col INT)
-        END
-        """
-        res = sql_server_handler.native_query(query)
-        check_valid_response(res)
-        tables = get_table_names(sql_server_handler)
-        assert (
-            self.table_for_creation in tables
-        ), f"expected to have {self.table_for_creation} in database, but got: {tables}"
-
-    def test_drop_table(self, sql_server_handler):
-        query = f"DROP TABLE IF EXISTS {self.table_for_creation}"
-        res = sql_server_handler.native_query(query)
-        check_valid_response(res)
-        tables = get_table_names(sql_server_handler)
-        assert self.table_for_creation not in tables
-
-
-@pytest.mark.usefixtures("sql_server_handler")
-class TestMSSQLHandlerQuery:
-    def test_select_native_query(self, sql_server_handler):
-        query = "SELECT * FROM test"
-        response = sql_server_handler.native_query(query)
-        assert type(response) is Response
-        assert response.resp_type == RESPONSE_TYPE.TABLE
-
-        expected_data = {
-            "col_one": [1, 2, 3],
-            "col_two": [-1, -2, -3],
-            "col_three": [0.1, 0.2, 0.3],
-            "col_four": ["A", "B", "C"],
-        }
-        expected_df = pd.DataFrame(expected_data)
-        assert response.data_frame.equals(
-            expected_df
-        ), "response does not contain the expected data"
-
-    def test_select_query(self, sql_server_handler):
-        limit = 3
-        query = "SELECT * FROM test"
-        res = sql_server_handler.query(query)
-        check_valid_response(res)
-        got_rows = res.data_frame.shape[0]
-        want_rows = limit
-        assert (
-            got_rows == want_rows
-        ), f"expected to have {want_rows} rows in response but got: {got_rows}"