Init sagemaker

metadata-ingestion/README.md

@@ -45,6 +45,7 @@ We use a plugin architecture so that you can install only the dependencies you a
 | oracle        | `pip install 'acryl-datahub[oracle]'`                      | Oracle source                       |
 | postgres      | `pip install 'acryl-datahub[postgres]'`                    | Postgres source                     |
 | redshift      | `pip install 'acryl-datahub[redshift]'`                    | Redshift source                     |
+| sagemaker     | `pip install 'acryl-datahub[sagemaker]'`                   | AWS SageMaker source                |
 | sqlalchemy    | `pip install 'acryl-datahub[sqlalchemy]'`                  | Generic SQLAlchemy source           |
 | snowflake     | `pip install 'acryl-datahub[snowflake]'`                   | Snowflake source                    |
 | superset      | `pip install 'acryl-datahub[superset]'`                    | Superset source                     |
@@ -344,6 +345,27 @@ source:
     # options is same as above
 ```
 
+### AWS SageMaker `sagemaker`
+
+Extracts:
+
+- Feature groups (support for models and jobs coming soon!)
+
+```yml
+source:
+  type: glue
+  config:
+    aws_region: # aws_region_name, i.e. "eu-west-1"
+    env: # environment for the DatasetSnapshot URN, one of "DEV", "EI", "PROD" or "CORP". Defaults to "PROD".
+
+    # Credentials. If not specified here, these are picked up according to boto3 rules.
+    # (see https://boto3.amazonaws.com/v1/documentation/api/latest/guide/credentials.html)
+    aws_access_key_id: # Optional.
+    aws_secret_access_key: # Optional.
+    aws_session_token: # Optional.
+    aws_role: # Optional (Role chaining supported by using a sorted list).
+```
+
 ### Snowflake `snowflake`
 
 Extracts:

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

@@ -77,11 +77,3 @@ def get_client(self, service: str) -> boto3.client:
             )
         else:
             return boto3.client(service, region_name=self.aws_region)
-
-    @property
-    def glue_client(self):
-        return self.get_client("glue")
-
-    @property
-    def s3_client(self):
-        return self.get_client("s3")

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

@@ -50,6 +50,14 @@ class GlueSourceConfig(AwsSourceConfig):
 
     extract_transforms: Optional[bool] = True
 
+    @property
+    def glue_client(self):
+        return self.get_client("glue")
+
+    @property
+    def s3_client(self):
+        return self.get_client("s3")
+
 
 @dataclass
 class GlueSourceReport(SourceReport):

metadata-ingestion/src/datahub/ingestion/source/sagemaker.py

@@ -0,0 +1,56 @@
+from dataclasses import dataclass
+from dataclasses import field as dataclass_field
+from typing import List
+
+from datahub.ingestion.api.common import PipelineContext
+from datahub.ingestion.api.source import Source, SourceReport
+from datahub.ingestion.source.aws_common import AwsSourceConfig
+
+
+class SagemakerSourceConfig(AwsSourceConfig):
+    @property
+    def sagemaker_client(self):
+        return self.get_client("samgemaker")
+
+
+@dataclass
+class SagemakerSourceReport(SourceReport):
+    tables_scanned = 0
+    filtered: List[str] = dataclass_field(default_factory=list)
+
+    def report_table_scanned(self) -> None:
+        self.tables_scanned += 1
+
+    def report_table_dropped(self, table: str) -> None:
+        self.filtered.append(table)
+
+
+class GlueSource(Source):
+    source_config: SagemakerSourceConfig
+    report = SagemakerSourceReport()
+
+    def __init__(self, config: SagemakerSourceConfig, ctx: PipelineContext):
+        super().__init__(ctx)
+        self.source_config = config
+        self.report = SagemakerSourceReport()
+        self.sagemaker_client = config.sagemaker_client
+        self.env = config.env
+
+    @classmethod
+    def create(cls, config_dict, ctx):
+        config = SagemakerSourceConfig.parse_obj(config_dict)
+        return cls(config, ctx)
+
+    def get_all_feature_groups(self):
+        """
+        List all feature groups in SageMaker.
+        """
+
+        feature_groups = []
+
+        # see https://boto3.amazonaws.com/v1/documentation/api/latest/reference/services/sagemaker.html#SageMaker.Client.list_feature_groups
+        paginator = self.sagemaker_client.get_paginator("list_feature_groups")
+        for page in paginator.paginate():
+            feature_groups += page["FeatureGroupSummaries"]
+
+        return feature_groups