diff --git a/README.md b/README.md index d5005ed6..28922c51 100644 --- a/README.md +++ b/README.md @@ -226,6 +226,30 @@ Postgres MCP Pro supports multiple *access modes* to give you control over the o To use restricted mode, replace `--access-mode=unrestricted` with `--access-mode=restricted` in the configuration examples above. +#### Autocommit Mode + +Some database proxies (e.g. [QueryPie](https://www.querypie.com/), [PgBouncer](https://www.pgbouncer.org/) in transaction mode) block explicit transaction control statements (`BEGIN`, `COMMIT`, `ROLLBACK`). The `--autocommit` flag disables implicit transactions so that each statement executes independently: + +```json +{ + "mcpServers": { + "postgres": { + "command": "uvx", + "args": [ + "postgres-mcp", + "--access-mode=restricted", + "--autocommit" + ], + "env": { + "DATABASE_URI": "postgresql://username:password@localhost:5432/dbname" + } + } + } +} +``` + +You can also enable autocommit via the `DATABASE_AUTOCOMMIT` environment variable (`true`, `1`, or `yes`). + #### Other MCP Clients diff --git a/src/postgres_mcp/server.py b/src/postgres_mcp/server.py index f3ba8f8b..c5708abd 100644 --- a/src/postgres_mcp/server.py +++ b/src/postgres_mcp/server.py @@ -596,6 +596,14 @@ async def main(): default=8000, help="Port for streamable HTTP server (default: 8000)", ) + parser.add_argument( + "--autocommit", + action="store_true", + default=False, + help="Use autocommit mode, disabling implicit transactions (BEGIN/COMMIT/ROLLBACK). " + "Useful for database proxies that block transaction control statements. " + "Can also be set via DATABASE_AUTOCOMMIT environment variable (true/1/yes).", + ) args = parser.parse_args() @@ -633,9 +641,15 @@ async def main(): "Error: No database URL provided. Please specify via 'DATABASE_URI' environment variable or command-line argument.", ) + # Determine autocommit mode from CLI flag or environment variable + use_autocommit = args.autocommit or os.environ.get("DATABASE_AUTOCOMMIT", "").lower() in ("true", "1", "yes") + + if use_autocommit: + logger.info("Autocommit mode enabled - transactions (BEGIN/COMMIT/ROLLBACK) will be skipped") + # Initialize database connection pool try: - await db_connection.pool_connect(database_url) + await db_connection.pool_connect(database_url, autocommit=use_autocommit) logger.info("Successfully connected to database and initialized connection pool") except Exception as e: logger.warning( diff --git a/src/postgres_mcp/sql/sql_driver.py b/src/postgres_mcp/sql/sql_driver.py index 5beacb03..81ddd0e2 100644 --- a/src/postgres_mcp/sql/sql_driver.py +++ b/src/postgres_mcp/sql/sql_driver.py @@ -68,8 +68,15 @@ def __init__(self, connection_url: Optional[str] = None): self._is_valid = False self._last_error = None - async def pool_connect(self, connection_url: Optional[str] = None) -> AsyncConnectionPool: - """Initialize connection pool with retry logic.""" + async def pool_connect(self, connection_url: Optional[str] = None, autocommit: bool = False) -> AsyncConnectionPool: + """Initialize connection pool with retry logic. + + Args: + connection_url: PostgreSQL connection URL + autocommit: When True, connections use autocommit mode (no implicit BEGIN). + Useful for database proxies (e.g. QueryPie) that block + transaction control statements. + """ # If we already have a valid pool, return it if self.pool and self._is_valid: return self.pool @@ -86,11 +93,13 @@ async def pool_connect(self, connection_url: Optional[str] = None) -> AsyncConne try: # Configure connection pool with appropriate settings + pool_kwargs = {"autocommit": True} if autocommit else {} self.pool = AsyncConnectionPool( conninfo=url, min_size=1, max_size=5, open=False, # Don't connect immediately, let's do it explicitly + kwargs=pool_kwargs, ) # Open the pool explicitly @@ -223,11 +232,12 @@ async def execute_query( async def _execute_with_connection(self, connection, query, params, force_readonly) -> Optional[List[RowResult]]: """Execute query with the given connection.""" + is_autocommit = connection.autocommit transaction_started = False try: async with connection.cursor(row_factory=dict_row) as cursor: - # Start read-only transaction - if force_readonly: + # Start read-only transaction (skip in autocommit mode) + if force_readonly and not is_autocommit: await cursor.execute("BEGIN TRANSACTION READ ONLY") transaction_started = True @@ -241,28 +251,30 @@ async def _execute_with_connection(self, connection, query, params, force_readon pass if cursor.description is None: # No results (like DDL statements) - if not force_readonly: - await cursor.execute("COMMIT") - elif transaction_started: - await cursor.execute("ROLLBACK") - transaction_started = False + if not is_autocommit: + if not force_readonly: + await cursor.execute("COMMIT") + elif transaction_started: + await cursor.execute("ROLLBACK") + transaction_started = False return None # Get results from the last statement only rows = await cursor.fetchall() - # End the transaction appropriately - if not force_readonly: - await cursor.execute("COMMIT") - elif transaction_started: - await cursor.execute("ROLLBACK") - transaction_started = False + # End the transaction appropriately (skip in autocommit mode) + if not is_autocommit: + if not force_readonly: + await cursor.execute("COMMIT") + elif transaction_started: + await cursor.execute("ROLLBACK") + transaction_started = False return [SqlDriver.RowResult(cells=dict(row)) for row in rows] except Exception as e: # Try to roll back the transaction if it's still active - if transaction_started: + if transaction_started and not is_autocommit: try: await connection.rollback() except Exception as rollback_error: diff --git a/tests/unit/sql/test_sql_driver.py b/tests/unit/sql/test_sql_driver.py index 4033537d..e077ea9b 100644 --- a/tests/unit/sql/test_sql_driver.py +++ b/tests/unit/sql/test_sql_driver.py @@ -365,3 +365,105 @@ async def test_engine_url_connection(): # Verify driver state assert driver.is_pool is True assert driver.conn is not None + + +def _make_autocommit_connection(autocommit=True, has_results=True): + """Create a properly mocked async connection for _execute_with_connection tests.""" + cursor = AsyncMock() + cursor.nextset = MagicMock(return_value=False) + cursor.fetchall = AsyncMock(return_value=[ + {"id": 1, "name": "test1"}, + {"id": 2, "name": "test2"}, + ]) + cursor.description = ["column1", "column2"] if has_results else None + + cursor_ctx = AsyncMock() + cursor_ctx.__aenter__ = AsyncMock(return_value=cursor) + cursor_ctx.__aexit__ = AsyncMock(return_value=False) + + connection = MagicMock() + connection.autocommit = autocommit + connection.cursor = MagicMock(return_value=cursor_ctx) + connection.rollback = AsyncMock() + + return connection, cursor + + +@pytest.mark.asyncio +async def test_execute_query_autocommit_skips_transaction(): + """Test that autocommit mode skips BEGIN/COMMIT/ROLLBACK statements.""" + connection, cursor = _make_autocommit_connection(autocommit=True) + + driver = SqlDriver(conn=connection) + + result = await driver._execute_with_connection( + connection, "SELECT * FROM test", None, force_readonly=True + ) + + # Verify no transaction statements were issued + for c in cursor.execute.call_args_list: + sql = c[0][0] if c[0] else "" + assert "BEGIN" not in sql, "BEGIN should not be called in autocommit mode" + assert "COMMIT" not in sql, "COMMIT should not be called in autocommit mode" + assert "ROLLBACK" not in sql, "ROLLBACK should not be called in autocommit mode" + + # Verify the actual query was still executed + assert call("SELECT * FROM test") in cursor.execute.call_args_list + + # Verify results were returned + assert result is not None + assert len(result) == 2 + + +@pytest.mark.asyncio +async def test_execute_query_autocommit_no_results(): + """Test autocommit mode with DDL statements that return no results.""" + connection, cursor = _make_autocommit_connection(autocommit=True, has_results=False) + + driver = SqlDriver(conn=connection) + + result = await driver._execute_with_connection( + connection, "CREATE TABLE test (id int)", None, force_readonly=False + ) + + # Verify no transaction statements were issued + for c in cursor.execute.call_args_list: + sql = c[0][0] if c[0] else "" + assert "COMMIT" not in sql, "COMMIT should not be called in autocommit mode" + + assert result is None + + +@pytest.mark.asyncio +async def test_execute_query_non_autocommit_uses_transaction(): + """Test that non-autocommit mode still uses BEGIN/COMMIT/ROLLBACK as before.""" + connection, cursor = _make_autocommit_connection(autocommit=False) + + driver = SqlDriver(conn=connection) + + await driver._execute_with_connection( + connection, "SELECT * FROM test", None, force_readonly=True + ) + + # Verify transaction statements WERE issued + assert call("BEGIN TRANSACTION READ ONLY") in cursor.execute.call_args_list + assert call("ROLLBACK") in cursor.execute.call_args_list + + +@pytest.mark.asyncio +async def test_execute_query_autocommit_error_skips_rollback(): + """Test that autocommit mode skips rollback on error.""" + connection, cursor = _make_autocommit_connection(autocommit=True) + + # Make query execution fail + cursor.execute = AsyncMock(side_effect=Exception("Query failed")) + + driver = SqlDriver(conn=connection) + + with pytest.raises(Exception, match="Query failed"): + await driver._execute_with_connection( + connection, "SELECT * FROM test", None, force_readonly=True + ) + + # Verify rollback was NOT called (autocommit mode) + connection.rollback.assert_not_awaited()