snapshot-restore-py

import os from snapshot_restore import register_before_snapshot, register_after_restore # Global state that will be snapshotted config = None db_connection = None def load_config(): """Simulate loading configuration from environment or Parameter Store.""" global config if config is None: print("Loading config...") # In a real scenario, fetch config from a secure source config = os.environ.get('MY_APP_CONFIG', 'default_config_value') print(f"Config loaded: {config}") return config def connect_to_db(): """Simulate connecting to a database.""" global db_connection if db_connection is None: print("Connecting to database...") # In a real scenario, use actual DB connection, e.g., psycopg2.connect db_connection = f"DB_Connection_to_{os.environ.get('DB_HOST', 'localhost')}" print(f"Database connected: {db_connection}") return db_connection @register_before_snapshot def before_snapshot_hook(): """Hook to run before Lambda takes a snapshot. Use this to clean up resources that cannot be snapshotted (e.g., open network connections). """ global db_connection print("Executing before_snapshot_hook: Closing DB connection.") if db_connection: # Proper closing of connection in a real app db_connection = None # Simulate closing @register_after_restore def after_restore_hook(): """Hook to run after Lambda restores from a snapshot. Use this to re-establish connections or refresh volatile data. """ global config, db_connection print("Executing after_restore_hook: Re-establishing resources.") # Refresh config if it contains time-sensitive data load_config() # Re-establish DB connection connect_to_db() def lambda_handler(event, context): """Main Lambda handler function.""" load_config() db = connect_to_db() print(f"Handler executed with config: {config} and db: {db}") return { 'statusCode': 200, 'body': 'Function executed successfully!' } # Example of local testing (not part of actual Lambda deployment process) if __name__ == "__main__": os.environ['MY_APP_CONFIG'] = 'local_test_config' os.environ['DB_HOST'] = 'test_db_instance' print("--- First 'invocation' (initialization) ---") lambda_handler({}, {}) print("\n--- Simulating Snapshot and Restore (local only, no actual snapshot) ---") before_snapshot_hook() after_restore_hook() print("\n--- Second 'invocation' (after restore) ---") lambda_handler({}, {})