DBA-732 clean shutdown

Author: greglearns

src/postgres_mcp/server.py

@@ -49,6 +49,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]:
@@ -545,32 +546,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(0)
 
 
 if __name__ == "__main__":