yujunjun/datahub
2
0
Fork 0
mirror of https://github.com/datahub-project/datahub.git synced 2025-09-26 09:35:23 +00:00
Code Issues Packages Projects Releases Wiki Activity

feat(ingest): support advanced configs for aws (#9237)

Browse Source
This commit is contained in:
Harshal Sheth 2023-11-16 13:34:18 -05:00 committed by GitHub
parent 9d41a8f9f0
commit 417ffb12d8
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
1 changed files with 23 additions and 6 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

29
metadata-ingestion/src/datahub/ingestion/source/aws/aws_common.py
View File

@ -1,8 +1,8 @@
from typing import TYPE_CHECKING, Dict, List, Optional, Union from typing import TYPE_CHECKING, Any, Dict, List, Optional, Union
import boto3 import boto3
from boto3.session import Session from boto3.session import Session
from botocore.config import Config from botocore.config import DEFAULT_TIMEOUT, Config
from botocore.utils import fix_s3_host from botocore.utils import fix_s3_host
from pydantic.fields import Field from pydantic.fields import Field
@ -104,6 +104,16 @@ class AwsConnectionConfig(ConfigModel):
description="A set of proxy configs to use with AWS. See the [botocore.config](https://botocore.amazonaws.com/v1/documentation/api/latest/reference/config.html) docs for details.", description="A set of proxy configs to use with AWS. See the [botocore.config](https://botocore.amazonaws.com/v1/documentation/api/latest/reference/config.html) docs for details.",
) )
read_timeout: float = Field(
default=DEFAULT_TIMEOUT,
description="The timeout for reading from the connection (in seconds).",
)
aws_advanced_config: Dict[str, Any] = Field(
default_factory=dict,
description="Advanced AWS configuration options. These are passed directly to [botocore.config.Config](https://botocore.amazonaws.com/v1/documentation/api/latest/reference/config.html).",
)
def _normalized_aws_roles(self) -> List[AwsAssumeRoleConfig]: def _normalized_aws_roles(self) -> List[AwsAssumeRoleConfig]:
if not self.aws_role: if not self.aws_role:
return [] return []
@ -167,13 +177,20 @@ class AwsConnectionConfig(ConfigModel):
} }
return {} return {}
def _aws_config(self) -> Config:
return Config(
proxies=self.aws_proxy,
read_timeout=self.read_timeout,
**self.aws_advanced_config,
)
def get_s3_client( def get_s3_client(
self, verify_ssl: Optional[Union[bool, str]] = None self, verify_ssl: Optional[Union[bool, str]] = None
) -> "S3Client": ) -> "S3Client":
return self.get_session().client( return self.get_session().client(
"s3", "s3",
endpoint_url=self.aws_endpoint_url, endpoint_url=self.aws_endpoint_url,
config=Config(proxies=self.aws_proxy), config=self._aws_config(),
verify=verify_ssl, verify=verify_ssl,
) )
@ -183,7 +200,7 @@ class AwsConnectionConfig(ConfigModel):
resource = self.get_session().resource( resource = self.get_session().resource(
"s3", "s3",
endpoint_url=self.aws_endpoint_url, endpoint_url=self.aws_endpoint_url,
config=Config(proxies=self.aws_proxy), config=self._aws_config(),
verify=verify_ssl, verify=verify_ssl,
) )
# according to: https://stackoverflow.com/questions/32618216/override-s3-endpoint-using-boto3-configuration-file # according to: https://stackoverflow.com/questions/32618216/override-s3-endpoint-using-boto3-configuration-file
@ -195,10 +212,10 @@ class AwsConnectionConfig(ConfigModel):
return resource return resource
def get_glue_client(self) -> "GlueClient": def get_glue_client(self) -> "GlueClient":
return self.get_session().client("glue") return self.get_session().client("glue", config=self._aws_config())
def get_sagemaker_client(self) -> "SageMakerClient": def get_sagemaker_client(self) -> "SageMakerClient":
return self.get_session().client("sagemaker") return self.get_session().client("sagemaker", config=self._aws_config())
class AwsSourceConfig(EnvConfigMixin, AwsConnectionConfig): class AwsSourceConfig(EnvConfigMixin, AwsConnectionConfig):