diff --git a/skyvern/forge/sdk/api/aws.py b/skyvern/forge/sdk/api/aws.py index 58ba2a0a..278e3f1a 100644 --- a/skyvern/forge/sdk/api/aws.py +++ b/skyvern/forge/sdk/api/aws.py @@ -4,6 +4,7 @@ from urllib.parse import urlparse import aioboto3 import structlog +from types_boto3_ec2.client import EC2Client from types_boto3_ecs.client import ECSClient from types_boto3_s3.client import S3Client from types_boto3_secretsmanager.client import SecretsManagerClient @@ -29,6 +30,7 @@ class AWSClientType(StrEnum): S3 = "s3" SECRETS_MANAGER = "secretsmanager" ECS = "ecs" + EC2 = "ec2" class AsyncAWSClient: @@ -57,6 +59,9 @@ class AsyncAWSClient: def _s3_client(self) -> S3Client: return self.session.client(AWSClientType.S3, region_name=self.region_name, endpoint_url=self._endpoint_url) + def _ec2_client(self) -> EC2Client: + return self.session.client(AWSClientType.EC2, region_name=self.region_name, endpoint_url=self._endpoint_url) + def _create_tag_string(self, tags: dict[str, str]) -> str: return "&".join([f"{k}={v}" for k, v in tags.items()]) @@ -312,6 +317,12 @@ class AsyncAWSClient: async with self._ecs_client() as client: return await client.deregister_task_definition(taskDefinition=task_definition) + ###### EC2 ###### + async def describe_network_interfaces(self, network_interface_ids: list[str]) -> dict: + # https://boto3.amazonaws.com/v1/documentation/api/latest/reference/services/ec2/client/describe_network_interfaces.html + async with self._ec2_client() as client: + return await client.describe_network_interfaces(NetworkInterfaceIds=network_interface_ids) + class S3Uri: # From: https://stackoverflow.com/questions/42641315/s3-urls-get-bucket-name-and-path