feat(vpc): allow Vpc.fromLookup() to discover asymmetric subnets

Previously, Vpc.fromLookup() required every subnet group to be in the same Availability Zones,
and for all subnets in all groups to cover all of those Availability Zones.
This loosens this requirement, so that now Vpc.fromLookup() works for VPCs of all shapes.

This also add a way to customize which tag of the subnet is considered when grouping subnets into groups when calling fromLookup().

Fixes aws#3407

packages/@aws-cdk/aws-ec2/lib/vpc-lookup.ts

@@ -38,4 +38,15 @@ export interface VpcLookupOptions {
    * @default Don't care whether we return the default VPC
    */
   readonly isDefault?: boolean;
+
+  /**
+   * Optional tag whose value is used as the name of the subnet group
+   * the subnet with the tag belongs to.
+   * If not provided, we'll look at the aws-cdk:subnet-name tag.
+   * If the subnet does not have the tag,
+   * we'll use its type as the name.
+   *
+   * @default aws-cdk:subnet-name
+   */
+  readonly subnetGroupNameTag?: string;
 }

packages/@aws-cdk/aws-ec2/lib/vpc.ts

@@ -801,13 +801,16 @@ export class Vpc extends VpcBase {
       filter.isDefault = options.isDefault ? 'true' : 'false';
     }
 
-    const attributes = ContextProvider.getValue(scope, {
+    const attributes: cxapi.VpcContextResponse = ContextProvider.getValue(scope, {
       provider: cxapi.VPC_PROVIDER,
-      props: { filter } as cxapi.VpcContextQuery,
-      dummyValue: undefined
+      props: {
+        filter,
+        subnetGroupNameTag: options.subnetGroupNameTag,
+      } as cxapi.VpcContextQuery,
+      dummyValue: undefined,
     }).value;
 
-    return new ImportedVpc(scope, id, attributes || DUMMY_VPC_PROPS, attributes === undefined);
+    return new LookedUpVpc(scope, id, attributes || DUMMY_VPC_PROPS, attributes === undefined);
 
     /**
      * Prefixes all keys in the argument with `tag:`.`
@@ -1446,6 +1449,60 @@ class ImportedVpc extends VpcBase {
   }
 }
 
+class LookedUpVpc extends VpcBase {
+  public readonly vpcId: string;
+  public readonly vpnGatewayId?: string;
+  public readonly internetConnectivityEstablished: IDependable = new ConcreteDependable();
+  public readonly availabilityZones: string[];
+  public readonly publicSubnets: ISubnet[];
+  public readonly privateSubnets: ISubnet[];
+  public readonly isolatedSubnets: ISubnet[];
+
+  constructor(scope: Construct, id: string, props: cxapi.VpcContextResponse, isIncomplete: boolean) {
+    super(scope, id);
+
+    this.vpcId = props.vpcId;
+    this.vpnGatewayId = props.vpnGatewayId;
+    this.incompleteSubnetDefinition = isIncomplete;
+
+    const availabilityZones = Array.from(new Set<string>(flatMap(props.subnetGroups, subnetGroup => {
+      return subnetGroup.subnets.map(subnet => subnet.availabilityZone);
+    })));
+    availabilityZones.sort((az1, az2) => az1.localeCompare(az2));
+    this.availabilityZones = availabilityZones;
+
+    this.publicSubnets = this.extractSubnetsOfType(props.subnetGroups, cxapi.VpcSubnetGroupType.PUBLIC);
+    this.privateSubnets = this.extractSubnetsOfType(props.subnetGroups, cxapi.VpcSubnetGroupType.PRIVATE);
+    this.isolatedSubnets = this.extractSubnetsOfType(props.subnetGroups, cxapi.VpcSubnetGroupType.ISOLATED);
+  }
+
+  private extractSubnetsOfType(subnetGroups: cxapi.VpcSubnetGroup[], subnetGroupType: cxapi.VpcSubnetGroupType): ISubnet[] {
+    return flatMap(subnetGroups.filter(subnetGroup => subnetGroup.type === subnetGroupType),
+        subnetGroup => this.subnetGroupToSubnets(subnetGroup));
+  }
+
+  private subnetGroupToSubnets(subnetGroup: cxapi.VpcSubnetGroup): ISubnet[] {
+    const ret = new Array<ISubnet>();
+    for (let i = 0; i < subnetGroup.subnets.length; i++) {
+      const vpcSubnet = subnetGroup.subnets[i];
+      ret.push(Subnet.fromSubnetAttributes(this, `${subnetGroup.name}Subnet${i + 1}`, {
+        availabilityZone: vpcSubnet.availabilityZone,
+        subnetId: vpcSubnet.subnetId,
+        routeTableId: vpcSubnet.routeTableId,
+      }));
+    }
+    return ret;
+  }
+}
+
+function flatMap<T, U>(xs: T[], fn: (x: T) => U[]): U[] {
+  const ret = new Array<U>();
+  for (const x of xs) {
+    ret.push(...fn(x));
+  }
+  return ret;
+}
+
 /**
  * If the placement strategy is completely "default", reify the defaults so
  * consuming code doesn't have to reimplement the same analysis every time.
@@ -1544,10 +1601,39 @@ class ImportedSubnet extends Resource implements ISubnet, IPublicSubnet, IPrivat
  * It's only used for testing and on the first run-through.
  */
 const DUMMY_VPC_PROPS: cxapi.VpcContextResponse = {
-  availabilityZones: ['dummy-1a', 'dummy-1b'],
+  subnetGroups: [
+    {
+      name: 'Public',
+      type: cxapi.VpcSubnetGroupType.PUBLIC,
+      subnets: [
+        {
+          availabilityZone: 'dummy-1a',
+          subnetId: 's-12345',
+          routeTableId: 'rtb-12345s',
+        },
+        {
+          availabilityZone: 'dummy-1b',
+          subnetId: 's-67890',
+          routeTableId: 'rtb-67890s',
+        },
+      ],
+    },
+    {
+      name: 'Private',
+      type: cxapi.VpcSubnetGroupType.PRIVATE,
+      subnets: [
+        {
+          availabilityZone: 'dummy-1a',
+          subnetId: 'p-12345',
+          routeTableId: 'rtb-12345p',
+        },
+        {
+          availabilityZone: 'dummy-1b',
+          subnetId: 'p-67890',
+          routeTableId: 'rtb-57890p',
+        },
+      ],
+    },
+  ],
   vpcId: 'vpc-12345',
-  publicSubnetIds: ['s-12345', 's-67890'],
-  publicSubnetRouteTableIds: ['rtb-12345s', 'rtb-67890s'],
-  privateSubnetIds: ['p-12345', 'p-67890'],
-  privateSubnetRouteTableIds: ['rtb-12345p', 'rtb-57890p'],
 };

packages/@aws-cdk/aws-ec2/test/test.vpc.from-lookup.ts

@@ -0,0 +1,114 @@
+import { Construct, ContextProvider, GetContextValueOptions, GetContextValueResult, Lazy, Stack } from "@aws-cdk/core";
+import cxapi = require('@aws-cdk/cx-api');
+import { Test } from 'nodeunit';
+import { Vpc } from "../lib";
+
+export = {
+  'Vpc.fromLookup()': {
+    'requires concrete values'(test: Test) {
+      // GIVEN
+      const stack = new Stack();
+
+      test.throws(() => {
+        Vpc.fromLookup(stack, 'Vpc', {
+          vpcId: Lazy.stringValue({ produce: () => 'some-id' })
+        });
+
+      }, 'All arguments to Vpc.fromLookup() must be concrete');
+
+      test.done();
+    },
+
+    'selecting subnets by name from a looked-up VPC does not throw'(test: Test) {
+      // GIVEN
+      const stack = new Stack(undefined, undefined, { env: { region: 'us-east-1', account: '123456789012' }});
+      const vpc = Vpc.fromLookup(stack, 'VPC', {
+        vpcId: 'vpc-1234'
+      });
+
+      // WHEN
+      vpc.selectSubnets({ subnetName: 'Bleep' });
+
+      // THEN: no exception
+
+      test.done();
+    },
+
+    'accepts asymmetric subnets'(test: Test) {
+      const previous = mockVpcContextProviderWith({
+        vpcId: 'vpc-1234',
+        subnetGroups: [
+          {
+            name: 'Public',
+            type: cxapi.VpcSubnetGroupType.PUBLIC,
+            subnets: [
+              {
+                subnetId: 'pub-sub-in-us-east-1a',
+                availabilityZone: 'us-east-1a',
+                routeTableId: 'rt-123',
+              },
+              {
+                subnetId: 'pub-sub-in-us-east-1b',
+                availabilityZone: 'us-east-1b',
+                routeTableId: 'rt-123',
+              },
+            ],
+          },
+          {
+            name: 'Private',
+            type: cxapi.VpcSubnetGroupType.PRIVATE,
+            subnets: [
+              {
+                subnetId: 'pri-sub-1-in-us-east-1c',
+                availabilityZone: 'us-east-1c',
+                routeTableId: 'rt-123',
+              },
+              {
+                subnetId: 'pri-sub-2-in-us-east-1c',
+                availabilityZone: 'us-east-1c',
+                routeTableId: 'rt-123',
+              },
+              {
+                subnetId: 'pri-sub-1-in-us-east-1d',
+                availabilityZone: 'us-east-1d',
+                routeTableId: 'rt-123',
+              },
+              {
+                subnetId: 'pri-sub-2-in-us-east-1d',
+                availabilityZone: 'us-east-1d',
+                routeTableId: 'rt-123',
+              },
+            ],
+          },
+        ],
+      });
+
+      const stack = new Stack();
+      const vpc = Vpc.fromLookup(stack, 'Vpc', {
+        isDefault: true,
+      });
+
+      test.deepEqual(vpc.availabilityZones, ['us-east-1a', 'us-east-1b', 'us-east-1c', 'us-east-1d']);
+      test.equal(vpc.publicSubnets.length, 2);
+      test.equal(vpc.privateSubnets.length, 4);
+      test.equal(vpc.isolatedSubnets.length, 0);
+
+      restoreContextProvider(previous);
+      test.done();
+    },
+  },
+};
+
+function mockVpcContextProviderWith(response: cxapi.VpcContextResponse) {
+  const previous = ContextProvider.getValue;
+  ContextProvider.getValue = (_scope: Construct, _options: GetContextValueOptions) => {
+    return {
+      value: response,
+    };
+  };
+  return previous;
+}
+
+function restoreContextProvider(previous: (scope: Construct, options: GetContextValueOptions) => GetContextValueResult): void {
+  ContextProvider.getValue = previous;
+}

packages/@aws-cdk/aws-ec2/test/test.vpc.ts

@@ -848,35 +848,6 @@ export = {
       test.done();
     }
   },
-
-  'fromLookup() requires concrete values'(test: Test) {
-    // GIVEN
-    const stack = new Stack();
-
-    test.throws(() => {
-      Vpc.fromLookup(stack, 'Vpc', {
-        vpcId: Lazy.stringValue({ produce: () => 'some-id' })
-      });
-
-    }, 'All arguments to Vpc.fromLookup() must be concrete');
-
-    test.done();
-  },
-
-  'selecting subnets by name from a looked-up VPC does not throw'(test: Test) {
-    // GIVEN
-    const stack = new Stack(undefined, undefined, { env: { region: 'us-east-1', account: '123456789012' }});
-    const vpc = Vpc.fromLookup(stack, 'VPC', {
-      vpcId: 'vpc-1234'
-    });
-
-    // WHEN
-    vpc.selectSubnets({ subnetName: 'Bleep' });
-
-    // THEN: no exception
-
-    test.done();
-  },
 };
 
 function getTestStack(): Stack {