diff --git a/src/postgres_mcp/top_queries/top_queries_calc.py b/src/postgres_mcp/top_queries/top_queries_calc.py index 87191073..54e44138 100644 --- a/src/postgres_mcp/top_queries/top_queries_calc.py +++ b/src/postgres_mcp/top_queries/top_queries_calc.py @@ -139,14 +139,21 @@ async def get_top_resource_queries(self, frac_threshold: float = 0.05) -> str: logger.debug(f"PostgreSQL version: {pg_version}") # Column names changed in PostgreSQL 13 + # Also, wal_bytes was added in PostgreSQL 13 if pg_version >= 13: # PostgreSQL 13 and newer total_time_col = "total_exec_time" mean_time_col = "mean_exec_time" + stddev_time_col = "stddev_exec_time" + wal_bytes_col = "wal_bytes" + wal_bytes_frac = "wal_bytes / NULLIF(SUM(wal_bytes) OVER (), 0) AS total_wal_bytes_frac" else: # PostgreSQL 12 and older total_time_col = "total_time" mean_time_col = "mean_time" + stddev_time_col = "stddev_time" + wal_bytes_col = "0 AS wal_bytes" # Column doesn't exist in PG12 + wal_bytes_frac = "0 AS total_wal_bytes_frac" query = cast( LiteralString, @@ -156,18 +163,23 @@ async def get_top_resource_queries(self, frac_threshold: float = 0.05) -> str: query, calls, rows, - {total_time_col} total_exec_time, - {mean_time_col} mean_exec_time, - stddev_exec_time, + {total_time_col} AS total_exec_time, + {mean_time_col} AS mean_exec_time, + {stddev_time_col} AS stddev_exec_time, shared_blks_hit, shared_blks_read, shared_blks_dirtied, - wal_bytes, - total_exec_time / SUM(total_exec_time) OVER () AS total_exec_time_frac, - (shared_blks_hit + shared_blks_read) / SUM(shared_blks_hit + shared_blks_read) OVER () AS shared_blks_accessed_frac, - shared_blks_read / SUM(shared_blks_read) OVER () AS shared_blks_read_frac, - shared_blks_dirtied / SUM(shared_blks_dirtied) OVER () AS shared_blks_dirtied_frac, - wal_bytes / SUM(wal_bytes) OVER () AS total_wal_bytes_frac + {wal_bytes_col}, + {total_time_col} / NULLIF(SUM({total_time_col}) OVER (), 0) + AS total_exec_time_frac, + (shared_blks_hit + shared_blks_read) + / NULLIF(SUM(shared_blks_hit + shared_blks_read) OVER (), 0) + AS shared_blks_accessed_frac, + shared_blks_read / NULLIF(SUM(shared_blks_read) OVER (), 0) + AS shared_blks_read_frac, + shared_blks_dirtied / NULLIF(SUM(shared_blks_dirtied) OVER (), 0) + AS shared_blks_dirtied_frac, + {wal_bytes_frac} FROM pg_stat_statements ) SELECT diff --git a/tests/conftest.py b/tests/conftest.py index f202ff67..e272d29b 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -17,7 +17,7 @@ def event_loop_policy(): return asyncio.DefaultEventLoopPolicy() -@pytest.fixture(scope="class", params=["postgres:15", "postgres:16"]) +@pytest.fixture(scope="class", params=["postgres:12", "postgres:15", "postgres:16"]) def test_postgres_connection_string(request) -> Generator[tuple[str, str], None, None]: yield from create_postgres_container(request.param) diff --git a/tests/unit/top_queries/test_top_queries_calc.py b/tests/unit/top_queries/test_top_queries_calc.py index 9dadba34..1f076045 100644 --- a/tests/unit/top_queries/test_top_queries_calc.py +++ b/tests/unit/top_queries/test_top_queries_calc.py @@ -212,3 +212,43 @@ async def test_error_handling(mock_pg13_driver, mock_extension_installed): # Check that the error is properly reported assert "Error getting slow queries: Database error" in result + + +@pytest.mark.asyncio +async def test_resource_queries_pg12(mock_pg12_driver, mock_extension_installed): + """Test resource queries on PostgreSQL 12 uses correct column names (stddev_time, no wal_bytes).""" + # Create the TopQueriesCalc instance with the mock driver + calc = TopQueriesCalc(sql_driver=mock_pg12_driver) + + # Get resource queries + _result = await calc.get_top_resource_queries(frac_threshold=0.05) + + # Get the executed query from the mock + call_args = str(mock_pg12_driver.execute_query.call_args) + + # Verify PG12 column names are used (aliased to consistent output names) + assert "stddev_time AS stddev_exec_time" in call_args, "Should use stddev_time aliased for PG12" + assert "total_time AS total_exec_time" in call_args, "Should use total_time aliased for PG12" + assert "mean_time AS mean_exec_time" in call_args, "Should use mean_time aliased for PG12" + # wal_bytes should be replaced with 0 for PG12 + assert "0 AS wal_bytes" in call_args, "Should use 0 AS wal_bytes for PG12" + + +@pytest.mark.asyncio +async def test_resource_queries_pg13(mock_pg13_driver, mock_extension_installed): + """Test resource queries on PostgreSQL 13 uses correct column names (stddev_exec_time, wal_bytes).""" + # Create the TopQueriesCalc instance with the mock driver + calc = TopQueriesCalc(sql_driver=mock_pg13_driver) + + # Get resource queries + _result = await calc.get_top_resource_queries(frac_threshold=0.05) + + # Get the executed query from the mock + call_args = str(mock_pg13_driver.execute_query.call_args) + + # Verify PG13 column names are used + assert "stddev_exec_time" in call_args, "Should use stddev_exec_time for PG13" + assert "total_exec_time" in call_args, "Should use total_exec_time for PG13" + assert "mean_exec_time" in call_args, "Should use mean_exec_time for PG13" + # wal_bytes should be the actual column for PG13 + assert "wal_bytes" in call_args, "Should use wal_bytes column for PG13"