MCP PostgreSQL Operations

""" PostgreSQL Version Compatibility Utilities Provides version detection and compatibility handling for MCP PostgreSQL tools. """ import re import logging from typing import Tuple, Optional from .functions import execute_single_query logger = logging.getLogger(__name__) class PostgreSQLVersion: """PostgreSQL version information and compatibility utilities.""" def __init__(self, major: int, minor: int = 0, patch: int = 0): self.major = major self.minor = minor self.patch = patch def __str__(self): return f"{self.major}.{self.minor}.{self.patch}" def __ge__(self, other): if isinstance(other, (int, float)): return self.major >= other if isinstance(other, PostgreSQLVersion): return (self.major, self.minor, self.patch) >= (other.major, other.minor, other.patch) return False def __lt__(self, other): if isinstance(other, (int, float)): return self.major < other if isinstance(other, PostgreSQLVersion): return (self.major, self.minor, self.patch) < (other.major, other.minor, other.patch) return False @property def is_modern(self) -> bool: """Check if this is a modern PostgreSQL version (12+).""" return self.major >= 12 @property def has_checkpointer_split(self) -> bool: """Check if checkpointer stats are in separate view (only PG15).""" return self.major == 15 @property def has_pg_stat_io(self) -> bool: """Check if pg_stat_io view is available (16+).""" return self.major >= 16 @property def has_enhanced_wal_receiver(self) -> bool: """Check if pg_stat_wal_receiver has written_lsn/flushed_lsn (16+).""" return self.major >= 16 @property def has_replication_slot_stats(self) -> bool: """Check if pg_stat_replication_slots is available (14+).""" return self.major >= 14 @property def has_parallel_leader_tracking(self) -> bool: """Check if pg_stat_activity has leader_pid column (14+).""" return self.major >= 14 @property def has_replication_slot_wal_status(self) -> bool: """Check if pg_replication_slots has wal_status and safe_wal_size columns (13+).""" return self.major >= 13 @property def has_table_stats_ins_since_vacuum(self) -> bool: """Check if pg_stat_*_tables has n_ins_since_vacuum column (13+).""" return self.major >= 13 @property def has_pg_stat_statements_exec_time(self) -> bool: """Check if pg_stat_statements uses total_exec_time and mean_exec_time columns (13+).""" return self.major >= 13 # Global version cache _cached_version: Optional[PostgreSQLVersion] = None async def get_postgresql_version(database: str = None, force_refresh: bool = False) -> PostgreSQLVersion: """ Get PostgreSQL server version with caching. Args: database: Database to connect to force_refresh: Force refresh cached version Returns: PostgreSQLVersion object """ global _cached_version if _cached_version is not None and not force_refresh: return _cached_version try: result = await execute_single_query("SELECT version()", database=database) version_string = result.get('version', '') # Parse version string like "PostgreSQL 16.1 on x86_64-pc-linux-gnu..." version_match = re.search(r'PostgreSQL\s+(\d+)\.?(\d*)\.?(\d*)', version_string) if version_match: major = int(version_match.group(1)) minor = int(version_match.group(2) or 0) patch = int(version_match.group(3) or 0) _cached_version = PostgreSQLVersion(major, minor, patch) logger.info(f"Detected PostgreSQL version: {_cached_version}") return _cached_version else: logger.warning(f"Could not parse version string: {version_string}") # Default to PostgreSQL 17 if parsing fails _cached_version = PostgreSQLVersion(17, 0, 0) return _cached_version except Exception as e: logger.error(f"Failed to get PostgreSQL version: {e}") # Default to PostgreSQL 17 if version detection fails _cached_version = PostgreSQLVersion(17, 0, 0) return _cached_version async def check_feature_availability(feature: str, database: str = None) -> bool: """ Check if a specific PostgreSQL feature is available. Args: feature: Feature name to check database: Database to connect to Returns: True if feature is available """ version = await get_postgresql_version(database) feature_requirements = { 'pg_stat_io': version.has_pg_stat_io, 'checkpointer_split': version.has_checkpointer_split, 'enhanced_wal_receiver': version.has_enhanced_wal_receiver, 'replication_slot_stats': version.has_replication_slot_stats, 'parallel_leader_tracking': version.has_parallel_leader_tracking, } return feature_requirements.get(feature, False) async def get_compatible_column_list(table_name: str, all_columns: list, version_specific_columns: dict, database: str = None) -> str: """ Generate version-compatible column list for SELECT queries. Args: table_name: PostgreSQL table/view name all_columns: List of all possible columns version_specific_columns: Dict mapping version requirements to columns database: Database to connect to Returns: Comma-separated column list for SQL SELECT """ version = await get_postgresql_version(database) available_columns = [] for col in all_columns: # Check if column has version requirements required_version = version_specific_columns.get(col) if required_version is None: # No version requirement - always available available_columns.append(col) elif version >= required_version: # Version requirement met available_columns.append(col) else: # Version requirement not met - provide NULL placeholder available_columns.append(f"NULL::text AS {col}") return ", ".join(available_columns) def get_version_appropriate_query(base_query: str, version_variants: dict, version: PostgreSQLVersion) -> str: """ Select version-appropriate query variant. Args: base_query: Default/fallback query version_variants: Dict mapping version requirements to query variants version: PostgreSQL version Returns: Most appropriate query for the version """ # Sort version variants by version (highest first) sorted_variants = sorted(version_variants.items(), key=lambda x: (x[0].major, x[0].minor), reverse=True) for required_version, query in sorted_variants: if version >= required_version: return query return base_query # Version-specific query builders class VersionAwareQueries: """Collection of version-aware query builders.""" @staticmethod async def get_bgwriter_checkpointer_stats(database: str = None) -> str: """Get background writer/checkpointer stats with version compatibility.""" version = await get_postgresql_version(database) if version.has_checkpointer_split: # PostgreSQL 15+: Use separate checkpointer view return """ SELECT 'checkpointer' as component, num_timed, num_requested, restartpoints_timed, restartpoints_req, restartpoints_done, write_time, sync_time, buffers_written, stats_reset FROM pg_stat_checkpointer UNION ALL SELECT 'bgwriter' as component, NULL::bigint as num_timed, NULL::bigint as num_requested, NULL::bigint as restartpoints_timed, NULL::bigint as restartpoints_req, NULL::bigint as restartpoints_done, NULL::double precision as write_time, NULL::double precision as sync_time, buffers_clean as buffers_written, stats_reset FROM pg_stat_bgwriter """ else: # PostgreSQL 10-14: All stats in bgwriter view return """ SELECT 'bgwriter_legacy' as component, buffers_clean, maxwritten_clean, buffers_alloc, stats_reset, NULL::bigint as num_timed, NULL::bigint as num_requested FROM pg_stat_bgwriter """ @staticmethod async def get_io_statistics(database: str = None) -> str: """Get I/O statistics with version compatibility.""" version = await get_postgresql_version(database) if version.has_pg_stat_io: # PostgreSQL 16+: Use comprehensive pg_stat_io return """ SELECT backend_type, object, context, reads, read_time, writes, write_time, extends, extend_time, hits, evictions, reuses, fsyncs, fsync_time FROM pg_stat_io WHERE reads > 0 OR writes > 0 OR hits > 0 """ else: # PostgreSQL 10-15: Fall back to pg_statio_* views return """ SELECT 'client backend' as backend_type, 'relation' as object, 'normal' as context, heap_blks_read as reads, 0::double precision as read_time, 0::bigint as writes, 0::double precision as write_time, 0::bigint as extends, 0::double precision as extend_time, heap_blks_hit as hits, 0::bigint as evictions, 0::bigint as reuses, 0::bigint as fsyncs, 0::double precision as fsync_time FROM pg_statio_all_tables WHERE heap_blks_read > 0 OR heap_blks_hit > 0 """ @staticmethod async def get_activity_with_leader_info(database: str = None) -> str: """Get activity info with parallel leader tracking if available.""" version = await get_postgresql_version(database) base_columns = [ "pid", "datname", "usename", "application_name", "client_addr", "state", "query_start", "query" ] version_columns = { "leader_pid": PostgreSQLVersion(14), "query_id": PostgreSQLVersion(13) } columns = await get_compatible_column_list( "pg_stat_activity", base_columns + list(version_columns.keys()), version_columns, database ) return f"SELECT {columns} FROM pg_stat_activity WHERE state = 'active'" @staticmethod async def get_replication_slots_query(database: str = None) -> str: """Get replication slots info with version compatibility.""" version = await get_postgresql_version(database) base_columns = [ "slot_name", "plugin", "slot_type", "datoid", "temporary", "active", "active_pid", "restart_lsn", "confirmed_flush_lsn" ] # wal_status and safe_wal_size are available from PostgreSQL 13+ if version.has_replication_slot_wal_status: return """ SELECT slot_name, plugin, slot_type, datoid, temporary, active, active_pid, restart_lsn, confirmed_flush_lsn, wal_status, safe_wal_size / 1024 / 1024 as safe_wal_size_mb FROM pg_replication_slots ORDER BY slot_name """ else: # PostgreSQL 12 and older - without wal_status and safe_wal_size return """ SELECT slot_name, plugin, slot_type, datoid, temporary, active, active_pid, restart_lsn, confirmed_flush_lsn, NULL::text as wal_status, NULL::numeric as safe_wal_size_mb FROM pg_replication_slots ORDER BY slot_name """ @staticmethod async def get_wal_receiver_query(database: str = None) -> str: """Get WAL receiver status with version compatibility.""" version = await get_postgresql_version(database) if version.has_enhanced_wal_receiver: # PostgreSQL 16+: has written_lsn/flushed_lsn columns return """ SELECT pid, status, receive_start_lsn, receive_start_tli, written_lsn, flushed_lsn, received_tli, last_msg_send_time, last_msg_receipt_time, latest_end_lsn, latest_end_time, slot_name, sender_host, sender_port, conninfo FROM pg_stat_wal_receiver """ else: # PostgreSQL 10-15: no written_lsn/flushed_lsn columns return """ SELECT pid, status, receive_start_lsn, receive_start_tli, NULL::text as written_lsn, NULL::text as flushed_lsn, received_tli, last_msg_send_time, last_msg_receipt_time, latest_end_lsn, latest_end_time, slot_name, sender_host, sender_port, conninfo FROM pg_stat_wal_receiver """ @staticmethod async def get_all_tables_stats_query(include_system: bool = False, database: str = None) -> str: """Get all tables statistics query with version compatibility.""" version = await get_postgresql_version(database) view_name = "pg_stat_all_tables" if include_system else "pg_stat_user_tables" # n_ins_since_vacuum is available from PostgreSQL 13+ if version.has_table_stats_ins_since_vacuum: return f""" SELECT schemaname as schema_name, relname as table_name, seq_scan as sequential_scans, seq_tup_read as seq_tuples_read, idx_scan as index_scans, idx_tup_fetch as idx_tuples_fetched, n_tup_ins as tuples_inserted, n_tup_upd as tuples_updated, n_tup_del as tuples_deleted, n_tup_hot_upd as hot_updates, n_live_tup as estimated_live_tuples, n_dead_tup as estimated_dead_tuples, CASE WHEN n_live_tup > 0 THEN ROUND((n_dead_tup::numeric / n_live_tup) * 100, 2) ELSE 0 END as dead_tuple_ratio_percent, n_mod_since_analyze as modified_since_analyze, n_ins_since_vacuum as inserted_since_vacuum, last_vacuum, last_autovacuum, last_analyze, last_autoanalyze, vacuum_count, autovacuum_count, analyze_count, autoanalyze_count FROM {view_name} ORDER BY seq_scan + COALESCE(idx_scan, 0) DESC, schemaname, relname """ else: # PostgreSQL 12 - without n_ins_since_vacuum return f""" SELECT schemaname as schema_name, relname as table_name, seq_scan as sequential_scans, seq_tup_read as seq_tuples_read, idx_scan as index_scans, idx_tup_fetch as idx_tuples_fetched, n_tup_ins as tuples_inserted, n_tup_upd as tuples_updated, n_tup_del as tuples_deleted, n_tup_hot_upd as hot_updates, n_live_tup as estimated_live_tuples, n_dead_tup as estimated_dead_tuples, CASE WHEN n_live_tup > 0 THEN ROUND((n_dead_tup::numeric / n_live_tup) * 100, 2) ELSE 0 END as dead_tuple_ratio_percent, n_mod_since_analyze as modified_since_analyze, NULL::bigint as inserted_since_vacuum, last_vacuum, last_autovacuum, last_analyze, last_autoanalyze, vacuum_count, autovacuum_count, analyze_count, autoanalyze_count FROM {view_name} ORDER BY seq_scan + COALESCE(idx_scan, 0) DESC, schemaname, relname """ # Utility functions for tool implementations async def execute_version_aware_query(queries_by_version: dict, fallback_query: str, database: str = None): """ Execute the most appropriate query based on PostgreSQL version. Args: queries_by_version: Dict mapping PostgreSQLVersion to query strings fallback_query: Default query if no version matches database: Database to connect to Returns: Query results """ from .functions import execute_query version = await get_postgresql_version(database) query = get_version_appropriate_query(fallback_query, queries_by_version, version) return await execute_query(query, database=database) # Version-aware pg_stat_statements queries async def get_pg_stat_statements_query(database: str = None) -> str: """ Get version-compatible pg_stat_statements query. Args: database: Database to connect to for version detection Returns: SQL query string compatible with the database version """ version = await get_postgresql_version(database) # Common base columns available in all versions base_columns = [ "queryid", "query", "calls", "rows" ] # Add version-specific timing columns if version.has_pg_stat_statements_exec_time: # PostgreSQL 13+: uses total_exec_time, mean_exec_time base_columns.extend([ "total_exec_time", "mean_exec_time", "min_exec_time", "max_exec_time", "stddev_exec_time" ]) else: # PostgreSQL 12: uses total_time, mean_time base_columns.extend([ "total_time as total_exec_time", "mean_time as mean_exec_time", "min_time as min_exec_time", "max_time as max_exec_time", "stddev_time as stddev_exec_time" ]) # Add remaining common columns base_columns.extend([ "shared_blks_hit", "shared_blks_read", "shared_blks_dirtied", "shared_blks_written", "local_blks_hit", "local_blks_read", "local_blks_dirtied", "local_blks_written", "temp_blks_read", "temp_blks_written" ]) columns_str = ",\n ".join(base_columns) return f""" SELECT {columns_str} FROM pg_stat_statements ORDER BY total_exec_time DESC """ # Version-aware pg_stat_monitor queries async def get_pg_stat_monitor_query(database: str = None) -> str: """ Get version-compatible pg_stat_monitor query. Args: database: Database to connect to for version detection Returns: SQL query string compatible with the database version """ version = await get_postgresql_version(database) # Common base columns available in all versions base_columns = [ "query", "calls", "rows" ] # Add version-specific timing columns if version.has_pg_stat_statements_exec_time: # PostgreSQL 13+: uses total_exec_time, mean_exec_time base_columns.extend([ "total_exec_time", "mean_exec_time" ]) else: # PostgreSQL 12: uses total_time, mean_time base_columns.extend([ "total_time as total_exec_time", "mean_time as mean_exec_time" ]) # Add remaining common columns base_columns.extend([ "shared_blks_hit", "shared_blks_read", "client_ip", "bucket_start_time" ]) columns_str = ",\n ".join(base_columns) return f""" SELECT {columns_str} FROM pg_stat_monitor ORDER BY total_exec_time DESC """