MindsDB 25.4.5.0__py3-none-any.whl → 25.5.4.0__py3-none-any.whl

mindsdb/integrations/handlers/paypal_handler/paypal_tables.py

@@ -104,7 +104,7 @@ class SubscriptionsTable(APITable):
         return subscriptions_df
 
     def get_columns(self) -> List[Text]:
-        return pd.json_normalize(self.get_subscriptions(count = 1)).columns.tolist()
+        return pd.json_normalize(self.get_subscriptions(count=1)).columns.tolist()
 
     def get_subscriptions(self, **kwargs) -> List[Dict]:
         connection = self.handler.connect()
@@ -114,6 +114,7 @@ class SubscriptionsTable(APITable):
 
 class OrdersTable(APITable):
     """The PayPal Orders Table implementation"""
+
     def select(self, query: ast.Select) -> pd.DataFrame:
         """
         Pulls PayPal Orders data.
@@ -137,18 +138,18 @@ class OrdersTable(APITable):
         )
         selected_columns, where_conditions, order_by_conditions, result_limit = select_statement_parser.parse_query()
 
-        id=None
+        id = None
         subset_where_conditions = []
         for op, arg1, arg2 in where_conditions:
             if arg1 == 'id':
                 if op == '=':
-                    id=arg2
+                    id = arg2
                 else:
                     raise NotImplementedError("Only '=' operator is supported for 'ids' column")
             elif arg1 in ['state', 'amount', 'create_time', 'update_time', 'links', 'pending_reason', 'parent_payment']:
                 subset_where_conditions.append([op, arg1, arg2])
 
-        if not id :
+        if not id:
             raise NotImplementedError("id column is required for this table")
 
         orders_df = pd.json_normalize(self.get_orders(id))
@@ -160,16 +161,17 @@ class OrdersTable(APITable):
         )
         orders_df = select_statement_executor.execute_query()
         return orders_df
+
     def get_columns(self) -> List[Text]:
-         return ["id",
-                 "status",
-                 "intent",
-                 "purchase_units",
-                 "links",
-                 "create_time"]
+        return ["id",
+                "status",
+                "intent",
+                "purchase_units",
+                "links",
+                "create_time"]
 
     # restore this or similar header list for API 2.0 refactor
-    #restore this list when restore paypalsdk api, and retired the request call
+    # restore this list when restore paypalsdk api, and retired the request call
         # return ["id",
         #         "status",
         #         "intent",
@@ -183,7 +185,7 @@ class OrdersTable(APITable):
         #         "create_time"]
 
     def get_orders(self, id) -> List[Dict]:
-        #we can use the paypalrestsdk api to get the order if they refactor their code
+        # we can use the paypalrestsdk api to get the order if they refactor their code
         connection = self.handler.connect()
         endpoint = f"v2/checkout/orders/{id}"
         order = connection.get(endpoint)
@@ -259,7 +261,7 @@ class PayoutsTable(APITable):
             "fees_value",
         ]
 
-    def get_payout(self, payout_batch_id:str) -> List[Dict]:
+    def get_payout(self, payout_batch_id: str) -> List[Dict]:
         connection = self.handler.connect()
         endpoint = f"v1/payments/payouts/{payout_batch_id}"
         payout = connection.get(endpoint)
@@ -280,4 +282,3 @@ class PayoutsTable(APITable):
         }
 
         return [payout_data]
-

mindsdb/integrations/handlers/pgvector_handler/pgvector_handler.py

@@ -1,7 +1,6 @@
 import os
 import json
-from enum import Enum
-from typing import Dict, List, Union
+from typing import Dict, List, Union, Literal
 from urllib.parse import urlparse
 
 import pandas as pd
@@ -39,7 +38,7 @@ class PgVectorHandler(PostgresHandler, VectorStoreHandler):
         self._is_vector_registered = False
         # we get these from the connection args on PostgresHandler parent
         self._is_sparse = self.connection_args.get('is_sparse', False)
-        self._vector_size = self.connection_args.get('vector_size', None) 
+        self._vector_size = self.connection_args.get('vector_size', None)
 
         if self._is_sparse:
             if not self._vector_size:
@@ -67,6 +66,21 @@ class PgVectorHandler(PostgresHandler, VectorStoreHandler):
         self.distance_op = distance_op
         self.connect()
 
+    def get_metric_type(self) -> str:
+        """
+        Get the metric type from the distance ops
+
+        """
+        distance_ops_to_metric_type_map = {
+            '<->': 'vector_l2_ops',
+            '<#>': 'vector_ip_ops',
+            '<=>': 'vector_cosine_ops',
+            '<+>': 'vector_l1_ops',
+            '<~>': 'bit_hamming_ops',
+            '<%>': 'bit_jaccard_ops'
+        }
+        return distance_ops_to_metric_type_map.get(self.distance_op, 'vector_cosine_ops')
+
     def _make_connection_args(self):
         cloud_pgvector_url = os.environ.get('KB_PGVECTOR_URL')
         # if no connection args and shared pg vector defined - use it
@@ -234,7 +248,6 @@ class PgVectorHandler(PostgresHandler, VectorStoreHandler):
 
         targets = ', '.join(modified_columns)
 
-
         if filter_conditions:
 
             if embedding_search:
@@ -255,7 +268,7 @@ class PgVectorHandler(PostgresHandler, VectorStoreHandler):
                 # Calculate distance as part of the query if needed
                 if has_distance:
                     targets = f"{targets}, (embeddings {self.distance_op} '{search_vector}') as distance"
-                
+
                 return f"SELECT {targets} FROM {table_name} {where_clause} ORDER BY embeddings {self.distance_op} '{search_vector}' ASC {limit_clause} {offset_clause} "
 
             else:
@@ -304,7 +317,7 @@ class PgVectorHandler(PostgresHandler, VectorStoreHandler):
         embeddings: List[float],
         query: str = None,
         metadata: Dict[str, str] = None,
-        distance_function = DistanceFunction.COSINE_DISTANCE,
+        distance_function=DistanceFunction.COSINE_DISTANCE,
         **kwargs
     ) -> pd.DataFrame:
         '''
@@ -348,7 +361,7 @@ class PgVectorHandler(PostgresHandler, VectorStoreHandler):
         # See https://docs.pgvecto.rs/use-case/hybrid-search.html#advanced-search-merge-the-results-of-full-text-search-and-vector-search.
         #
         # We can break down the below query as follows:
-        # 
+        #
         # Start with a CTE (Common Table Expression) called semantic_search (https://www.postgresql.org/docs/current/queries-with.html).
         # This expression calculates rank by the defined distance function, which measures the distance between the
         # embeddings column and the given embeddings vector. Results are ordered by this rank.
@@ -409,11 +422,11 @@ class PgVectorHandler(PostgresHandler, VectorStoreHandler):
         with self.connection.cursor() as cur:
             # For sparse vectors, use sparsevec type
             vector_column_type = 'sparsevec' if self._is_sparse else 'vector'
-            
+
             # Vector size is required for sparse vectors, optional for dense
             if self._is_sparse and not self._vector_size:
                 raise ValueError("vector_size is required for sparse vectors")
-            
+
             # Add vector size specification only if provided
             size_spec = f"({self._vector_size})" if self._vector_size is not None else "()"
             if vector_column_type == 'vector':
@@ -438,7 +451,7 @@ class PgVectorHandler(PostgresHandler, VectorStoreHandler):
         table_name = self._check_table(table_name)
 
         if 'metadata' in data.columns:
-             data['metadata'] = data['metadata'].apply(json.dumps)
+            data['metadata'] = data['metadata'].apply(json.dumps)
 
         resp = super().insert(table_name, data)
         if resp.resp_type == RESPONSE_TYPE.ERROR:
@@ -521,3 +534,33 @@ class PgVectorHandler(PostgresHandler, VectorStoreHandler):
         """
         table_name = self._check_table(table_name)
         self.raw_query(f"DROP TABLE IF EXISTS {table_name}")
+
+    def create_index(self, table_name: str, column_name: str = "embeddings", index_type: Literal['ivfflat', 'hnsw'] = "hnsw", metric_type: str = None):
+        """
+        Create an index on the pgvector table.
+        Args:
+            table_name (str): Name of the table to create the index on.
+            column_name (str): Name of the column to create the index on.
+            index_type (str): Type of the index to create. Supported types are 'ivfflat' and 'hnsw'.
+            metric_type (str): Metric type for the index. Supported types are 'vector_l2_ops', 'vector_ip_ops', and 'vector_cosine_ops'.
+        """
+        if metric_type is None:
+            metric_type = self.get_metric_type()
+        # Check if the index type is supported
+        if index_type not in ['ivfflat', 'hnsw']:
+            raise ValueError("Invalid index type. Supported types are 'ivfflat' and 'hnsw'.")
+        table_name = self._check_table(table_name)
+        # first we make sure embedding dimension is set
+        embedding_dim_size_df = self.raw_query(f"SELECT vector_dims({column_name}) FROM {table_name} LIMIT 1")
+        # check if answer is empty
+        if embedding_dim_size_df.empty:
+            raise ValueError("Could not determine embedding dimension size. Make sure that knowledge base isn't empty")
+        try:
+            embedding_dim = int(embedding_dim_size_df.iloc[0, 0])
+            # alter table to add dimension
+            self.raw_query(f"ALTER TABLE {table_name} ALTER COLUMN {column_name} TYPE vector({embedding_dim})")
+        except Exception:
+            raise ValueError("Could not determine embedding dimension size. Make sure that knowledge base isn't empty")
+
+        # Create the index
+        self.raw_query(f"CREATE INDEX ON {table_name} USING {index_type} ({column_name} {metric_type})")

mindsdb/integrations/handlers/phoenix_handler/__init__.py

@@ -17,4 +17,4 @@ icon_path = 'icon.png'
 __all__ = [
     'Handler', 'version', 'name', 'type', 'title', 'description',
     'connection_args', 'connection_args_example', 'import_error', 'icon_path'
-]
+]

mindsdb/integrations/handlers/phoenix_handler/phoenix_handler.py

@@ -20,6 +20,7 @@ from mindsdb.integrations.libs.response import (
 
 logger = log.getLogger(__name__)
 
+
 class PhoenixHandler(DatabaseHandler):
     """
     This handler handles connection and execution of the Apache Phoenix statements.

mindsdb/integrations/handlers/pinot_handler/__init__.py

@@ -17,4 +17,4 @@ icon_path = 'icon.svg'
 __all__ = [
     'Handler', 'version', 'name', 'type', 'title', 'description',
     'connection_args', 'connection_args_example', 'import_error', 'icon_path'
-]
+]

mindsdb/integrations/handlers/pinot_handler/pinot_handler.py

@@ -23,6 +23,7 @@ from mindsdb.integrations.libs.response import (
 
 logger = log.getLogger(__name__)
 
+
 class PinotHandler(DatabaseHandler):
     """
     This handler handles connection and execution of the Apache Pinot statements.
@@ -182,7 +183,7 @@ class PinotHandler(DatabaseHandler):
         api_url = f"{self.connection_data['host']}:{self.connection_data['controller_port']}/tables"
         try:
             result = requests.get(api_url)
-        except InvalidSchema as e:
+        except InvalidSchema:
             api_url = f"{self.connection_data['scheme']}://{api_url}"
             result = requests.get(api_url)
 
@@ -208,7 +209,7 @@ class PinotHandler(DatabaseHandler):
         api_url = f"{self.connection_data['host']}:{self.connection_data['controller_port']}/tables/{table_name}/schema"
         try:
             result = requests.get(api_url)
-        except InvalidSchema as e:
+        except InvalidSchema:
             api_url = f"{self.connection_data['scheme']}://{api_url}"
             result = requests.get(api_url)
 

mindsdb/integrations/handlers/plaid_handler/plaid_handler.py

@@ -1,6 +1,5 @@
 import pandas as pd
 from mindsdb.utilities import log
-from mindsdb_sql_parser import parse_sql
 from mindsdb.integrations.libs.api_handler import APIHandler, FuncParser
 from mindsdb.integrations.libs.response import (
     HandlerStatusResponse as StatusResponse,
@@ -27,12 +26,13 @@ PLAID_ENV = {
 
 logger = log.getLogger(__name__)
 
+
 class PlaidHandler(APIHandler):
     '''A class for handling connections and interactions with the Plaid API.
 
     Attributes:
         plaid_env (str): Enviroment used by user [ 'sandbox'(default) OR 'development' OR 'production' ].
-        client_id (str): Your Plaid API client_id. 
+        client_id (str): Your Plaid API client_id.
         secret (str): Your Plaid API secret
         access_token (str): The access token for the Plaid account.
     '''
@@ -119,7 +119,7 @@ class PlaidHandler(APIHandler):
             DataFrame
         '''
 
-        result=pd.DataFrame()
+        result = pd.DataFrame()
         if method_name == 'get_balance':
             result = self.get_balance(params=params)
             result = BalanceTable(self).filter_columns(result=result)
@@ -143,9 +143,9 @@ class PlaidHandler(APIHandler):
         self.connect()
         if params.get('last_updated_datetime') is not None:
             options = AccountsBalanceGetRequestOptions(
-                min_last_updated_datetime=datetime.strptime( 
+                min_last_updated_datetime=datetime.strptime(
                     params.get('last_updated_datetime')
-                    )
+                )
             )
 
             response = self.api.accounts_balance_get(
@@ -192,7 +192,7 @@ class PlaidHandler(APIHandler):
             end_date = datetime.strptime(params.get('end_date'), '%Y-%m-%d').date()
         else:
             raise Exception('start_date and end_date is required in format YYYY-MM-DD ')
-      
+
         request = TransactionsGetRequest(
             access_token=self.access_token,
             start_date=start_date,
@@ -207,17 +207,17 @@ class PlaidHandler(APIHandler):
         # transactions and retrieve all available data
         while len(transactions) < response['total_transactions']:
             request = TransactionsGetRequest(
-                    access_token=self.access_token,
-                    start_date=start_date,
-                    end_date=end_date,
-                    options=TransactionsGetRequestOptions(
-                        offset=len(transactions)
-                    )
+                access_token=self.access_token,
+                start_date=start_date,
+                end_date=end_date,
+                options=TransactionsGetRequestOptions(
+                    offset=len(transactions)
+                )
             )
             response = self.api.transactions_get(request)
             transactions.extend(parse_transaction(response['transactions']))
 
-        # Converting date column from str 
+        # Converting date column from str
         df = pd.DataFrame(transactions)
         for i in ['date', 'authorized_date']:
             df[i] = pd.to_datetime(df[i]).dt.date

mindsdb/integrations/handlers/plaid_handler/plaid_tables.py

@@ -1,5 +1,5 @@
 import pandas as pd
-from mindsdb.integrations.libs.api_handler import  APITable
+from mindsdb.integrations.libs.api_handler import APITable
 from mindsdb.integrations.utilities.sql_utils import extract_comparison_conditions
 from mindsdb_sql_parser import ast
 
@@ -15,8 +15,8 @@ class BalanceTable(APITable):
         get_columns(): Get the list of column names for the balance table.
 
     '''
-    
-    def select(self, query: ast.Select) :
+
+    def select(self, query: ast.Select):
         '''Select data from the balance table and return it as a pandas DataFrame.
 
         Args:
@@ -34,21 +34,21 @@ class BalanceTable(APITable):
                     params[i[1]] = i[2]
                 else:
                     raise Exception("Only equals to '=' is Supported with 'last_updated_datetime'")
-                
+
         result = self.handler.call_plaid_api(method_name='get_balance', params=params)
 
         self.filter_columns(query=query, result=result)
         return result
-    
+
     def get_columns(self):
         '''Get the list of column names for the balance table.
-          
+
         Returns:
             list: A list of column names for the balance table.
 
         '''
         return [
-            'account_id', 
+            'account_id',
             'account_name',
             'account_mask',
             'account_type',
@@ -58,11 +58,11 @@ class BalanceTable(APITable):
             'balance_unofficial_currency_code',
             'balance_available',
             'balance_current',
-            'balance_limit' 
+            'balance_limit'
         ]
 
     def filter_columns(self, result: pd.DataFrame, query: ast.Select = None):
-        
+
         columns = []
         if query is not None:
             for target in query.targets:
@@ -103,7 +103,7 @@ class TransactionTable(BalanceTable):
 
     '''
 
-    def select(self, query: ast.Select) :
+    def select(self, query: ast.Select):
         '''Select data from the transaction table and return it as a pandas DataFrame.
 
         Args:
@@ -155,5 +155,3 @@ class TransactionTable(BalanceTable):
             'payment_channel',
             'pending',
         ]
-
-    

mindsdb/integrations/handlers/plaid_handler/utils.py

@@ -1,9 +1,7 @@
-def parse_transaction(res:list):
-    parsed=[]
+def parse_transaction(res: list):
+    parsed = []
     for dic in res:
-        dic=dic.to_dict()
+        dic = dic.to_dict()
         parsed.append(dic)
-        
-    return parsed
-
 
+    return parsed

mindsdb/integrations/handlers/planetscale_handler/planetscale_handler.py

@@ -1,6 +1,3 @@
-from collections import OrderedDict
-
-from mindsdb.integrations.libs.const import HANDLER_CONNECTION_ARG_TYPE as ARG_TYPE
 from mindsdb.integrations.handlers.mysql_handler import Handler as MySQLHandler
 
 
@@ -9,6 +6,6 @@ class PlanetScaleHandler(MySQLHandler):
     This handler handles the connection and execution of queries against PlanetScale.
     """
     name = 'planet_scale'
-    
+
     def __init__(self, name, **kwargs):
         super().__init__(name, **kwargs)

mindsdb/integrations/handlers/portkey_handler/__init__.py

@@ -1,10 +1,10 @@
+from .__about__ import __version__ as version
+from .__about__ import __description__ as description
 from mindsdb.integrations.libs.const import HANDLER_TYPE
 from mindsdb.utilities import log
 
 logger = log.getLogger(__name__)
 
-from .__about__ import __description__ as description
-from .__about__ import __version__ as version
 
 try:
     from .portkey_handler import PortkeyHandler as Handler

mindsdb/integrations/handlers/postgres_handler/postgres_handler.py

@@ -1,12 +1,13 @@
 import time
 import json
-from typing import Optional
+from typing import Optional, Any
 
 import pandas as pd
+from pandas import DataFrame
 import psycopg
-from psycopg.postgres import types
+from psycopg import Column as PGColumn, Cursor
+from psycopg.postgres import TypeInfo, types as pg_types
 from psycopg.pq import ExecStatus
-from pandas import DataFrame
 
 from mindsdb_sql_parser import parse_sql
 from mindsdb.utilities.render.sqlalchemy_render import SqlalchemyRender
@@ -27,7 +28,7 @@ logger = log.getLogger(__name__)
 SUBSCRIBE_SLEEP_INTERVAL = 1
 
 
-def _map_type(internal_type_name: str) -> MYSQL_DATA_TYPE:
+def _map_type(internal_type_name: str | None) -> MYSQL_DATA_TYPE:
     """Map Postgres types to MySQL types.
 
     Args:
@@ -36,14 +37,22 @@ def _map_type(internal_type_name: str) -> MYSQL_DATA_TYPE:
     Returns:
         MYSQL_DATA_TYPE: The MySQL type that corresponds to the Postgres type.
     """
+    fallback_type = MYSQL_DATA_TYPE.VARCHAR
+
+    if internal_type_name is None:
+        return fallback_type
+
     internal_type_name = internal_type_name.lower()
     types_map = {
-        ('smallint', 'integer', 'bigint', 'int', 'smallserial', 'serial', 'bigserial'): MYSQL_DATA_TYPE.INT,
-        ('real', 'money', 'float'): MYSQL_DATA_TYPE.FLOAT,
+        ('smallint', 'smallserial'): MYSQL_DATA_TYPE.SMALLINT,
+        ('integer', 'int', 'serial'): MYSQL_DATA_TYPE.INT,
+        ('bigint', 'bigserial'): MYSQL_DATA_TYPE.BIGINT,
+        ('real', 'float'): MYSQL_DATA_TYPE.FLOAT,
         ('numeric', 'decimal'): MYSQL_DATA_TYPE.DECIMAL,
         ('double precision',): MYSQL_DATA_TYPE.DOUBLE,
         ('character varying', 'varchar'): MYSQL_DATA_TYPE.VARCHAR,
-        ('character', 'char', 'bpchar', 'bpchar', 'text'): MYSQL_DATA_TYPE.TEXT,
+        # NOTE: if return chars-types as mysql's CHAR, then response will be padded with spaces, so return as TEXT
+        ('money', 'character', 'char', 'bpchar', 'bpchar', 'text'): MYSQL_DATA_TYPE.TEXT,
         ('timestamp', 'timestamp without time zone', 'timestamp with time zone'): MYSQL_DATA_TYPE.DATETIME,
         ('date', ): MYSQL_DATA_TYPE.DATE,
         ('time', 'time without time zone', 'time with time zone'): MYSQL_DATA_TYPE.TIME,
@@ -55,8 +64,51 @@ def _map_type(internal_type_name: str) -> MYSQL_DATA_TYPE:
         if internal_type_name in db_types_list:
             return mysql_data_type
 
-    logger.warning(f"Postgres handler type mapping: unknown type: {internal_type_name}, use VARCHAR as fallback.")
-    return MYSQL_DATA_TYPE.VARCHAR
+    logger.debug(f"Postgres handler type mapping: unknown type: {internal_type_name}, use VARCHAR as fallback.")
+    return fallback_type
+
+
+def _make_table_response(result: list[tuple[Any]], cursor: Cursor) -> Response:
+    """Build response from result and cursor.
+
+    Args:
+        result (list[tuple[Any]]): result of the query.
+        cursor (psycopg.Cursor): cursor object.
+
+    Returns:
+        Response: response object.
+    """
+    description: list[PGColumn] = cursor.description
+    mysql_types: list[MYSQL_DATA_TYPE] = []
+    for column in description:
+        pg_type_info: TypeInfo = pg_types.get(column.type_code)
+        if pg_type_info is None:
+            logger.warning(f'Postgres handler: unknown type: {column.type_code}')
+        regtype: str = pg_type_info.regtype if pg_type_info is not None else None
+        mysql_type = _map_type(regtype)
+        mysql_types.append(mysql_type)
+
+    # region cast int and bool to nullable types
+    serieses = []
+    for i, mysql_type in enumerate(mysql_types):
+        expected_dtype = None
+        if mysql_type in (
+            MYSQL_DATA_TYPE.SMALLINT, MYSQL_DATA_TYPE.INT, MYSQL_DATA_TYPE.MEDIUMINT,
+            MYSQL_DATA_TYPE.BIGINT, MYSQL_DATA_TYPE.TINYINT
+        ):
+            expected_dtype = 'Int64'
+        elif mysql_type in (MYSQL_DATA_TYPE.BOOL, MYSQL_DATA_TYPE.BOOLEAN):
+            expected_dtype = 'boolean'
+        serieses.append(pd.Series([row[i] for row in result], dtype=expected_dtype, name=description[i].name))
+    df = pd.concat(serieses, axis=1, copy=False)
+    # endregion
+
+    return Response(
+        RESPONSE_TYPE.TABLE,
+        data_frame=df,
+        affected_rows=cursor.rowcount,
+        mysql_types=mysql_types
+    )
 
 
 class PostgresHandler(DatabaseHandler):
@@ -199,13 +251,13 @@ class PostgresHandler(DatabaseHandler):
         for column_index, column_name in enumerate(df.columns):
             col = df[column_name]
             if str(col.dtype) == 'object':
-                pg_type = types.get(description[column_index].type_code)
-                if pg_type is not None and pg_type.name in types_map:
-                    col = col.fillna(0)
+                pg_type_info: TypeInfo = pg_types.get(description[column_index].type_code)        # type_code is int!?
+                if pg_type_info is not None and pg_type_info.name in types_map:
+                    col = col.fillna(0)   # TODO rework
                     try:
-                        df[column_name] = col.astype(types_map[pg_type.name])
+                        df[column_name] = col.astype(types_map[pg_type_info.name])
                     except ValueError as e:
-                        logger.error(f'Error casting column {col.name} to {types_map[pg_type.name]}: {e}')
+                        logger.error(f'Error casting column {col.name} to {types_map[pg_type_info.name]}: {e}')
         df.columns = columns
 
     @profiler.profile()
@@ -232,16 +284,7 @@ class PostgresHandler(DatabaseHandler):
                     response = Response(RESPONSE_TYPE.OK, affected_rows=cur.rowcount)
                 else:
                     result = cur.fetchall()
-                    df = DataFrame(
-                        result,
-                        columns=[x.name for x in cur.description]
-                    )
-                    self._cast_dtypes(df, cur.description)
-                    response = Response(
-                        RESPONSE_TYPE.TABLE,
-                        data_frame=df,
-                        affected_rows=cur.rowcount
-                    )
+                    response = _make_table_response(result, cur)
                 connection.commit()
             except Exception as e:
                 logger.error(f'Error running query: {query} on {self.database}, {e}!')
@@ -257,6 +300,44 @@ class PostgresHandler(DatabaseHandler):
 
         return response
 
+    def query_stream(self, query: ASTNode, fetch_size: int = 1000):
+        """
+        Executes a SQL query and stream results outside by batches
+
+        :param query: An ASTNode representing the SQL query to be executed.
+        :param fetch_size: size of the batch
+        :return: generator with query results
+        """
+        query_str, params = self.renderer.get_exec_params(query, with_failback=True)
+
+        need_to_close = not self.is_connected
+
+        connection = self.connect()
+        with connection.cursor() as cur:
+            try:
+                if params is not None:
+                    cur.executemany(query_str, params)
+                else:
+                    cur.execute(query_str)
+
+                if cur.pgresult is not None and ExecStatus(cur.pgresult.status) != ExecStatus.COMMAND_OK:
+                    while True:
+                        result = cur.fetchmany(fetch_size)
+                        if not result:
+                            break
+                        df = DataFrame(
+                            result,
+                            columns=[x.name for x in cur.description]
+                        )
+                        self._cast_dtypes(df, cur.description)
+                        yield df
+                connection.commit()
+            finally:
+                connection.rollback()
+
+        if need_to_close:
+            self.disconnect()
+
     def insert(self, table_name: str, df: pd.DataFrame) -> Response:
         need_to_close = not self.is_connected