Skip to content

Commit 163ba59

Browse files
hrodmnhrodmn
committed
feat: enable SnapStart on lambda functions
1 parent 616b890 commit 163ba59

6 files changed

Lines changed: 214 additions & 12 deletions

File tree

‎infrastructure/app.py‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -181,6 +181,7 @@ def __init__(
181181
"handler": "handler.handler",
182182
"runtime": aws_lambda.Runtime.PYTHON_3_12,
183183
},
184+
enable_snap_start=True,
184185
)
185186

186187
#######################################################################
@@ -228,6 +229,7 @@ def __init__(
228229
"handler": "handler.handler",
229230
"runtime": aws_lambda.Runtime.PYTHON_3_12,
230231
},
232+
enable_snap_start=True,
231233
)
232234

233235
#######################################################################
@@ -274,6 +276,7 @@ def __init__(
274276
"handler": "handler.handler",
275277
"runtime": aws_lambda.Runtime.PYTHON_3_12,
276278
},
279+
enable_snap_start=True,
277280
)
278281

279282
if app_config.stac_ingestor:

‎infrastructure/handlers/raster_handler.py‎

Lines changed: 55 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -7,16 +7,70 @@
77
from eoapi.raster.app import app
88
from eoapi.raster.config import PostgresSettings
99
from mangum import Mangum
10+
from snapshot_restore_py import register_after_restore, register_before_snapshot
1011
from titiler.pgstac.db import connect_to_db
1112

1213
logging.getLogger("mangum.lifespan").setLevel(logging.ERROR)
1314
logging.getLogger("mangum.http").setLevel(logging.ERROR)
1415

16+
postgres_settings = PostgresSettings()
17+
_connection_initialized = False
18+
19+
20+
@register_before_snapshot
21+
def on_snapshot():
22+
"""
23+
Runtime hook called by Lambda before taking a snapshot.
24+
We close database connections that shouldn't be in the snapshot.
25+
"""
26+
27+
if hasattr(app, "state") and hasattr(app.state, "dbpool") and app.state.dbpool:
28+
try:
29+
app.state.dbpool.close()
30+
app.state.dbpool = None
31+
except Exception as e:
32+
print(f"SnapStart: Error closing database pool: {e}")
33+
34+
return {"statusCode": 200}
35+
36+
37+
@register_after_restore
38+
def on_snap_restore():
39+
"""
40+
Runtime hook called by Lambda after restoring from a snapshot.
41+
We recreate database connections that were closed before the snapshot.
42+
"""
43+
global _connection_initialized
44+
45+
try:
46+
try:
47+
loop = asyncio.get_running_loop()
48+
except RuntimeError:
49+
loop = asyncio.new_event_loop()
50+
asyncio.set_event_loop(loop)
51+
52+
if hasattr(app.state, "dbpool") and app.state.dbpool:
53+
try:
54+
app.state.dbpool.close()
55+
except Exception as e:
56+
print(f"SnapStart: Error closing stale pool: {e}")
57+
app.state.dbpool = None
58+
59+
loop.run_until_complete(connect_to_db(app, settings=postgres_settings))
60+
61+
_connection_initialized = True
62+
63+
except Exception as e:
64+
print(f"SnapStart: Failed to initialize database connection: {e}")
65+
raise
66+
67+
return {"statusCode": 200}
68+
1569

1670
@app.on_event("startup")
1771
async def startup_event() -> None:
1872
"""Connect to database on startup."""
19-
await connect_to_db(app, settings=PostgresSettings())
73+
await connect_to_db(app, settings=postgres_settings)
2074

2175

2276
handler = Mangum(app, lifespan="off")

‎infrastructure/handlers/stac_handler.py‎

Lines changed: 77 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -7,13 +7,90 @@
77
from eoapi.stac.app import app
88
from eoapi.stac.config import PostgresSettings, Settings
99
from mangum import Mangum
10+
from snapshot_restore_py import register_after_restore, register_before_snapshot
1011
from stac_fastapi.pgstac.db import connect_to_db
1112

1213
logging.getLogger("mangum.lifespan").setLevel(logging.ERROR)
1314
logging.getLogger("mangum.http").setLevel(logging.ERROR)
1415

1516
settings = Settings()
1617

18+
_connection_initialized = False
19+
20+
21+
@register_before_snapshot
22+
def on_snapshot():
23+
"""
24+
Runtime hook called by Lambda before taking a snapshot.
25+
We close database connections that shouldn't be in the snapshot.
26+
"""
27+
28+
if hasattr(app, "state") and hasattr(app.state, "readpool") and app.state.readpool:
29+
try:
30+
app.state.readpool.close()
31+
app.state.readpool = None
32+
except Exception as e:
33+
print(f"SnapStart: Error closing database readpool: {e}")
34+
35+
if (
36+
hasattr(app, "state")
37+
and hasattr(app.state, "writepool")
38+
and app.state.writepool
39+
):
40+
try:
41+
app.state.writepool.close()
42+
app.state.writepool = None
43+
except Exception as e:
44+
print(f"SnapStart: Error closing database writepool: {e}")
45+
46+
return {"statusCode": 200}
47+
48+
49+
@register_after_restore
50+
def on_snap_restore():
51+
"""
52+
Runtime hook called by Lambda after restoring from a snapshot.
53+
We recreate database connections that were closed before the snapshot.
54+
"""
55+
global _connection_initialized
56+
57+
try:
58+
try:
59+
loop = asyncio.get_running_loop()
60+
except RuntimeError:
61+
loop = asyncio.new_event_loop()
62+
asyncio.set_event_loop(loop)
63+
64+
if hasattr(app.state, "readpool") and app.state.readpool:
65+
try:
66+
app.state.readpool.close()
67+
except Exception as e:
68+
print(f"SnapStart: Error closing stale readpool: {e}")
69+
app.state.readpool = None
70+
71+
if hasattr(app.state, "writepool") and app.state.writepool:
72+
try:
73+
app.state.writepool.close()
74+
except Exception as e:
75+
print(f"SnapStart: Error closing stale writepool: {e}")
76+
app.state.writepool = None
77+
78+
loop.run_until_complete(
79+
connect_to_db(
80+
app,
81+
postgres_settings=PostgresSettings(),
82+
add_write_connection_pool=settings.enable_transaction,
83+
)
84+
)
85+
86+
_connection_initialized = True
87+
88+
except Exception as e:
89+
print(f"SnapStart: Failed to initialize database connection: {e}")
90+
raise
91+
92+
return {"statusCode": 200}
93+
1794

1895
@app.on_event("startup")
1996
async def startup_event() -> None:

‎infrastructure/handlers/vector_handler.py‎

Lines changed: 69 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -8,6 +8,7 @@
88
from eoapi.vector.app import app
99
from eoapi.vector.config import PostgresSettings
1010
from mangum import Mangum
11+
from snapshot_restore_py import register_after_restore, register_before_snapshot
1112
from tipg.collections import register_collection_catalog
1213
from tipg.database import connect_to_db
1314
from tipg.settings import DatabaseSettings
@@ -26,6 +27,73 @@
2627
# We allow non-spatial tables
2728
spatial=False,
2829
)
30+
postgres_settings = PostgresSettings()
31+
32+
_connection_initialized = False
33+
34+
35+
@register_before_snapshot
36+
def on_snapshot():
37+
"""
38+
Runtime hook called by Lambda before taking a snapshot.
39+
We close database connections that shouldn't be in the snapshot.
40+
"""
41+
42+
if hasattr(app, "state") and hasattr(app.state, "pool") and app.state.pool:
43+
try:
44+
app.state.pool.close()
45+
app.state.pool = None
46+
except Exception as e:
47+
print(f"SnapStart: Error closing database pool: {e}")
48+
49+
return {"statusCode": 200}
50+
51+
52+
@register_after_restore
53+
def on_snap_restore():
54+
"""
55+
Runtime hook called by Lambda after restoring from a snapshot.
56+
We recreate database connections that were closed before the snapshot.
57+
"""
58+
global _connection_initialized
59+
60+
try:
61+
try:
62+
loop = asyncio.get_running_loop()
63+
except RuntimeError:
64+
loop = asyncio.new_event_loop()
65+
asyncio.set_event_loop(loop)
66+
67+
if hasattr(app.state, "pool") and app.state.pool:
68+
try:
69+
app.state.pool.close()
70+
except Exception as e:
71+
print(f"SnapStart: Error closing stale pool: {e}")
72+
app.state.pool = None
73+
74+
loop.run_until_complete(
75+
connect_to_db(
76+
app,
77+
schemas=["pgstac", "public"],
78+
user_sql_files=list(CUSTOM_SQL_DIRECTORY.glob("*.sql")), # type: ignore
79+
settings=postgres_settings,
80+
)
81+
)
82+
83+
loop.run_until_complete(
84+
register_collection_catalog(
85+
app,
86+
db_settings=db_settings,
87+
)
88+
)
89+
90+
_connection_initialized = True
91+
92+
except Exception as e:
93+
print(f"SnapStart: Failed to initialize database connection: {e}")
94+
raise
95+
96+
return {"statusCode": 200}
2997

3098

3199
@app.on_event("startup")
@@ -36,7 +104,7 @@ async def startup_event() -> None:
36104
# We enable both pgstac and public schemas (pgstac will be used by custom functions)
37105
schemas=["pgstac", "public"],
38106
user_sql_files=list(CUSTOM_SQL_DIRECTORY.glob("*.sql")), # type: ignore
39-
settings=PostgresSettings(),
107+
settings=postgres_settings,
40108
)
41109
await register_collection_catalog(app, db_settings=db_settings)
42110

‎pyproject.toml‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -9,7 +9,7 @@ dependencies = []
99
[dependency-groups]
1010
deploy = [
1111
"boto3==1.24.15",
12-
"eoapi-cdk==10.2.5",
12+
"eoapi-cdk==10.3.0",
1313
"pydantic-settings[yaml]==2.2.1",
1414
"pydantic==2.7",
1515
"typing-extensions>=4.12.2",

‎uv.lock‎

Lines changed: 9 additions & 9 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

0 commit comments

Comments
 (0)