feat(ingest): type stubs for boto3

metadata-ingestion/setup.py

@@ -122,6 +122,7 @@ def get_long_description():
     "types-cachetools",
     # versions 0.1.13 and 0.1.14 seem to have issues
     "types-click==0.1.12",
+    "boto3-stubs[s3,glue,sagemaker]",
 }
 
 base_dev_requirements = {

metadata-ingestion/src/datahub/ingestion/source/aws/aws_common.py

@@ -2,6 +2,10 @@
 from typing import List, Optional, Union
 
 import boto3
+from boto3.session import Session
+from mypy_boto3_glue import GlueClient
+from mypy_boto3_s3 import S3Client
+from mypy_boto3_sagemaker import SageMakerClient
 
 from datahub.configuration import ConfigModel
 from datahub.configuration.common import AllowDenyPattern
@@ -46,22 +50,20 @@ class AwsSourceConfig(ConfigModel):
     aws_role: Optional[Union[str, List[str]]] = None
     aws_region: str
 
-    def get_client(self, service: str) -> boto3.client:
+    def get_session(self) -> Session:
         if (
             self.aws_access_key_id
             and self.aws_secret_access_key
             and self.aws_session_token
         ):
-            return boto3.client(
-                service,
+            return Session(
                 aws_access_key_id=self.aws_access_key_id,
                 aws_secret_access_key=self.aws_secret_access_key,
                 aws_session_token=self.aws_session_token,
                 region_name=self.aws_region,
             )
         elif self.aws_access_key_id and self.aws_secret_access_key:
-            return boto3.client(
-                service,
+            return Session(
                 aws_access_key_id=self.aws_access_key_id,
                 aws_secret_access_key=self.aws_secret_access_key,
                 region_name=self.aws_region,
@@ -77,15 +79,23 @@ def get_client(self, service: str) -> boto3.client:
                     self.aws_role,
                     {},
                 )
-            return boto3.client(
-                service,
+            return Session(
                 aws_access_key_id=credentials["AccessKeyId"],
                 aws_secret_access_key=credentials["SecretAccessKey"],
                 aws_session_token=credentials["SessionToken"],
                 region_name=self.aws_region,
             )
         else:
-            return boto3.client(service, region_name=self.aws_region)
+            return Session(region_name=self.aws_region)
+
+    def get_s3_client(self) -> S3Client:
+        return self.get_session().client("s3")
+
+    def get_glue_client(self) -> GlueClient:
+        return self.get_session().client("glue")
+
+    def get_sagemaker_client(self) -> SageMakerClient:
+        return self.get_session().client("sagemaker")
 
 
 def make_s3_urn(s3_uri: str, env: str, suffix: Optional[str] = None) -> str:

metadata-ingestion/src/datahub/ingestion/source/aws/glue.py

@@ -50,11 +50,11 @@ class GlueSourceConfig(AwsSourceConfig):
 
     @property
     def glue_client(self):
-        return self.get_client("glue")
+        return self.get_glue_client()
 
     @property
     def s3_client(self):
-        return self.get_client("s3")
+        return self.get_s3_client()
 
 
 @dataclass

metadata-ingestion/src/datahub/ingestion/source/aws/sagemaker_processors/common.py

@@ -13,7 +13,7 @@ class SagemakerSourceConfig(AwsSourceConfig):
 
     @property
     def sagemaker_client(self):
-        return self.get_client("sagemaker")
+        return self.get_sagemaker_client()
 
 
 @dataclass

metadata-ingestion/src/datahub/ingestion/source/aws/sagemaker_processors/feature_groups.py

@@ -1,5 +1,12 @@
 from dataclasses import dataclass
-from typing import Any, Dict, Iterable, List
+from typing import Iterable, List
+
+from mypy_boto3_sagemaker import SageMakerClient
+from mypy_boto3_sagemaker.type_defs import (
+    DescribeFeatureGroupResponseTypeDef,
+    FeatureDefinitionTypeDef,
+    FeatureGroupSummaryTypeDef,
+)
 
 import datahub.emitter.mce_builder as builder
 from datahub.ingestion.api.workunit import MetadataWorkUnit
@@ -23,11 +30,11 @@
 
 @dataclass
 class FeatureGroupProcessor:
-    sagemaker_client: Any
+    sagemaker_client: SageMakerClient
     env: str
     report: SagemakerSourceReport
 
-    def get_all_feature_groups(self) -> List[Dict[str, Any]]:
+    def get_all_feature_groups(self) -> List[FeatureGroupSummaryTypeDef]:
         """
         List all feature groups in SageMaker.
         """
@@ -41,7 +48,9 @@ def get_all_feature_groups(self) -> List[Dict[str, Any]]:
 
         return feature_groups
 
-    def get_feature_group_details(self, feature_group_name: str) -> Dict[str, Any]:
+    def get_feature_group_details(
+        self, feature_group_name: str
+    ) -> DescribeFeatureGroupResponseTypeDef:
         """
         Get details of a feature group (including list of component features).
         """
@@ -59,15 +68,13 @@ def get_feature_group_details(self, feature_group_name: str) -> Dict[str, Any]:
             next_features = self.sagemaker_client.describe_feature_group(
                 FeatureGroupName=feature_group_name, NextToken=next_token
             )
-            feature_group["FeatureDefinitions"].append(
-                next_features["FeatureDefinitions"]
-            )
+            feature_group["FeatureDefinitions"] += next_features["FeatureDefinitions"]
             next_token = feature_group.get("NextToken", "")
 
         return feature_group
 
     def get_feature_group_wu(
-        self, feature_group_details: Dict[str, Any]
+        self, feature_group_details: DescribeFeatureGroupResponseTypeDef
     ) -> MetadataWorkUnit:
         """
         Generate an MLFeatureTable workunit for a SageMaker feature group.
@@ -138,7 +145,9 @@ def get_feature_type(self, aws_type: str, feature_name: str) -> str:
         return mapped_type
 
     def get_feature_wu(
-        self, feature_group_details: Dict[str, Any], feature: Dict[str, Any]
+        self,
+        feature_group_details: DescribeFeatureGroupResponseTypeDef,
+        feature: FeatureDefinitionTypeDef,
     ) -> MetadataWorkUnit:
         """
         Generate an MLFeature workunit for a SageMaker feature.

metadata-ingestion/src/datahub/ingestion/source/aws/sagemaker_processors/job_classes.py

@@ -0,0 +1,232 @@
+from typing import Dict
+
+from typing_extensions import Final
+
+from datahub.metadata.schema_classes import JobStatusClass
+
+
+class SageMakerJobInfo:
+
+    # Note: The weird typing for the following commands is because the child classes
+    # use a Final type to ensure that they're interpreted by the boto3 stubs correctly.
+    # If we try to type these as plain strings, we get a TypeError because it's being converted
+    # to a non-overwritable type.
+    # See https://mypy.readthedocs.io/en/stable/final_attrs.html#details-of-using-final
+
+    # boto3 command to get list of jobs
+    @property
+    def list_command(self) -> str:
+        raise NotImplementedError
+
+    # field in job listing response containing actual list
+    @property
+    def list_key(self) -> str:
+        raise NotImplementedError
+
+    # field in job listing response element corresponding to job name
+    @property
+    def list_name_key(self) -> str:
+        raise NotImplementedError
+
+    # field in job listing response element corresponding to job ARN
+    @property
+    def list_arn_key(self) -> str:
+        raise NotImplementedError
+
+    # boto3 command to get job details
+    @property
+    def describe_command(self) -> str:
+        raise NotImplementedError
+
+    # field in job description response corresponding to job name
+    @property
+    def describe_name_key(self) -> str:
+        raise NotImplementedError
+
+    # field in job description response corresponding to job ARN
+    @property
+    def describe_arn_key(self) -> str:
+        raise NotImplementedError
+
+    # field in job description response corresponding to job status
+    @property
+    def describe_status_key(self) -> str:
+        raise NotImplementedError
+
+    # job-specific mapping from boto3 status strings to DataHub-native enum
+    status_map: Dict[str, str]
+
+    # name of function for processing job for ingestion
+    processor: str
+
+
+class AutoMlJobInfo(SageMakerJobInfo):
+    # https://boto3.amazonaws.com/v1/documentation/api/latest/reference/services/sagemaker.html#SageMaker.Client.list_auto_ml_jobs
+    list_command: Final = "list_auto_ml_jobs"
+    list_key: Final = "AutoMLJobSummaries"
+    list_name_key: Final = "AutoMLJobName"
+    list_arn_key: Final = "AutoMLJobArn"
+    # https://boto3.amazonaws.com/v1/documentation/api/latest/reference/services/sagemaker.html#SageMaker.Client.describe_auto_ml_job
+    describe_command: Final = "describe_auto_ml_job"
+    describe_name_key: Final = "AutoMLJobName"
+    describe_arn_key: Final = "AutoMLJobArn"
+    describe_status_key: Final = "AutoMLJobStatus"
+    status_map = {
+        "Completed": JobStatusClass.COMPLETED,
+        "InProgress": JobStatusClass.IN_PROGRESS,
+        "Failed": JobStatusClass.FAILED,
+        "Stopped": JobStatusClass.STOPPED,
+        "Stopping": JobStatusClass.STOPPING,
+    }
+
+    processor = "process_auto_ml_job"
+
+
+class CompilationJobInfo(SageMakerJobInfo):
+    # https://boto3.amazonaws.com/v1/documentation/api/latest/reference/services/sagemaker.html#SageMaker.Client.list_compilation_jobs
+    list_command: Final = "list_compilation_jobs"
+    list_key: Final = "CompilationJobSummaries"
+    list_name_key: Final = "CompilationJobName"
+    list_arn_key: Final = "CompilationJobArn"
+    # https://boto3.amazonaws.com/v1/documentation/api/latest/reference/services/sagemaker.html#SageMaker.Client.describe_compilation_job
+    describe_command: Final = "describe_compilation_job"
+    describe_name_key: Final = "CompilationJobName"
+    describe_arn_key: Final = "CompilationJobArn"
+    describe_status_key: Final = "CompilationJobStatus"
+    status_map = {
+        "INPROGRESS": JobStatusClass.IN_PROGRESS,
+        "COMPLETED": JobStatusClass.COMPLETED,
+        "FAILED": JobStatusClass.FAILED,
+        "STARTING": JobStatusClass.STARTING,
+        "STOPPING": JobStatusClass.STOPPING,
+        "STOPPED": JobStatusClass.STOPPED,
+    }
+    processor = "process_compilation_job"
+
+
+class EdgePackagingJobInfo(SageMakerJobInfo):
+    # https://boto3.amazonaws.com/v1/documentation/api/latest/reference/services/sagemaker.html#SageMaker.Client.list_edge_packaging_jobs
+    list_command: Final = "list_edge_packaging_jobs"
+    list_key: Final = "EdgePackagingJobSummaries"
+    list_name_key: Final = "EdgePackagingJobName"
+    list_arn_key: Final = "EdgePackagingJobArn"
+    # https://boto3.amazonaws.com/v1/documentation/api/latest/reference/services/sagemaker.html#SageMaker.Client.describe_edge_packaging_job
+    describe_command: Final = "describe_edge_packaging_job"
+    describe_name_key: Final = "EdgePackagingJobName"
+    describe_arn_key: Final = "EdgePackagingJobArn"
+    describe_status_key: Final = "EdgePackagingJobStatus"
+    status_map = {
+        "INPROGRESS": JobStatusClass.IN_PROGRESS,
+        "COMPLETED": JobStatusClass.COMPLETED,
+        "FAILED": JobStatusClass.FAILED,
+        "STARTING": JobStatusClass.STARTING,
+        "STOPPING": JobStatusClass.STOPPING,
+        "STOPPED": JobStatusClass.STOPPED,
+    }
+    processor = "process_edge_packaging_job"
+
+
+class HyperParameterTuningJobInfo(SageMakerJobInfo):
+    # https://boto3.amazonaws.com/v1/documentation/api/latest/reference/services/sagemaker.html#SageMaker.Client.list_hyper_parameter_tuning_jobs
+    list_command: Final = "list_hyper_parameter_tuning_jobs"
+    list_key: Final = "HyperParameterTuningJobSummaries"
+    list_name_key: Final = "HyperParameterTuningJobName"
+    list_arn_key: Final = "HyperParameterTuningJobArn"
+    # https://boto3.amazonaws.com/v1/documentation/api/latest/reference/services/sagemaker.html#SageMaker.Client.describe_hyper_parameter_tuning_job
+    describe_command: Final = "describe_hyper_parameter_tuning_job"
+    describe_name_key: Final = "HyperParameterTuningJobName"
+    describe_arn_key: Final = "HyperParameterTuningJobArn"
+    describe_status_key: Final = "HyperParameterTuningJobStatus"
+    status_map = {
+        "InProgress": JobStatusClass.IN_PROGRESS,
+        "Completed": JobStatusClass.COMPLETED,
+        "Failed": JobStatusClass.FAILED,
+        "Stopping": JobStatusClass.STOPPING,
+        "Stopped": JobStatusClass.STOPPED,
+    }
+    processor = "process_hyper_parameter_tuning_job"
+
+
+class LabelingJobInfo(SageMakerJobInfo):
+    # https://boto3.amazonaws.com/v1/documentation/api/latest/reference/services/sagemaker.html#SageMaker.Client.list_labeling_jobs
+    list_command: Final = "list_labeling_jobs"
+    list_key: Final = "LabelingJobSummaryList"
+    list_name_key: Final = "LabelingJobName"
+    list_arn_key: Final = "LabelingJobArn"
+    # https://boto3.amazonaws.com/v1/documentation/api/latest/reference/services/sagemaker.html#SageMaker.Client.describe_labeling_job
+    describe_command: Final = "describe_labeling_job"
+    describe_name_key: Final = "LabelingJobName"
+    describe_arn_key: Final = "LabelingJobArn"
+    describe_status_key: Final = "LabelingJobStatus"
+    status_map = {
+        "Initializing": JobStatusClass.STARTING,
+        "InProgress": JobStatusClass.IN_PROGRESS,
+        "Completed": JobStatusClass.COMPLETED,
+        "Failed": JobStatusClass.FAILED,
+        "Stopping": JobStatusClass.STOPPING,
+        "Stopped": JobStatusClass.STOPPED,
+    }
+    processor = "process_labeling_job"
+
+
+class ProcessingJobInfo(SageMakerJobInfo):
+    # https://boto3.amazonaws.com/v1/documentation/api/latest/reference/services/sagemaker.html#SageMaker.Client.list_processing_jobs
+    list_command: Final = "list_processing_jobs"
+    list_key: Final = "ProcessingJobSummaries"
+    list_name_key: Final = "ProcessingJobName"
+    list_arn_key: Final = "ProcessingJobArn"
+    # https://boto3.amazonaws.com/v1/documentation/api/latest/reference/services/sagemaker.html#SageMaker.Client.describe_processing_job
+    describe_command: Final = "describe_processing_job"
+    describe_name_key: Final = "ProcessingJobName"
+    describe_arn_key: Final = "ProcessingJobArn"
+    describe_status_key: Final = "ProcessingJobStatus"
+    status_map = {
+        "InProgress": JobStatusClass.IN_PROGRESS,
+        "Completed": JobStatusClass.COMPLETED,
+        "Failed": JobStatusClass.FAILED,
+        "Stopping": JobStatusClass.STOPPING,
+        "Stopped": JobStatusClass.STOPPED,
+    }
+    processor = "process_processing_job"
+
+
+class TrainingJobInfo(SageMakerJobInfo):
+    # https://boto3.amazonaws.com/v1/documentation/api/latest/reference/services/sagemaker.html#SageMaker.Client.list_training_jobs
+    list_command: Final = "list_training_jobs"
+    list_key: Final = "TrainingJobSummaries"
+    list_name_key: Final = "TrainingJobName"
+    list_arn_key: Final = "TrainingJobArn"
+    # https://boto3.amazonaws.com/v1/documentation/api/latest/reference/services/sagemaker.html#SageMaker.Client.describe_training_job
+    describe_command: Final = "describe_training_job"
+    describe_name_key: Final = "TrainingJobName"
+    describe_arn_key: Final = "TrainingJobArn"
+    describe_status_key: Final = "TrainingJobStatus"
+    status_map = {
+        "InProgress": JobStatusClass.IN_PROGRESS,
+        "Completed": JobStatusClass.COMPLETED,
+        "Failed": JobStatusClass.FAILED,
+        "Stopping": JobStatusClass.STOPPING,
+        "Stopped": JobStatusClass.STOPPED,
+    }
+    processor = "process_training_job"
+
+
+class TransformJobInfo(SageMakerJobInfo):
+    # https://boto3.amazonaws.com/v1/documentation/api/latest/reference/services/sagemaker.html#SageMaker.Client.list_transform_jobs
+    list_command: Final = "list_transform_jobs"
+    list_key: Final = "TransformJobSummaries"
+    list_name_key: Final = "TransformJobName"
+    list_arn_key: Final = "TransformJobArn"
+    # https://boto3.amazonaws.com/v1/documentation/api/latest/reference/services/sagemaker.html#SageMaker.Client.describe_transform_job
+    describe_command: Final = "describe_transform_job"
+    describe_name_key: Final = "TransformJobName"
+    describe_arn_key: Final = "TransformJobArn"
+    describe_status_key: Final = "TransformJobStatus"
+    status_map = {
+        "InProgress": JobStatusClass.IN_PROGRESS,
+        "Completed": JobStatusClass.COMPLETED,
+        "Failed": JobStatusClass.FAILED,
+        "Stopping": JobStatusClass.STOPPING,
+        "Stopped": JobStatusClass.STOPPED,
+    }
+    processor = "process_transform_job"