Use pgbouncer's public IP if it is in a public subnet

Author: hrodmn

integration_tests/cdk/app.py

@@ -83,18 +83,21 @@ def __init__(
             removal_policy=RemovalPolicy.DESTROY,
         )
 
-        pgstac_db.db.connections.allow_default_port_from_any_ipv4()
+        assert pgstac_db.security_group
+
+        pgstac_db.security_group.add_ingress_rule(
+            aws_ec2.Peer.any_ipv4(), aws_ec2.Port.tcp(5432)
+        )
 
         PgStacApiLambda(
             self,
             "pgstac-api",
+            connection_target=pgstac_db.connection_target,
+            db_secret=pgstac_db.pgstac_secret,
             api_env={
                 "NAME": app_config.build_service_name("STAC API"),
                 "description": f"{app_config.stage} STAC API",
             },
-            vpc=vpc,
-            connection_target=pgstac_db.connection_target,
-            db_secret=pgstac_db.pgstac_secret,
         )
 
         TitilerPgstacApiLambda(
@@ -104,7 +107,6 @@ def __init__(
                 "NAME": app_config.build_service_name("titiler pgSTAC API"),
                 "description": f"{app_config.stage} titiler pgstac API",
             },
-            vpc=vpc,
             connection_target=pgstac_db.connection_target,
             db_secret=pgstac_db.pgstac_secret,
             buckets=[],
@@ -116,7 +118,6 @@ def __init__(
         TiPgApiLambda(
             self,
             "tipg-api",
-            vpc=vpc,
             connection_target=pgstac_db.connection_target,
             db_secret=pgstac_db.pgstac_secret,
             api_env={

lib/database/PgBouncer.ts

@@ -70,6 +70,7 @@ export interface PgBouncerProps {
 export class PgBouncer extends Construct {
   public readonly instance: ec2.Instance;
   public readonly pgbouncerSecret: secretsmanager.Secret;
+  public readonly securityGroup: ec2.SecurityGroup;
 
   // The max_connections parameter in PgBouncer determines the maximum number of
   // connections to open on the actual database instance. We want that number to
@@ -134,6 +135,13 @@ export class PgBouncer extends Construct {
       })
     );
 
+    // Create a security group and allow connections from the Lambda IP ranges for this region
+    this.securityGroup = new ec2.SecurityGroup(this, "PgBouncerSecurityGroup", {
+      vpc: props.vpc,
+      description: "Security group for PgBouncer instance",
+      allowAllOutbound: true,
+    });
+
     // Create PgBouncer instance
     this.instance = new ec2.Instance(this, "Instance", {
       vpc: props.vpc,
@@ -142,6 +150,7 @@ export class PgBouncer extends Construct {
           ? ec2.SubnetType.PUBLIC
           : ec2.SubnetType.PRIVATE_WITH_EGRESS,
       },
+      securityGroup: this.securityGroup,
       instanceType,
       instanceName: props.instanceName,
       machineImage: ec2.MachineImage.fromSsmParameter(
@@ -161,6 +170,7 @@ export class PgBouncer extends Construct {
       ],
       userData: this.loadUserDataScript(pgBouncerConfig, props.database),
       userDataCausesReplacement: true,
+      associatePublicIpAddress: props.usePublicSubnet,
     });
 
     // Allow PgBouncer to connect to RDS
@@ -201,7 +211,9 @@ export class PgBouncer extends Construct {
     new CustomResource(this, "pgbouncerSecretBootstrapper", {
       serviceToken: secretUpdaterFn.functionArn,
       properties: {
-        instanceIp: this.instance.instancePrivateIp,
+        instanceIp: props.usePublicSubnet
+          ? this.instance.instancePublicIp
+          : this.instance.instancePrivateIp,
       },
     });
   }

lib/database/index.ts

@@ -36,7 +36,9 @@ export class PgStacDatabase extends Construct {
   db: rds.DatabaseInstance;
   pgstacSecret: secretsmanager.ISecret;
   private _pgBouncerServer?: PgBouncer;
+
   public readonly connectionTarget: rds.IDatabaseInstance | ec2.Instance;
+  public readonly securityGroup?: ec2.SecurityGroup;
 
   constructor(scope: Construct, id: string, props: PgStacDatabaseProps) {
     super(scope, id);
@@ -153,7 +155,9 @@ export class PgStacDatabase extends Construct {
           secret: this.pgstacSecret,
         },
         dbMaxConnections: parseInt(defaultParameters.maxConnections),
-        usePublicSubnet: true,
+        usePublicSubnet:
+          !props.vpcSubnets ||
+          props.vpcSubnets.subnetType === ec2.SubnetType.PUBLIC,
         pgBouncerConfig: {
           poolMode: "transaction",
           maxClientConn: 1000,
@@ -168,6 +172,7 @@ export class PgStacDatabase extends Construct {
 
       this.pgstacSecret = this._pgBouncerServer.pgbouncerSecret;
       this.connectionTarget = this._pgBouncerServer.instance;
+      this.securityGroup = this._pgBouncerServer.securityGroup;
     } else {
       this.connectionTarget = this.db;
     }

lib/stac-api/index.ts

@@ -20,8 +20,6 @@ export class PgStacApiLambda extends Construct {
   constructor(scope: Construct, id: string, props: PgStacApiLambdaProps) {
     super(scope, id);
 
-    console.log(props)
-    console.log(props.lambdaFunctionOptions);
     this.stacApiLambdaFunction = new lambda.Function(this, "lambda", {
       // defaults
       runtime: lambda.Runtime.PYTHON_3_11,
@@ -47,9 +45,11 @@ export class PgStacApiLambda extends Construct {
     });
 
     props.dbSecret.grantRead(this.stacApiLambdaFunction);
+
     if (props.vpc){
-      this.stacApiLambdaFunction.connections.allowTo(props.connectionTarget, ec2.Port.tcp(5432));
+      this.stacApiLambdaFunction.connections.allowTo(props.connectionTarget, ec2.Port.tcp(5432), "allow connections from stac-fastapi-pgstac");
     }
+
     const stacApi = new HttpApi(this, `${Stack.of(this).stackName}-stac-api`, {
       defaultDomainMapping: props.stacApiDomainName ? {
         domainName: props.stacApiDomainName

lib/tipg-api/index.ts

@@ -1,9 +1,9 @@
 import {
     Stack,
     aws_ec2 as ec2,
-    aws_rds as rds,
     aws_lambda as lambda,
     aws_logs as logs,
+    aws_rds as rds,
     aws_secretsmanager as secretsmanager,
     CfnOutput,
     Duration,
@@ -45,9 +45,11 @@ import {
       });
 
       props.dbSecret.grantRead(this.tiPgLambdaFunction);
+
       if (props.vpc){
         this.tiPgLambdaFunction.connections.allowTo(props.connectionTarget, ec2.Port.tcp(5432), "allow connections from tipg");
       }
+
       const tipgApi = new HttpApi(this, `${Stack.of(this).stackName}-tipg-api`, {
         defaultDomainMapping: props.tipgApiDomainName ? {
           domainName: props.tipgApiDomainName
@@ -92,7 +94,6 @@ import {
      */
     readonly dbSecret: secretsmanager.ISecret;
 
-
     /**
      * Customized environment variables to send to titiler-pgstac runtime.
      */

lib/titiler-pgstac-api/index.ts

@@ -103,10 +103,10 @@ import { CustomLambdaFunctionProps } from "../utils";
      */
     readonly vpc?: ec2.IVpc;
 
-  /**
-   * RDS Instance with installed pgSTAC or pgbouncer server.
-   */
-  readonly connectionTarget: rds.IDatabaseInstance | ec2.IInstance;
+    /**
+     * RDS Instance with installed pgSTAC or pgbouncer server.
+     */
+    readonly connectionTarget: rds.IDatabaseInstance | ec2.IInstance;
 
     /**
      * Subnet into which the lambda should be deployed.