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

mindsdb/integrations/handlers/postgres_handler/postgres_handler.py

@@ -1,3 +1,5 @@
+import csv
+import io
 import time
 import json
 from typing import Optional, Any
@@ -13,12 +15,12 @@ from mindsdb_sql_parser import parse_sql
 from mindsdb.utilities.render.sqlalchemy_render import SqlalchemyRender
 from mindsdb_sql_parser.ast.base import ASTNode
 
-from mindsdb.integrations.libs.base import DatabaseHandler
+from mindsdb.integrations.libs.base import MetaDatabaseHandler
 from mindsdb.utilities import log
 from mindsdb.integrations.libs.response import (
     HandlerStatusResponse as StatusResponse,
     HandlerResponse as Response,
-    RESPONSE_TYPE
+    RESPONSE_TYPE,
 )
 import mindsdb.utilities.profiler as profiler
 from mindsdb.api.mysql.mysql_proxy.libs.constants.mysql import MYSQL_DATA_TYPE
@@ -44,20 +46,21 @@ def _map_type(internal_type_name: str | None) -> MYSQL_DATA_TYPE:
 
     internal_type_name = internal_type_name.lower()
     types_map = {
-        ('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,
+        ("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,
         # 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,
-        ('boolean',): MYSQL_DATA_TYPE.BOOL,
-        ('bytea',): MYSQL_DATA_TYPE.BINARY,
+        ("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,
+        ("boolean",): MYSQL_DATA_TYPE.BOOL,
+        ("bytea",): MYSQL_DATA_TYPE.BINARY,
+        ("json", "jsonb"): MYSQL_DATA_TYPE.JSON,
     }
 
     for db_types_list, mysql_data_type in types_map.items():
@@ -81,10 +84,28 @@ def _make_table_response(result: list[tuple[Any]], cursor: Cursor) -> Response:
     description: list[PGColumn] = cursor.description
     mysql_types: list[MYSQL_DATA_TYPE] = []
     for column in description:
+        if column.type_display == "vector":
+            # 'vector' is type of pgvector extension, added here as text to not import pgvector
+            # NOTE: data returned as numpy array
+            mysql_types.append(MYSQL_DATA_TYPE.VECTOR)
+            continue
         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
+            # postgres may return 'polymorphic type', which are not present in the pg_types
+            # list of 'polymorphic type' can be obtained:
+            # SELECT oid, typname, typcategory FROM pg_type WHERE typcategory = 'P' ORDER BY oid;
+            if column.type_code in (2277, 5078):
+                # anyarray, anycompatiblearray
+                regtype = "json"
+            else:
+                logger.warning(f"Postgres handler: unknown type: {column.type_code}")
+                mysql_types.append(MYSQL_DATA_TYPE.TEXT)
+                continue
+        elif pg_type_info.array_oid == column.type_code:
+            # it is any array, handle is as json
+            regtype: str = "json"
+        else:
+            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)
 
@@ -93,38 +114,37 @@ def _make_table_response(result: list[tuple[Any]], cursor: Cursor) -> Response:
     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
+            MYSQL_DATA_TYPE.SMALLINT,
+            MYSQL_DATA_TYPE.INT,
+            MYSQL_DATA_TYPE.MEDIUMINT,
+            MYSQL_DATA_TYPE.BIGINT,
+            MYSQL_DATA_TYPE.TINYINT,
         ):
-            expected_dtype = 'Int64'
+            expected_dtype = "Int64"
         elif mysql_type in (MYSQL_DATA_TYPE.BOOL, MYSQL_DATA_TYPE.BOOLEAN):
-            expected_dtype = '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
-    )
+    return Response(RESPONSE_TYPE.TABLE, data_frame=df, affected_rows=cursor.rowcount, mysql_types=mysql_types)
 
 
-class PostgresHandler(DatabaseHandler):
+class PostgresHandler(MetaDatabaseHandler):
     """
     This handler handles connection and execution of the PostgreSQL statements.
     """
-    name = 'postgres'
 
-    @profiler.profile('init_pg_handler')
+    name = "postgres"
+
+    @profiler.profile("init_pg_handler")
     def __init__(self, name=None, **kwargs):
         super().__init__(name)
         self.parser = parse_sql
-        self.connection_args = kwargs.get('connection_data')
-        self.dialect = 'postgresql'
-        self.database = self.connection_args.get('database')
-        self.renderer = SqlalchemyRender('postgres')
+        self.connection_args = kwargs.get("connection_data")
+        self.dialect = "postgresql"
+        self.database = self.connection_args.get("database")
+        self.renderer = SqlalchemyRender("postgres")
 
         self.connection = None
         self.is_connected = False
@@ -136,30 +156,30 @@ class PostgresHandler(DatabaseHandler):
 
     def _make_connection_args(self):
         config = {
-            'host': self.connection_args.get('host'),
-            'port': self.connection_args.get('port'),
-            'user': self.connection_args.get('user'),
-            'password': self.connection_args.get('password'),
-            'dbname': self.connection_args.get('database')
+            "host": self.connection_args.get("host"),
+            "port": self.connection_args.get("port"),
+            "user": self.connection_args.get("user"),
+            "password": self.connection_args.get("password"),
+            "dbname": self.connection_args.get("database"),
         }
 
         # https://www.postgresql.org/docs/current/libpq-connect.html#LIBPQ-PARAMKEYWORDS
-        connection_parameters = self.connection_args.get('connection_parameters')
+        connection_parameters = self.connection_args.get("connection_parameters")
         if isinstance(connection_parameters, dict) is False:
             connection_parameters = {}
-        if 'connect_timeout' not in connection_parameters:
-            connection_parameters['connect_timeout'] = 10
+        if "connect_timeout" not in connection_parameters:
+            connection_parameters["connect_timeout"] = 10
         config.update(connection_parameters)
 
-        if self.connection_args.get('sslmode'):
-            config['sslmode'] = self.connection_args.get('sslmode')
+        if self.connection_args.get("sslmode"):
+            config["sslmode"] = self.connection_args.get("sslmode")
 
-        if self.connection_args.get('autocommit'):
-            config['autocommit'] = self.connection_args.get('autocommit')
+        if self.connection_args.get("autocommit"):
+            config["autocommit"] = self.connection_args.get("autocommit")
 
         # If schema is not provided set public as default one
-        if self.connection_args.get('schema'):
-            config['options'] = f'-c search_path={self.connection_args.get("schema")},public'
+        if self.connection_args.get("schema"):
+            config["options"] = f"-c search_path={self.connection_args.get('schema')},public"
         return config
 
     @profiler.profile()
@@ -182,7 +202,7 @@ class PostgresHandler(DatabaseHandler):
             self.is_connected = True
             return self.connection
         except psycopg.Error as e:
-            logger.error(f'Error connecting to PostgreSQL {self.database}, {e}!')
+            logger.error(f"Error connecting to PostgreSQL {self.database}, {e}!")
             self.is_connected = False
             raise
 
@@ -209,10 +229,10 @@ class PostgresHandler(DatabaseHandler):
             connection = self.connect()
             with connection.cursor() as cur:
                 # Execute a simple query to test the connection
-                cur.execute('select 1;')
+                cur.execute("select 1;")
             response.success = True
         except psycopg.Error as e:
-            logger.error(f'Error connecting to PostgreSQL {self.database}, {e}!')
+            logger.error(f"Error connecting to PostgreSQL {self.database}, {e}!")
             response.error_message = str(e)
 
         if response.success and need_to_close:
@@ -239,25 +259,25 @@ class PostgresHandler(DatabaseHandler):
                 description (list): psycopg cursor description
         """
         types_map = {
-            'int2': 'int16',
-            'int4': 'int32',
-            'int8': 'int64',
-            'numeric': 'float64',
-            'float4': 'float32',
-            'float8': 'float64'
+            "int2": "int16",
+            "int4": "int32",
+            "int8": "int64",
+            "numeric": "float64",
+            "float4": "float32",
+            "float8": "float64",
         }
         columns = df.columns
         df.columns = list(range(len(columns)))
         for column_index, column_name in enumerate(df.columns):
             col = df[column_name]
-            if str(col.dtype) == 'object':
-                pg_type_info: TypeInfo = pg_types.get(description[column_index].type_code)        # type_code is int!?
+            if str(col.dtype) == "object":
+                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
+                    col = col.fillna(0)  # TODO rework
                     try:
                         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_info.name]}: {e}')
+                        logger.error(f"Error casting column {col.name} to {types_map[pg_type_info.name]}: {e}")
         df.columns = columns
 
     @profiler.profile()
@@ -287,12 +307,8 @@ class PostgresHandler(DatabaseHandler):
                     response = _make_table_response(result, cur)
                 connection.commit()
             except Exception as e:
-                logger.error(f'Error running query: {query} on {self.database}, {e}!')
-                response = Response(
-                    RESPONSE_TYPE.ERROR,
-                    error_code=0,
-                    error_message=str(e)
-                )
+                logger.error(f"Error running query: {query} on {self.database}, {e}!")
+                response = Response(RESPONSE_TYPE.ERROR, error_code=0, error_message=str(e))
                 connection.rollback()
 
         if need_to_close:
@@ -325,10 +341,7 @@ class PostgresHandler(DatabaseHandler):
                         result = cur.fetchmany(fetch_size)
                         if not result:
                             break
-                        df = DataFrame(
-                            result,
-                            columns=[x.name for x in cur.description]
-                        )
+                        df = DataFrame(result, columns=[x.name for x in cur.description])
                         self._cast_dtypes(df, cur.description)
                         yield df
                 connection.commit()
@@ -349,16 +362,10 @@ class PostgresHandler(DatabaseHandler):
 
         # copy requires precise cases of names: get current column names from table and adapt input dataframe columns
         if resp.data_frame is not None and not resp.data_frame.empty:
-            db_columns = {
-                c.lower(): c
-                for c in resp.data_frame['COLUMN_NAME']
-            }
+            db_columns = {c.lower(): c for c in resp.data_frame["COLUMN_NAME"]}
 
             # try to get case of existing column
-            columns = [
-                db_columns.get(c.lower(), c)
-                for c in columns
-            ]
+            columns = [db_columns.get(c.lower(), c) for c in columns]
 
         columns = [f'"{c}"' for c in columns]
         rowcount = None
@@ -370,7 +377,7 @@ class PostgresHandler(DatabaseHandler):
 
                 connection.commit()
             except Exception as e:
-                logger.error(f'Error running insert to {table_name} on {self.database}, {e}!')
+                logger.error(f"Error running insert to {table_name} on {self.database}, {e}!")
                 connection.rollback()
                 raise e
             rowcount = cur.rowcount
@@ -402,9 +409,9 @@ class PostgresHandler(DatabaseHandler):
         Returns:
             Response: A response object containing the list of tables and views, formatted as per the `Response` class.
         """
-        all_filter = 'and table_schema = current_schema()'
+        all_filter = "and table_schema = current_schema()"
         if all is True:
-            all_filter = ''
+            all_filter = ""
         query = f"""
             SELECT
                 table_schema,
@@ -439,7 +446,7 @@ class PostgresHandler(DatabaseHandler):
         if isinstance(schema_name, str):
             schema_name = f"'{schema_name}'"
         else:
-            schema_name = 'current_schema()'
+            schema_name = "current_schema()"
         query = f"""
             SELECT
                 COLUMN_NAME,
@@ -467,33 +474,33 @@ class PostgresHandler(DatabaseHandler):
 
     def subscribe(self, stop_event, callback, table_name, columns=None, **kwargs):
         config = self._make_connection_args()
-        config['autocommit'] = True
+        config["autocommit"] = True
 
         conn = psycopg.connect(connect_timeout=10, **config)
 
         # create db trigger
-        trigger_name = f'mdb_notify_{table_name}'
+        trigger_name = f"mdb_notify_{table_name}"
 
-        before, after = '', ''
+        before, after = "", ""
 
         if columns:
             # check column exist
-            conn.execute(f'select {",".join(columns)} from {table_name} limit 0')
+            conn.execute(f"select {','.join(columns)} from {table_name} limit 0")
 
             columns = set(columns)
-            trigger_name += '_' + '_'.join(columns)
+            trigger_name += "_" + "_".join(columns)
 
             news, olds = [], []
             for column in columns:
-                news.append(f'NEW.{column}')
-                olds.append(f'OLD.{column}')
+                news.append(f"NEW.{column}")
+                olds.append(f"OLD.{column}")
 
-            before = f'IF ({", ".join(news)}) IS DISTINCT FROM ({", ".join(olds)}) then\n'
-            after = '\nEND IF;'
+            before = f"IF ({', '.join(news)}) IS DISTINCT FROM ({', '.join(olds)}) then\n"
+            after = "\nEND IF;"
         else:
             columns = set()
 
-        func_code = f'''
+        func_code = f"""
              CREATE OR REPLACE FUNCTION {trigger_name}()
                RETURNS trigger AS $$
              DECLARE
@@ -504,16 +511,16 @@ class PostgresHandler(DatabaseHandler):
                RETURN NEW;
              END;
              $$ LANGUAGE plpgsql;
-         '''
+         """
         conn.execute(func_code)
 
         # for after update - new and old have the same values
-        conn.execute(f'''
+        conn.execute(f"""
              CREATE OR REPLACE TRIGGER {trigger_name}
                BEFORE INSERT OR UPDATE ON {table_name}
                FOR EACH ROW
                EXECUTE PROCEDURE {trigger_name}();
-        ''')
+        """)
         conn.commit()
 
         # start listen
@@ -544,8 +551,208 @@ class PostgresHandler(DatabaseHandler):
                 time.sleep(SUBSCRIBE_SLEEP_INTERVAL)
 
         finally:
-            conn.execute(f'drop TRIGGER {trigger_name} on {table_name}')
-            conn.execute(f'drop FUNCTION {trigger_name}')
+            conn.execute(f"drop TRIGGER {trigger_name} on {table_name}")
+            conn.execute(f"drop FUNCTION {trigger_name}")
             conn.commit()
 
             conn.close()
+
+    def meta_get_tables(self, table_names: Optional[list] = None) -> Response:
+        """
+        Retrieves metadata information about the tables in the PostgreSQL database to be stored in the data catalog.
+
+        Args:
+            table_names (list): A list of table names for which to retrieve metadata information.
+
+        Returns:
+            Response: A response object containing the metadata information, formatted as per the `Response` class.
+        """
+        query = """
+            SELECT
+                t.table_name,
+                t.table_schema,
+                t.table_type,
+                obj_description(pgc.oid, 'pg_class') AS table_description,
+                pgc.reltuples AS row_count
+            FROM information_schema.tables t
+            JOIN pg_catalog.pg_class pgc ON pgc.relname = t.table_name
+            JOIN pg_catalog.pg_namespace pgn ON pgn.oid = pgc.relnamespace
+            WHERE t.table_schema = current_schema()
+            AND t.table_type in ('BASE TABLE', 'VIEW')
+            AND t.table_name NOT LIKE 'pg_%'
+            AND t.table_name NOT LIKE 'sql_%'
+        """
+
+        if table_names is not None and len(table_names) > 0:
+            table_names = [f"'{t}'" for t in table_names]
+            query += f" AND t.table_name IN ({','.join(table_names)})"
+
+        result = self.native_query(query)
+        return result
+
+    def meta_get_columns(self, table_names: Optional[list] = None) -> Response:
+        """
+        Retrieves column metadata for the specified tables (or all tables if no list is provided).
+
+        Args:
+            table_names (list): A list of table names for which to retrieve column metadata.
+
+        Returns:
+            Response: A response object containing the column metadata.
+        """
+        query = """
+            SELECT
+                c.table_name,
+                c.column_name,
+                c.data_type,
+                col_description(pgc.oid, c.ordinal_position) AS column_description,
+                c.column_default,
+                (c.is_nullable = 'YES') AS is_nullable
+            FROM information_schema.columns c
+            JOIN pg_catalog.pg_class pgc ON pgc.relname = c.table_name
+            JOIN pg_catalog.pg_namespace pgn ON pgn.oid = pgc.relnamespace
+            WHERE c.table_schema = current_schema()
+            AND pgc.relkind = 'r'  -- Only consider regular tables (avoids indexes, sequences, etc.)
+            AND c.table_name NOT LIKE 'pg_%'
+            AND c.table_name NOT LIKE 'sql_%'
+            AND pgn.nspname = c.table_schema
+        """
+
+        if table_names is not None and len(table_names) > 0:
+            table_names = [f"'{t}'" for t in table_names]
+            query += f" AND c.table_name IN ({','.join(table_names)})"
+
+        result = self.native_query(query)
+        return result
+
+    def meta_get_column_statistics(self, table_names: Optional[list] = None) -> dict:
+        """
+        Retrieves column statistics (e.g., most common values, frequencies, null percentage, and distinct value count)
+        for the specified tables or all tables if no list is provided.
+
+        Args:
+            table_names (list): A list of table names for which to retrieve column statistics.
+
+        Returns:
+            dict: A dictionary containing the column statistics.
+        """
+        query = """
+            SELECT
+                ps.attname AS column_name,
+                ps.tablename AS table_name,
+                ps.most_common_vals AS most_common_values,
+                ps.most_common_freqs::text AS most_common_frequencies,
+                ps.null_frac * 100 AS null_percentage,
+                ps.n_distinct AS distinct_values_count,
+                ps.histogram_bounds AS histogram_bounds
+            FROM pg_stats ps
+            WHERE ps.schemaname = current_schema()
+            AND ps.tablename NOT LIKE 'pg_%'
+            AND ps.tablename NOT LIKE 'sql_%'
+        """
+
+        if table_names is not None and len(table_names) > 0:
+            table_names = [f"'{t}'" for t in table_names]
+            query += f" AND ps.tablename IN ({','.join(table_names)})"
+
+        result = self.native_query(query)
+        df = result.data_frame
+
+        def parse_pg_array_string(x):
+            try:
+                return (
+                    [item.strip(" ,") for row in csv.reader(io.StringIO(x.strip("{}"))) for item in row if item.strip()]
+                    if x
+                    else []
+                )
+            except IndexError:
+                logger.error(f"Error parsing PostgreSQL array string: {x}")
+                return []
+
+        # Convert most_common_values and most_common_frequencies from string representation to lists.
+        df["most_common_values"] = df["most_common_values"].apply(lambda x: parse_pg_array_string(x))
+        df["most_common_frequencies"] = df["most_common_frequencies"].apply(lambda x: parse_pg_array_string(x))
+
+        # Get the minimum and maximum values from the histogram bounds.
+        df["minimum_value"] = df["histogram_bounds"].apply(lambda x: parse_pg_array_string(x)[0] if x else None)
+        df["maximum_value"] = df["histogram_bounds"].apply(lambda x: parse_pg_array_string(x)[-1] if x else None)
+
+        # Handle cases where distinct_values_count is negative (indicating an approximation).
+        df["distinct_values_count"] = df["distinct_values_count"].apply(lambda x: x if x >= 0 else None)
+
+        result.data_frame = df.drop(columns=["histogram_bounds"])
+
+        return result
+
+    def meta_get_primary_keys(self, table_names: Optional[list] = None) -> Response:
+        """
+        Retrieves primary key information for the specified tables (or all tables if no list is provided).
+
+        Args:
+            table_names (list): A list of table names for which to retrieve primary key information.
+
+        Returns:
+            Response: A response object containing the primary key information.
+        """
+        query = """
+            SELECT
+                tc.table_name,
+                kcu.column_name,
+                kcu.ordinal_position,
+                tc.constraint_name
+            FROM
+                information_schema.table_constraints AS tc
+            JOIN
+                information_schema.key_column_usage AS kcu
+            ON
+                tc.constraint_name = kcu.constraint_name
+            WHERE
+                tc.constraint_type = 'PRIMARY KEY'
+                AND tc.table_schema = current_schema()
+        """
+
+        if table_names is not None and len(table_names) > 0:
+            table_names = [f"'{t}'" for t in table_names]
+            query += f" AND tc.table_name IN ({','.join(table_names)})"
+
+        result = self.native_query(query)
+        return result
+
+    def meta_get_foreign_keys(self, table_names: Optional[list] = None) -> Response:
+        """
+        Retrieves foreign key information for the specified tables (or all tables if no list is provided).
+
+        Args:
+            table_names (list): A list of table names for which to retrieve foreign key information.
+
+        Returns:
+            Response: A response object containing the foreign key information.
+        """
+        query = """
+            SELECT
+                ccu.table_name AS parent_table_name,
+                ccu.column_name AS parent_column_name,
+                tc.table_name AS child_table_name,
+                kcu.column_name AS child_column_name,
+                tc.constraint_name
+            FROM
+                information_schema.table_constraints AS tc
+            JOIN
+                information_schema.key_column_usage AS kcu
+            ON
+                tc.constraint_name = kcu.constraint_name
+            JOIN
+                information_schema.constraint_column_usage AS ccu
+            ON
+                ccu.constraint_name = tc.constraint_name
+            WHERE
+                tc.constraint_type = 'FOREIGN KEY'
+                AND tc.table_schema = current_schema()
+        """
+
+        if table_names is not None and len(table_names) > 0:
+            table_names = [f"'{t}'" for t in table_names]
+            query += f" AND tc.table_name IN ({','.join(table_names)})"
+
+        result = self.native_query(query)
+        return result