DBA-732 graceful shutdown (#42)

Author: greglearns

justfile

@@ -14,13 +14,22 @@ dev:
 nix-claude-desktop:
   NIXPKGS_ALLOW_UNFREE=1 nix run "github:k3d3/claude-desktop-linux-flake#claude-desktop-with-fhs" --impure
 
-release version note="Release v{{version}}" extra="": # NOTE version format should be 0.0.0
+release-help:
+  @echo "- update version in pyproject.toml"
+  @echo "- uv sync"
+  @echo "- git commit"
+  @echo "- git push && merge to main"
+  @echo '- just release 0.0.0 "note"'
+  @echo 'OR'
+  @echo '- just prerelease 0.0.0 1 "note"'
+
+release version note extra="": # ="Release v{{version}}" extra="": # NOTE version format should be 0.0.0
   #!/usr/bin/env bash
   if [[ "{{version}}" == v* ]]; then
     echo "Error: Do not include 'v' prefix in version. It will be added automatically."
     exit 1
   fi
-  uv build && git tag -a "v{{version}}" -m "Release v{{version}}" && git push --tags && gh release candidate "v{{version}}" --title "PostgreSQL MCP v{{version}}" --notes "{{note}}" {{extra}} dist/*.whl dist/*.tar.gz
+  uv build && git tag -a "v{{version}}" -m "Release v{{version}}" || true && git push --tags && gh release create "v{{version}}" --title "PostgreSQL MCP v{{version}}" --notes "{{note}}" {{extra}} dist/*.whl dist/*.tar.gz
 
-prerelease version rc note="Release candidate {{rc}} for version {{version}}":
+prerelease version rc note: #="Release candidate {{rc}} for version {{version}}":
   just release "{{version}}rc{{rc}}" "{{note}}" "--prerelease"

pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "postgres-mcp"
-version = "0.1.1rc1"
+version = "0.1.1"
 description = "PostgreSQL Tuning and Analysis Tool"
 readme = "README.md"
 requires-python = ">=3.12"

src/postgres_mcp/server.py

@@ -48,6 +48,7 @@ class AccessMode(str, Enum):
 # Global variables
 db_connection = DbConnPool()
 current_access_mode = AccessMode.UNRESTRICTED
+shutdown_in_progress = False
 
 
 async def get_sql_driver() -> Union[SqlDriver, SafeSqlDriver]:
@@ -542,32 +543,26 @@ async def main():
         logger.warning("Signal handling not supported on Windows")
         pass
 
-    # Run the app with FastMCP's stdio method
-    try:
-        await mcp.run_stdio_async()
-    finally:
-        # Close the connection pool when exiting
-        await shutdown()
+    await mcp.run_stdio_async()
 
 
 async def shutdown(sig=None):
     """Clean shutdown of the server."""
-    if sig:
-        logger.info(f"Received exit signal {sig.name}")
+    global shutdown_in_progress
 
-    logger.info("Closing database connections...")
-    await db_connection.close()
+    import os
 
-    # Give tasks a chance to complete
-    try:
-        tasks = [t for t in asyncio.all_tasks() if t is not asyncio.current_task()]
-        if tasks:
-            logger.info(f"Waiting for {len(tasks)} tasks to complete...")
-            await asyncio.gather(*tasks, return_exceptions=True)
-    except Exception as e:
-        logger.warning(f"Error during shutdown: {e}")
+    if shutdown_in_progress:
+        logger.warning("Forcing immediate exit")
+
+        os._exit(1)  # Use immediate process termination instead of sys.exit
+
+    shutdown_in_progress = True
+
+    if sig:
+        logger.info(f"Received exit signal {sig.name}")
 
-    logger.info("Shutdown complete.")
+    os._exit(128 + sig if sig is not None else 0)
 
 
 if __name__ == "__main__":