8000 Add pro clients, allow setting of request params through a "proxy" (#… · codeperl/localstack@85a00cb · GitHub
[go: up one dir, main page]
Skip to content 8000

Commit 85a00cb

Browse files
dfangldominikschubert
dfangl
and
dominikschubert
authored
Add pro clients, allow setting of request params through a "proxy" (localstack#7923)
Co-authored-by: Dominik Schubert <dominik.schubert91@gmail.com>
1 parent b6a456b commit 85a00cb

File tree

3 files changed

+307
-87
lines changed

3 files changed

+307
-87
lines changed

‎localstack/aws/connect.py

Lines changed: 50 additions & 74 deletions
Original file line numberDiff line numberDiff line change
@@ -8,8 +8,8 @@
88
import logging
99
import threading
1010
from abc import ABC, abstractmethod
11-
from functools import cache
12-
from typing import TYPE_CHECKING, Any, Optional, TypedDict
11+
from functools import cache, partial
12+
from typing import Any, Callable, Generic, Optional, TypedDict, TypeVar
1313

1414
from boto3.session import Session
1515
from botocore.client import BaseClient
@@ -22,44 +22,9 @@
2222
MAX_POOL_CONNECTIONS,
2323
)
2424
from localstack.utils.aws.aws_stack import get_local_service_url
25+
from localstack.utils.aws.client_types import TypedServiceClientFactory
2526
from localstack.utils.aws.request_context import get_region_from_request_context
2627

27-
if TYPE_CHECKING:
28-
from mypy_boto3_acm import ACMClient
29-
from mypy_boto3_apigateway import APIGatewayClient
30-
from mypy_boto3_cloudformation import CloudFormationClient
31-
from mypy_boto3_cloudwatch import CloudWatchClient
32-
from mypy_boto3_cognito_idp import CognitoIdentityProviderClient
33-
from mypy_boto3_dynamodb import DynamoDBClient
34-
from mypy_boto3_dynamodbstreams import DynamoDBStreamsClient
35-
from mypy_boto3_ec2 import EC2Client
36-
from mypy_boto3_ecr import ECRClient
37-
from mypy_boto3_es import ElasticsearchServiceClient
38-
from mypy_boto3_events import EventBridgeClient
39-
from mypy_boto3_firehose import FirehoseClient
40-
from mypy_boto3_iam import IAMClient
41-
from mypy_boto3_kinesis import KinesisClient
42-
from mypy_boto3_kms import KMSClient
43-
from mypy_boto3_lambda import LambdaClient
44-
from mypy_boto3_logs import CloudWatchLogsClient
45-
from mypy_boto3_opensearch import OpenSearchServiceClient
46-
from mypy_boto3_redshift import RedshiftClient
47-
from mypy_boto3_resource_groups import ResourceGroupsClient
48-
from mypy_boto3_resourcegroupstaggingapi import ResourceGroupsTaggingAPIClient
49-
from mypy_boto3_route53 import Route53Client
50-
from mypy_boto3_route53resolver import Route53ResolverClient
51-
from mypy_boto3_s3 import S3Client
52-
from mypy_boto3_s3control import S3ControlClient
53-
from mypy_boto3_secretsmanager import SecretsManagerClient
54-
from mypy_boto3_ses import SESClient
55-
from mypy_boto3_sns import SNSClient
56-
from mypy_boto3_sqs import SQSClient
57-
from mypy_boto3_ssm import SSMClient
58-
from mypy_boto3_stepfunctions import SFNClient
59-
from mypy_boto3_sts import STSClient
60-
from mypy_boto3_transcribe import TranscribeClient
61-
62-
6328
LOG = logging.getLogger(__name__)
6429

6530

@@ -115,49 +80,58 @@ def load_dto(data: str) -> InternalRequestParameters:
11580
return json.loads(data)
11681

11782

83+
T = TypeVar("T")
84+
85+
86+
class MetadataRequestInjector(Generic[T]):
87+
def __init__(self, client: T, params: dict[str, str] | None = None):
88+
self._client = client
89+
self._params = params
90+
91+
def __getattr__(self, item):
92+
target = getattr(self._client, item)
93+
if not isinstance(target, Callable):
94+
return target
95+
if self._params:
96+
return partial(target, **self._params)
97+
else:
98+
return target
99+
100+
def request_metadata(
101+
self, source_arn: str | None = None, service_principal: str | None = None
102+
) -> T:
103+
"""
104+
Provides request metadata to this client.
105+
Identical to providing _ServicePrincipal and _SourceArn directly as operation arguments but typing
106+
compatible.
107+
108+
Raw example: lambda_client.invoke(FunctionName="fn", _SourceArn="...")
109+
Injector example: lambda_client.request_metadata(source_arn="...").invoke(FunctionName="fn")
110+
Cannot be called on objects where the parameters are already set.
111+
112+
:param source_arn: Arn on which behalf the calls of this client shall be made
113+
:param service_principal: Service principal on which behalf the calls of this client shall be made
114+
:return: A new version of the MetadataRequestInjector
115+
"""
116+
if self._params is not None:
117+
raise TypeError("Request_data cannot be called on it's own return value")
118+
params = {}
119+
if source_arn:
120+
params["_SourceArn"] = source_arn
121+
if service_principal:
122+
params["_ServicePrincipal"] = service_principal
123+
return MetadataRequestInjector(client=self._client, params=params)
124+
125+
118126
#
119127
# Factory
120128
#
121-
class ServiceLevelClientFactory:
129+
class ServiceLevelClientFactory(TypedServiceClientFactory):
122130
"""
123131
A service level client factory, preseeded with parameters for the boto3 client creation.
124132
Will create any service client with parameters already provided by the ClientFactory.
125133
"""
126134

127-
acm: "ACMClient"
128-
apigateway: "APIGatewayClient"
129-
awslambda: "LambdaClient"
130-
cloudformation: "CloudFormationClient"
131-
cloudwatch: "CloudWatchClient"
132-
cognito_idp: "CognitoIdentityProviderClient"
133-
dynamodb: "DynamoDBClient"
134-
dynamodbstreams: "DynamoDBStreamsClient"
135-
ec2: "EC2Client"
136-
ecr: "ECRClient"
137-
es: "ElasticsearchServiceClient"
138-
events: "EventBridgeClient"
139-
firehose: "FirehoseClient"
140-
iam: "IAMClient"
141-
kinesis: "KinesisClient"
142-
kms: "KMSClient"
143-
logs: "CloudWatchLogsClient"
144-
opensearch: "OpenSearchServiceClient"
145-
redshift: "RedshiftClient"
146-
resource_groups: "ResourceGroupsClient"
147-
resourcegroupstaggingapi: "ResourceGroupsTaggingAPIClient"
148-
route53: "Route53Client"
149-
route53resolver: "Route53ResolverClient"
150-
s3: "S3Client"
151-
s3control: "S3ControlClient"
152-
secretsmanager: "SecretsManagerClient"
153-
ses: "SESClient"
154-
sns: "SNSClient"
155-
sqs: "SQSClient"
156-
ssm: "SSMClient"
157-
stepfunctions: "SFNClient"
158-
sts: "STSClient"
159-
transcribe: "TranscribeClient"
160-
161135
def __init__(
162136
self, *, factory: "ClientFactory", client_creation_params: dict[str, str | Config | None]
163137
):
@@ -166,7 +140,9 @@ def __init__(
166140

167141
def __getattr__(self, service: str):
168142
service = attribute_name_to_service_name(service)
169-
return self._factory.get_client(service_name=service, **self._client_creation_params)
143+
return MetadataRequestInjector(
144+
client=self._factory.get_client(service_name=service, **self._client_creation_params)
145+
)
170146

171147

172148
class ClientFactory(ABC):

‎localstack/utils/aws/client_types.py

Lines changed: 194 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,194 @@
1+
import abc
2+
from typing import TYPE_CHECKING, Union
3+
4+
if TYPE_CHECKING:
5+
from mypy_boto3_acm import ACMClient
6+
from mypy_boto3_amplify import AmplifyClient
 F438
7+
from mypy_boto3_apigateway import APIGatewayClient
8+
from mypy_boto3_apigatewayv2 import ApiGatewayV2Client
9+
from mypy_boto3_appconfig import AppConfigClient
10+
from mypy_boto3_appsync import AppSyncClient
11+
from mypy_boto3_athena import AthenaClient
12+
from mypy_boto3_autoscaling import AutoScalingClient
13+
from mypy_boto3_backup import BackupClient
14+
from mypy_boto3_batch import BatchClient
15+
from mypy_boto3_ce import CostExplorerClient
16+
from mypy_boto3_cloudformation import CloudFormationClient
17+
from mypy_boto3_cloudfront import CloudFrontClient
18+
from mypy_boto3_cloudtrail import CloudTrailClient
19+
from mypy_boto3_cloudwatch import CloudWatchClient
20+
from mypy_boto3_codecommit import CodeCommitClient
21+
from mypy_boto3_cognito_identity import CognitoIdentityClient
22+
from mypy_boto3_cognito_idp import CognitoIdentityProviderClient
23+
from mypy_boto3_docdb import DocDBClient
24+
from mypy_boto3_dynamodb import DynamoDBClient
25+
from mypy_boto3_dynamodbstreams import DynamoDBStreamsClient
26+
from mypy_boto3_ec2 import EC2Client
27+
from mypy_boto3_ecr import ECRClient
28+
from mypy_boto3_ecs import ECSClient
29+
from mypy_boto3_eks import EKSClient
30+
from mypy_boto3_elasticbeanstalk import ElasticBeanstalkClient
31+
from mypy_boto3_elbv2 import ElasticLoadBalancingv2Client
32+
from mypy_boto3_emr import EMRClient
33+
from mypy_boto3_es import ElasticsearchServiceClient
34+
from mypy_boto3_events import EventBridgeClient
35+
from mypy_boto3_firehose import FirehoseClient
36+
from mypy_boto3_glacier import GlacierClient
37+
from mypy_boto3_glue import GlueClient
38+
from mypy_boto3_iam import IAMClient
39+
from mypy_boto3_iot import IoTClient
40+
from mypy_boto3_iot_data import IoTDataPlaneClient
41+
from mypy_boto3_iotanalytics import IoTAnalyticsClient
42+
from mypy_boto3_iotwireless import IoTWirelessClient
43+
from mypy_boto3_kafka import KafkaClient
44+
from mypy_boto3_kinesis import KinesisClient
45+
from mypy_boto3_kms import KMSClient
46+
from mypy_boto3_lakeformation import LakeFormationClient
47+
from mypy_boto3_lambda import LambdaClient
48+
from mypy_boto3_logs import CloudWatchLogsClient
49+
from mypy_boto3_mediastore import MediaStoreClient
50+
from mypy_boto3_mq import MQClient
51+
from mypy_boto3_mwaa import MWAAClient
52+
from mypy_boto3_neptune import NeptuneClient
53+
from mypy_boto3_opensearch import OpenSearchServiceClient
54+
from mypy_boto3_organizations import OrganizationsClient
55+
from mypy_boto3_pi import PIClient
56+
from mypy_boto3_qldb import QLDBClient
57+
from mypy_boto3_qldb_session import QLDBSessionClient
58+
from mypy_boto3_rds import RDSClient
59+
from mypy_boto3_rds_data import RDSDataServiceClient
60+
from mypy_boto3_redshift import RedshiftClient
61+
from mypy_boto3_redshift_data import RedshiftDataAPIServiceClient
62+
from mypy_boto3_resource_groups import ResourceGroupsClient
63+
from mypy_boto3_resourcegroupstaggingapi import ResourceGroupsTaggingAPIClient
64+
from mypy_boto3_route53 import Route53Client
65+
from mypy_boto3_route53resolver import Route53ResolverClient
66+
from mypy_boto3_s3 import S3Client
67+
from mypy_boto3_s3control import S3ControlClient
68+
from mypy_boto3_sagemaker import SageMakerClient
69+
from mypy_boto3_sagemaker_runtime import SageMakerRuntimeClient
70+
from mypy_boto3_secretsmanager import SecretsManagerClient
71+
from mypy_boto3_serverlessrepo import ServerlessApplicationRepositoryClient
72+
from mypy_boto3_servicediscovery import ServiceDiscoveryClient
73+
from mypy_boto3_ses import SESClient
74+
from mypy_boto3_sesv2 import SESV2Client
75+
from mypy_boto3_sns import SNSClient
76+
from mypy_boto3_sqs import SQSClient
77+
from mypy_boto3_ssm import SSMClient
78+
from mypy_boto3_stepfunctions import SFNClient
79+
from mypy_boto3_sts import STSClient
80+
from mypy_boto3_timestream_query import TimestreamQueryClient
81+
from mypy_boto3_timestream_write import TimestreamWriteClient
82+
from mypy_boto3_transcribe import TranscribeServiceClient
83+
from mypy_boto3_xray import XRayClient
84+
85+
from localstack.aws.connect import MetadataRequestInjector
86+
87+
88+
class TypedServiceClientFactory(abc.ABC):
89+
acm: Union["ACMClient", "MetadataRequestInjector[ACMClient]"]
90+
amplify: Union["AmplifyClient", "MetadataRequestInjector[AmplifyClient]"]
91+
apigateway: Union["APIGatewayClient", "MetadataRequestInjector[APIGatewayClient]"]
92+
apigatewayv2: Union["ApiGatewayV2Client", "MetadataRequestInjector[ApiGatewayV2Client]"]
93+
appconfig: Union["AppConfigClient", "MetadataRequestInjector[AppConfigClient]"]
94+
appsync: Union["AppSyncClient", "MetadataRequestInjector[AppSyncClient]"]
95+
athena: Union["AthenaClient", "MetadataRequestInjector[AthenaClient]"]
96+
autoscaling: Union["AutoScalingClient", "MetadataRequestInjector[AutoScalingClient]"]
97+
awslambda: Union["LambdaClient", "MetadataRequestInjector[LambdaClient]"]
98+
backup: Union["BackupClient", "MetadataRequestInjector[BackupClient]"]
99+
batch: Union["BatchClient", "MetadataRequestInjector[BatchClient]"]
100+
ce: Union["CostExplorerClient", "MetadataRequestInjector[CostExplorerClient]"]
101+
cloudformation: Union["CloudFormationClient", "MetadataRequestInjector[CloudFormationClient]"]
102+
cloudfront: Union["CloudFrontClient", "MetadataRequestInjector[CloudFrontClient]"]
103+
cloudtrail: Union["CloudTrailClient", "MetadataRequestInjector[CloudTrailClient]"]
104+
cloudwatch: Union["CloudWatchClient", "MetadataRequestInjector[CloudWatchClient]"]
105+
codecommit: Union["CodeCommitClient", "MetadataRequestInjector[CodeCommitClient]"]
106+
cognito_identity: Union[
107+
"CognitoIdentityClient", "MetadataRequestInjector[CognitoIdentityClient]"
108+
]
109+
cognito_idp: Union[
110+
"CognitoIdentityProviderClient", "MetadataRequestInjector[CognitoIdentityProviderClient]"
111+
]
112+
docdb: Union["DocDBClient", "MetadataRequestInjector[DocDBClient]"]
113+
dynamodb: Union["DynamoDBClient", "MetadataRequestInjector[DynamoDBClient]"]
114+
dynamodbstreams: Union[
115+
"DynamoDBStreamsClient", "MetadataRequestInjector[DynamoDBStreamsClient]"
116+
]
117+
ec2: Union["EC2Client", "MetadataRequestInjector[EC2Client]"]
118+
ecr: Union["ECRClient", "MetadataRequestInjector[ECRClient]"]
119+
ecs: Union["ECSClient", "MetadataRequestInjector[ECSClient]"]
120+
eks: Union["EKSClient", "MetadataRequestInjector[EKSClient]"]
121+
elasticbeanstalk: Union[
122+
"ElasticBeanstalkClient", "MetadataRequestInjector[ElasticBeanstalkClient]"
123+
]
124+
elbv2: Union[
125+
"ElasticLoadBalancingv2Client", "MetadataRequestInjector[ElasticLoadBalancingv2Client]"
126+
]
127+
emr: Union["EMRClient", "MetadataRequestInjector[EMRClient]"]
128+
es: Union["ElasticsearchServiceClient", "MetadataRequestInjector[ElasticsearchServiceClient]"]
129+
events: Union["EventBridgeClient", "MetadataRequestInjector[EventBridgeClient]"]
130+
firehose: Union["FirehoseClient", "MetadataRequestInjector[FirehoseClient]"]
131+
glacier: Union["GlacierClient", "MetadataRequestInjector[GlacierClient]"]
132+ glue: Union["GlueClient", "MetadataRequestInjector[GlueClient]"]
133+
iam: Union["IAMClient", "MetadataRequestInjector[IAMClient]"]
134+
iot: Union["IoTClient", "MetadataRequestInjector[IoTClient]"]
135+
iot_data: Union["IoTDataPlaneClient", "MetadataRequestInjector[IoTDataPlaneClient]"]
136+
iotanalytics: Union["IoTAnalyticsClient", "MetadataRequestInjector[IoTAnalyticsClient]"]
137+
iotwireless: Union["IoTWirelessClient", "MetadataRequestInjector[IoTWirelessClient]"]
138+
kafka: Union["KafkaClient", "MetadataRequestInjector[KafkaClient]"]
139+
kinesis: Union["KinesisClient", "MetadataRequestInjector[KinesisClient]"]
140+
kms: Union["KMSClient", "MetadataRequestInjector[KMSClient]"]
141+
lakeformation: Union["LakeFormationClient", "MetadataRequestInjector[LakeFormationClient]"]
142+
logs: Union["CloudWatchLogsClient", "MetadataRequestInjector[CloudWatchLogsClient]"]
143+
mediastore: Union["MediaStoreClient", "MetadataRequestInjector[MediaStoreClient]"]
144+
mq: Union["MQClient", "MetadataRequestInjector[MQClient]"]
145+
mwaa: Union["MWAAClient", "MetadataRequestInjector[MWAAClient]"]
146+
neptune: Union["NeptuneClient", "MetadataRequestInjector[NeptuneClient]"]
147+
opensearch: Union["OpenSearchServiceClient", "MetadataRequestInjector[OpenSearchServiceClient]"]
148+
organizations: Union["OrganizationsClient", "MetadataRequestInjector[OrganizationsClient]"]
149+
pi: Union["PIClient", "MetadataRequestInjector[PIClient]"]
150+
qldb: Union["QLDBClient", "MetadataRequestInjector[QLDBClient]"]
151+
qldb_session: Union["QLDBSessionClient", "MetadataRequestInjector[QLDBSessionClient]"]
152+
rds: Union["RDSClient", "MetadataRequestInjector[RDSClient]"]
153+
rds_data: Union["RDSDataServiceClient", "MetadataRequestInjector[RDSDataServiceClient]"]
154+
redshift: Union["RedshiftClient", "MetadataRequestInjector[RedshiftClient]"]
155+
redshift_data: Union[
156+
"RedshiftDataAPIServiceClient", "MetadataRequestInjector[RedshiftDataAPIServiceClient]"
157+
]
158+
resource_groups: Union["ResourceGroupsClient", "MetadataRequestInjector[ResourceGroupsClient]"]
159+
resourcegroupstaggingapi: Union[
160+
"ResourceGroupsTaggingAPIClient", "MetadataRequestInjector[ResourceGroupsTaggingAPIClient]"
161+
]
162+
route53: Union["Route53Client", "MetadataRequestInjector[Route53Client]"]
163+
route53resolver: Union[
164+
"Route53ResolverClient", "MetadataRequestInjector[Route53ResolverClient]"
165+
]
166+
s3: Union["S3Client", "MetadataRequestInjector[S3Client]"]
167+
s3control: Union["S3ControlClient", "MetadataRequestInjector[S3ControlClient]"]
168+
sagemaker: Union["SageMakerClient", "MetadataRequestInjector[SageMakerClient]"]
169+
sagemaker_runtime: Union[
170+
"SageMakerRuntimeClient", "MetadataRequestInjector[SageMakerRuntimeClient]"
171+
]
172+
secretsmanager: Union["SecretsManagerClient", "MetadataRequestInjector[SecretsManagerClient]"]
173+
serverlessrepo: Union[
174+
"ServerlessApplicationRepositoryClient",
175+
"MetadataRequestInjector[ServerlessApplicationRepositoryClient]",
176+
]
177+
servicediscovery: Union[
178+
"ServiceDiscoveryClient", "MetadataRequestInjector[ServiceDiscoveryClient]"
179+
]
180+
ses: Union["SESClient", "MetadataRequestInjector[SESClient]"]
181+
sesv2: Union["SESV2Client", "MetadataRequestInjector[SESV2Client]"]
182+
sns: Union["SNSClient", "MetadataRequestInjector[SNSClient]"]
183+
sqs: Union["SQSClient", "MetadataRequestInjector[SQSClient]"]
184+
ssm: Union["SSMClient", "MetadataRequestInjector[SSMClient]"]
185+
stepfunctions: Union["SFNClient", "MetadataRequestInjector[SFNClient]"]
186+
sts: Union["STSClient", "MetadataRequestInjector[STSClient]"]
187+
timestream_query: Union[
188+
"TimestreamQueryClient", "MetadataRequestInjector[TimestreamQueryClient]"
189+
]
190+
timestream_write: Union[
191+
"TimestreamWriteClient", "MetadataRequestInjector[TimestreamWriteClient]"
192+
]
193+
transcribe: Union["TranscribeServiceClient", "MetadataRequestInjector[TranscribeServiceClient]"]
194+
xray: Union["XRayClient", "MetadataRequestInjector[XRayClient]"]

0 commit comments

Comments
 (0)
0