How to get unique AWS subnet IDs per AZ using data source in terraform?

The VPC needs to attach with the transit gateway, the code:

resource "aws_ec2_transit_gateway_vpc_attachment" "vpc_live" {
  subnet_ids         = toset(data.aws_subnets.vpc_live.ids)
  transit_gateway_id = var.tgw_id
  vpc_id             = var.vpc_id
}

data source:

data "aws_subnets" "vpc_live" {
  filter {
    name   = "vpc-id"
    values = [var.vpc_id]
  }
}

Error:

Error: updating EC2 Transit Gateway VPC Attachment (tgw-attach-xxxxxxxxxxxxxxx): InvalidParameterValue: AddSubnets values count greater than region availability zones
│   status code: 400, request id: dnc3c5d-8927-029c-9fd5-311c3c5b5046
│
│   with aws_ec2_transit_gateway_vpc_attachment.vpc_live,
│   on main.tf line 7, in resource "aws_ec2_transit_gateway_vpc_attachment" "vpc_live":
│    7: resource "aws_ec2_transit_gateway_vpc_attachment" "vpc_live" {

I understand it requires only one subnet-id per availability zone (AZ). I have been searching everywhere in their GitHub issues, StackOverflow, but I am wondering if nobody has faced this primary issue.

So my question is, how can I filter my data source in such a way that it returns only 1 subnet ID per AZ?

Solution

Well, that's quite easy:

First, get all the availability zones attached to each subnet:

data "aws_subnet" "vpc_live" {
  for_each = toset(data.aws_subnets.vpc_live.ids)
  id       = each.key
}

Second, do the reverse. Get all the subnets belongs to those AZs:

{ for s in data.aws_subnet.vpc_live : s.availability_zone => s.id... }

Finally, pick the first subnet on the index from all AZs:

resource "aws_ec2_transit_gateway_vpc_attachment" "vpc_live" {
  subnet_ids         = [for subnet_ids in { for s in data.aws_subnet.ids : s.availability_zone => s.id... } : subnet_ids[0]]
  transit_gateway_id = var.tgw_id
  vpc_id             = var.vpc_id
}

I ran the code on my system, and check the output:

subnets_ids = [
    + "subnet-xxxx123",
    + "subnet-xxxx456",
    + "subnet-xxxx789",
  ]