fix: add pgbouncer secretBootstrapper as a dependency for lambda functions (#221)

hrodmn · web-flow · commit 6b141e93e0c2 · 2025-12-19T10:24:54.000-06:00
diff --git a/lib/stac-api/index.ts b/lib/stac-api/index.ts
@@ -12,7 +12,12 @@ import {
 import { Construct } from "constructs";
 import * as path from "path";
 import { LambdaApiGateway } from "../lambda-api-gateway";
-import { CustomLambdaFunctionProps, resolveLambdaCode } from "../utils";
+import {
+  CustomLambdaFunctionProps,
+  resolveLambdaCode,
+  extractDatabaseDependencies,
+  createLambdaVersionWithDependencies,
+} from "../utils";
 
 export const EXTENSIONS = {
   QUERY: "query",
@@ -204,10 +209,23 @@ export class PgStacApiLambda extends Construct {
     });
     this.stacApiLambdaFunction = this.lambdaFunction = runtime.lambdaFunction;
 
+    // Determine which lambda to use for API Gateway
+    let apiLambda: lambda.Function | lambda.Version;
+    if (props.enableSnapStart) {
+      // Extract dependencies from database if it's a PgStacDatabase with PgBouncer
+      const dbDependencies = extractDatabaseDependencies(props.db);
+
+      // Create version with dependencies to ensure snapshot creation waits
+      apiLambda = createLambdaVersionWithDependencies(
+        runtime.lambdaFunction,
+        dbDependencies,
+      );
+    } else {
+      apiLambda = runtime.lambdaFunction;
+    }
+
     const { api } = new LambdaApiGateway(this, "stac-api", {
-      lambdaFunction: props.enableSnapStart!
-        ? runtime.lambdaFunction.currentVersion
-        : runtime.lambdaFunction,
+      lambdaFunction: apiLambda,
       domainName: props.domainName ?? props.stacApiDomainName,
     });
 
diff --git a/lib/tipg-api/index.ts b/lib/tipg-api/index.ts
@@ -12,7 +12,12 @@ import {
 import { Construct } from "constructs";
 import * as path from "path";
 import { LambdaApiGateway } from "../lambda-api-gateway";
-import { CustomLambdaFunctionProps, resolveLambdaCode } from "../utils";
+import {
+  CustomLambdaFunctionProps,
+  resolveLambdaCode,
+  extractDatabaseDependencies,
+  createLambdaVersionWithDependencies,
+} from "../utils";
 
 export class TiPgApiLambdaRuntime extends Construct {
   public readonly lambdaFunction: lambda.Function;
@@ -146,10 +151,23 @@ export class TiPgApiLambda extends Construct {
     });
     this.tiPgLambdaFunction = this.lambdaFunction = runtime.lambdaFunction;
 
+    // Determine which lambda to use for API Gateway
+    let apiLambda: lambda.Function | lambda.Version;
+    if (props.enableSnapStart) {
+      // Extract dependencies from database if it's a PgStacDatabase with PgBouncer
+      const dbDependencies = extractDatabaseDependencies(props.db);
+
+      // Create version with dependencies to ensure snapshot creation waits
+      apiLambda = createLambdaVersionWithDependencies(
+        runtime.lambdaFunction,
+        dbDependencies,
+      );
+    } else {
+      apiLambda = runtime.lambdaFunction;
+    }
+
     const { api } = new LambdaApiGateway(this, "api", {
-      lambdaFunction: props.enableSnapStart!
-        ? runtime.lambdaFunction.currentVersion
-        : runtime.lambdaFunction,
+      lambdaFunction: apiLambda,
       domainName: props.domainName ?? props.tipgApiDomainName,
     });
 
diff --git a/lib/titiler-pgstac-api/index.ts b/lib/titiler-pgstac-api/index.ts
@@ -13,7 +13,12 @@ import {
 import { Construct } from "constructs";
 import * as path from "path";
 import { LambdaApiGateway } from "../lambda-api-gateway";
-import { CustomLambdaFunctionProps, resolveLambdaCode } from "../utils";
+import {
+  CustomLambdaFunctionProps,
+  resolveLambdaCode,
+  extractDatabaseDependencies,
+  createLambdaVersionWithDependencies,
+} from "../utils";
 
 // default settings that can be overridden by the user-provided environment.
 let defaultTitilerPgstacEnv: Record<string, string> = {
@@ -190,10 +195,23 @@ export class TitilerPgstacApiLambda extends Construct {
     this.titilerPgstacLambdaFunction = this.lambdaFunction =
       runtime.lambdaFunction;
 
+    // Determine which lambda to use for API Gateway
+    let apiLambda: lambda.Function | lambda.Version;
+    if (props.enableSnapStart) {
+      // Extract dependencies from database if it's a PgStacDatabase with PgBouncer
+      const dbDependencies = extractDatabaseDependencies(props.db);
+
+      // Create version with dependencies to ensure snapshot creation waits
+      apiLambda = createLambdaVersionWithDependencies(
+        runtime.lambdaFunction,
+        dbDependencies,
+      );
+    } else {
+      apiLambda = runtime.lambdaFunction;
+    }
+
     const { api } = new LambdaApiGateway(this, "titlier-pgstac-api", {
-      lambdaFunction: props.enableSnapStart
-        ? runtime.lambdaFunction.currentVersion
-        : runtime.lambdaFunction,
+      lambdaFunction: apiLambda,
       domainName: props.domainName ?? props.titilerPgstacApiDomainName,
     });
 
diff --git a/lib/utils/index.ts b/lib/utils/index.ts
@@ -1,4 +1,4 @@
-import { aws_lambda as lambda } from "aws-cdk-lib";
+import { aws_lambda as lambda, CustomResource } from "aws-cdk-lib";
 
 export type CustomLambdaFunctionProps = lambda.FunctionProps | any;
 export const DEFAULT_PGSTAC_VERSION = "0.9.5";
@@ -15,17 +15,88 @@ export const DEFAULT_PGSTAC_VERSION = "0.9.5";
 export function resolveLambdaCode(
   userCode?: lambda.Code,
   dockerBuildPath?: string,
-  dockerBuildOptions?: lambda.DockerBuildAssetOptions
+  dockerBuildOptions?: lambda.DockerBuildAssetOptions,
 ): lambda.Code {
   if (userCode) {
     return userCode;
   }
 
   if (!dockerBuildPath || !dockerBuildOptions) {
     throw new Error(
-      "dockerBuildPath and dockerBuildOptions are required when no custom code is provided"
+      "dockerBuildPath and dockerBuildOptions are required when no custom code is provided",
     );
   }
 
   return lambda.Code.fromDockerBuild(dockerBuildPath, dockerBuildOptions);
 }
+
+/**
+ * Type guard to check if a value has the secretBootstrapper property.
+ *
+ * This is used to detect PgStacDatabase constructs with PgBouncer enabled,
+ * which expose a secretBootstrapper CustomResource that completes after
+ * the database secret is updated with PgBouncer connection information.
+ *
+ * @param value - Value to check
+ * @returns true if value has a defined secretBootstrapper property
+ */
+function hasSecretBootstrapper(
+  value: any,
+): value is { secretBootstrapper: CustomResource } {
+  return (
+    value &&
+    typeof value === "object" &&
+    "secretBootstrapper" in value &&
+    value.secretBootstrapper !== undefined
+  );
+}
+
+/**
+ * Extracts database dependencies from a database construct if available.
+ *
+ * For PgStacDatabase with PgBouncer enabled, returns the secretBootstrapper
+ * CustomResource which ensures the database secret is fully initialized before
+ * dependent resources are created.
+ *
+ * This is critical for SnapStart Lambda functions, as the snapshot must not
+ * be created until the secret contains the correct connection information.
+ *
+ * @param db - Database construct (may be PgStacDatabase, IDatabaseInstance, etc.)
+ * @returns Array of CustomResources to depend on (empty if none available)
+ */
+export function extractDatabaseDependencies(db: any): CustomResource[] {
+  if (hasSecretBootstrapper(db)) {
+    return [db.secretBootstrapper];
+  }
+  return [];
+}
+
+/**
+ * Creates a Lambda version with dependencies for SnapStart.
+ *
+ * When SnapStart is enabled, the version creation triggers snapshot creation.
+ * Dependencies ensure the snapshot isn't created until prerequisites are met,
+ * such as database secrets being fully initialized.
+ *
+ * This prevents race conditions where the SnapStart snapshot might capture
+ * incorrect or incomplete configuration.
+ *
+ * @param lambdaFunction - Lambda function to create a version from
+ * @param dependencies - Array of resources to depend on (pass empty array if none)
+ * @returns Lambda version with dependencies applied
+ */
+export function createLambdaVersionWithDependencies(
+  lambdaFunction: lambda.Function,
+  dependencies: CustomResource[],
+): lambda.Version {
+  const version = lambdaFunction.currentVersion;
+
+  // Add dependencies to the version resource to ensure proper ordering
+  if (dependencies && dependencies.length > 0) {
+    dependencies.forEach((dep) => {
+      version.node.addDependency(dep);
+    });
+  }
+
+  return version;
+}
