feat: enable SnapStart on lambda functions

infrastructure/app.py

@@ -181,6 +181,7 @@ def __init__(
                 "handler": "handler.handler",
                 "runtime": aws_lambda.Runtime.PYTHON_3_12,
             },
+            enable_snap_start=True,
         )
 
         #######################################################################
@@ -228,6 +229,7 @@ def __init__(
                 "handler": "handler.handler",
                 "runtime": aws_lambda.Runtime.PYTHON_3_12,
             },
+            enable_snap_start=True,
         )
 
         #######################################################################
@@ -274,6 +276,7 @@ def __init__(
                 "handler": "handler.handler",
                 "runtime": aws_lambda.Runtime.PYTHON_3_12,
             },
+            enable_snap_start=True,
         )
 
         if app_config.stac_ingestor:

infrastructure/handlers/raster_handler.py

@@ -7,16 +7,70 @@
 from eoapi.raster.app import app
 from eoapi.raster.config import PostgresSettings
 from mangum import Mangum
+from snapshot_restore_py import register_after_restore, register_before_snapshot
 from titiler.pgstac.db import connect_to_db
 
 logging.getLogger("mangum.lifespan").setLevel(logging.ERROR)
 logging.getLogger("mangum.http").setLevel(logging.ERROR)
 
+postgres_settings = PostgresSettings()
+_connection_initialized = False
+
+
+@register_before_snapshot
+def on_snapshot():
+    """
+    Runtime hook called by Lambda before taking a snapshot.
+    We close database connections that shouldn't be in the snapshot.
+    """
+
+    if hasattr(app, "state") and hasattr(app.state, "dbpool") and app.state.dbpool:
+        try:
+            app.state.dbpool.close()
+            app.state.dbpool = None
+        except Exception as e:
+            print(f"SnapStart: Error closing database pool: {e}")
+
+    return {"statusCode": 200}
+
+
+@register_after_restore
+def on_snap_restore():
+    """
+    Runtime hook called by Lambda after restoring from a snapshot.
+    We recreate database connections that were closed before the snapshot.
+    """
+    global _connection_initialized
+
+    try:
+        try:
+            loop = asyncio.get_running_loop()
+        except RuntimeError:
+            loop = asyncio.new_event_loop()
+            asyncio.set_event_loop(loop)
+
+        if hasattr(app.state, "dbpool") and app.state.dbpool:
+            try:
+                app.state.dbpool.close()
+            except Exception as e:
+                print(f"SnapStart: Error closing stale pool: {e}")
+            app.state.dbpool = None
+
+        loop.run_until_complete(connect_to_db(app, settings=postgres_settings))
+
+        _connection_initialized = True
+
+    except Exception as e:
+        print(f"SnapStart: Failed to initialize database connection: {e}")
+        raise
+
+    return {"statusCode": 200}
+
 
 @app.on_event("startup")
 async def startup_event() -> None:
     """Connect to database on startup."""
-    await connect_to_db(app, settings=PostgresSettings())
+    await connect_to_db(app, settings=postgres_settings)
 
 
 handler = Mangum(app, lifespan="off")

infrastructure/handlers/stac_handler.py

@@ -7,20 +7,98 @@
 from eoapi.stac.app import app
 from eoapi.stac.config import PostgresSettings, Settings
 from mangum import Mangum
+from snapshot_restore_py import register_after_restore, register_before_snapshot
 from stac_fastapi.pgstac.db import connect_to_db
 
 logging.getLogger("mangum.lifespan").setLevel(logging.ERROR)
 logging.getLogger("mangum.http").setLevel(logging.ERROR)
 
 settings = Settings()
+postgres_settings = PostgresSettings()
+
+_connection_initialized = False
+
+
+@register_before_snapshot
+def on_snapshot():
+    """
+    Runtime hook called by Lambda before taking a snapshot.
+    We close database connections that shouldn't be in the snapshot.
+    """
+
+    if hasattr(app, "state") and hasattr(app.state, "readpool") and app.state.readpool:
+        try:
+            app.state.readpool.close()
+            app.state.readpool = None
+        except Exception as e:
+            print(f"SnapStart: Error closing database readpool: {e}")
+
+    if (
+        hasattr(app, "state")
+        and hasattr(app.state, "writepool")
+        and app.state.writepool
+    ):
+        try:
+            app.state.writepool.close()
+            app.state.writepool = None
+        except Exception as e:
+            print(f"SnapStart: Error closing database writepool: {e}")
+
+    return {"statusCode": 200}
+
+
+@register_after_restore
+def on_snap_restore():
+    """
+    Runtime hook called by Lambda after restoring from a snapshot.
+    We recreate database connections that were closed before the snapshot.
+    """
+    global _connection_initialized
+
+    try:
+        try:
+            loop = asyncio.get_running_loop()
+        except RuntimeError:
+            loop = asyncio.new_event_loop()
+            asyncio.set_event_loop(loop)
+
+        if hasattr(app.state, "readpool") and app.state.readpool:
+            try:
+                app.state.readpool.close()
+            except Exception as e:
+                print(f"SnapStart: Error closing stale readpool: {e}")
+            app.state.readpool = None
+
+        if hasattr(app.state, "writepool") and app.state.writepool:
+            try:
+                app.state.writepool.close()
+            except Exception as e:
+                print(f"SnapStart: Error closing stale writepool: {e}")
+            app.state.writepool = None
+
+        loop.run_until_complete(
+            connect_to_db(
+                app,
+                postgres_settings=postgres_settings,
+                add_write_connection_pool=settings.enable_transaction,
+            )
+        )
+
+        _connection_initialized = True
+
+    except Exception as e:
+        print(f"SnapStart: Failed to initialize database connection: {e}")
+        raise
+
+    return {"statusCode": 200}
 
 
 @app.on_event("startup")
 async def startup_event() -> None:
     """Connect to database on startup."""
     await connect_to_db(
         app,
-        postgres_settings=PostgresSettings(),
+        postgres_settings=postgres_settings,
         add_write_connection_pool=settings.enable_transaction,
     )
 

infrastructure/handlers/vector_handler.py

@@ -8,6 +8,7 @@
 from eoapi.vector.app import app
 from eoapi.vector.config import PostgresSettings
 from mangum import Mangum
+from snapshot_restore_py import register_after_restore, register_before_snapshot
 from tipg.collections import register_collection_catalog
 from tipg.database import connect_to_db
 from tipg.settings import DatabaseSettings
@@ -26,6 +27,73 @@
     # We allow non-spatial tables
     spatial=False,
 )
+postgres_settings = PostgresSettings()
+
+_connection_initialized = False
+
+
+@register_before_snapshot
+def on_snapshot():
+    """
+    Runtime hook called by Lambda before taking a snapshot.
+    We close database connections that shouldn't be in the snapshot.
+    """
+
+    if hasattr(app, "state") and hasattr(app.state, "pool") and app.state.pool:
+        try:
+            app.state.pool.close()
+            app.state.pool = None
+        except Exception as e:
+            print(f"SnapStart: Error closing database pool: {e}")
+
+    return {"statusCode": 200}
+
+
+@register_after_restore
+def on_snap_restore():
+    """
+    Runtime hook called by Lambda after restoring from a snapshot.
+    We recreate database connections that were closed before the snapshot.
+    """
+    global _connection_initialized
+
+    try:
+        try:
+            loop = asyncio.get_running_loop()
+        except RuntimeError:
+            loop = asyncio.new_event_loop()
+            asyncio.set_event_loop(loop)
+
+        if hasattr(app.state, "pool") and app.state.pool:
+            try:
+                app.state.pool.close()
+            except Exception as e:
+                print(f"SnapStart: Error closing stale pool: {e}")
+            app.state.pool = None
+
+        loop.run_until_complete(
+            connect_to_db(
+                app,
+                schemas=["pgstac", "public"],
+                user_sql_files=list(CUSTOM_SQL_DIRECTORY.glob("*.sql")),  # type: ignore
+                settings=postgres_settings,
+            )
+        )
+
+        loop.run_until_complete(
+            register_collection_catalog(
+                app,
+                db_settings=db_settings,
+            )
+        )
+
+        _connection_initialized = True
+
+    except Exception as e:
+        print(f"SnapStart: Failed to initialize database connection: {e}")
+        raise
+
+    return {"statusCode": 200}
 
 
 @app.on_event("startup")
@@ -36,7 +104,7 @@ async def startup_event() -> None:
         # We enable both pgstac and public schemas (pgstac will be used by custom functions)
         schemas=["pgstac", "public"],
         user_sql_files=list(CUSTOM_SQL_DIRECTORY.glob("*.sql")),  # type: ignore
-        settings=PostgresSettings(),
+        settings=postgres_settings,
     )
     await register_collection_catalog(app, db_settings=db_settings)
 

pyproject.toml

@@ -9,7 +9,7 @@ dependencies = []
 [dependency-groups]
 deploy = [
     "boto3==1.24.15",
-    "eoapi-cdk==10.2.5",
+    "eoapi-cdk==10.3.0",
     "pydantic-settings[yaml]==2.2.1",
     "pydantic==2.7",
     "typing-extensions>=4.12.2",