Add test stub with offline store

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

@@ -79,6 +79,7 @@ def get_feature_group_details(self, feature_group_name: str) -> Dict[str, Any]:
             FeatureGroupName=feature_group_name
         )
 
+        # use falsy fallback since AWS stubs require this to be a string in tests
         next_token = feature_group.get("NextToken", "")
 
         # paginate over feature group features

metadata-ingestion/tests/unit/sagemaker/sagemaker_mces_golden.json

@@ -1,4 +1,98 @@
 [
+  {
+    "auditHeader": null,
+    "proposedSnapshot": {
+      "com.linkedin.pegasus2avro.metadata.snapshot.MLFeatureSnapshot": {
+        "urn": "urn:li:mlFeature:(test-2,some-feature-1)",
+        "aspects": [
+          {
+            "com.linkedin.pegasus2avro.ml.metadata.MLFeatureProperties": {
+              "description": null,
+              "dataType": "TEXT",
+              "version": null,
+              "sources": [
+                "urn:li:dataset:(urn:li:dataPlatform:s3,datahub-sagemaker-outputs,PROD)",
+                "urn:li:dataset:(urn:li:dataPlatform:glue,sagemaker_featurestore.test-2-123412341234,PROD)"
+              ]
+            }
+          }
+        ]
+      }
+    },
+    "proposedDelta": null
+  },
+  {
+    "auditHeader": null,
+    "proposedSnapshot": {
+      "com.linkedin.pegasus2avro.metadata.snapshot.MLPrimaryKeySnapshot": {
+        "urn": "urn:li:mlPrimaryKey:(test-2,some-feature-2)",
+        "aspects": [
+          {
+            "com.linkedin.pegasus2avro.ml.metadata.MLPrimaryKeyProperties": {
+              "description": null,
+              "dataType": "ORDINAL",
+              "version": null,
+              "sources": [
+                "urn:li:dataset:(urn:li:dataPlatform:s3,datahub-sagemaker-outputs,PROD)",
+                "urn:li:dataset:(urn:li:dataPlatform:glue,sagemaker_featurestore.test-2-123412341234,PROD)"
+              ]
+            }
+          }
+        ]
+      }
+    },
+    "proposedDelta": null
+  },
+  {
+    "auditHeader": null,
+    "proposedSnapshot": {
+      "com.linkedin.pegasus2avro.metadata.snapshot.MLFeatureSnapshot": {
+        "urn": "urn:li:mlFeature:(test-2,some-feature-3)",
+        "aspects": [
+          {
+            "com.linkedin.pegasus2avro.ml.metadata.MLFeatureProperties": {
+              "description": null,
+              "dataType": "CONTINUOUS",
+              "version": null,
+              "sources": [
+                "urn:li:dataset:(urn:li:dataPlatform:s3,datahub-sagemaker-outputs,PROD)",
+                "urn:li:dataset:(urn:li:dataPlatform:glue,sagemaker_featurestore.test-2-123412341234,PROD)"
+              ]
+            }
+          }
+        ]
+      }
+    },
+    "proposedDelta": null
+  },
+  {
+    "auditHeader": null,
+    "proposedSnapshot": {
+      "com.linkedin.pegasus2avro.metadata.snapshot.MLFeatureTableSnapshot": {
+        "urn": "urn:li:mlFeatureTable:(urn:li:dataPlatform:sagemaker,test-2)",
+        "aspects": [
+          {
+            "com.linkedin.pegasus2avro.ml.metadata.MLFeatureTableProperties": {
+              "customProperties": {
+                "arn": "arn:aws:sagemaker:us-west-2:123412341234:feature-group/test-2",
+                "creation_time": "2021-06-24 09:48:37.035000",
+                "status": "Created"
+              },
+              "description": "Yet another test feature group",
+              "mlFeatures": [
+                "urn:li:mlFeature:(test-2,some-feature-1)",
+                "urn:li:mlFeature:(test-2,some-feature-3)"
+              ],
+              "mlPrimaryKeys": [
+                "urn:li:mlPrimaryKey:(test-2,some-feature-2)"
+              ]
+            }
+          }
+        ]
+      }
+    },
+    "proposedDelta": null
+  },
   {
     "auditHeader": null,
     "proposedSnapshot": {

metadata-ingestion/tests/unit/test_sagemaker_source.py

@@ -9,6 +9,7 @@
 from tests.unit.test_sagemaker_source_stubs import (
     describe_feature_group_response_1,
     describe_feature_group_response_2,
+    describe_feature_group_response_3,
     list_feature_groups_response,
 )
 
@@ -34,6 +35,13 @@ def test_sagemaker_ingest(tmp_path, pytestconfig):
             list_feature_groups_response,
             {},
         )
+        sagemaker_stubber.add_response(
+            "describe_feature_group",
+            describe_feature_group_response_1,
+            {
+                "FeatureGroupName": "test-2",
+            },
+        )
         sagemaker_stubber.add_response(
             "describe_feature_group",
             describe_feature_group_response_2,
@@ -43,7 +51,7 @@ def test_sagemaker_ingest(tmp_path, pytestconfig):
         )
         sagemaker_stubber.add_response(
             "describe_feature_group",
-            describe_feature_group_response_1,
+            describe_feature_group_response_3,
             {
                 "FeatureGroupName": "test",
             },

metadata-ingestion/tests/unit/test_sagemaker_source_stubs.py

@@ -2,6 +2,12 @@
 
 list_feature_groups_response = {
     "FeatureGroupSummaries": [
+        {
+            "FeatureGroupName": "test-2",
+            "FeatureGroupArn": "arn:aws:sagemaker:us-west-2:123412341234:feature-group/test-2",
+            "CreationTime": datetime.datetime(2021, 6, 24, 9, 48, 37, 35000),
+            "FeatureGroupStatus": "Created",
+        },
         {
             "FeatureGroupName": "test-1",
             "FeatureGroupArn": "arn:aws:sagemaker:us-west-2:123412341234:feature-group/test-1",
@@ -19,26 +25,32 @@
 }
 
 describe_feature_group_response_1 = {
-    "FeatureGroupArn": "arn:aws:sagemaker:us-west-2:123412341234:feature-group/test",
-    "FeatureGroupName": "test",
-    "RecordIdentifierFeatureName": "feature_1",
-    "EventTimeFeatureName": "feature_3",
+    "FeatureGroupArn": "arn:aws:sagemaker:us-west-2:123412341234:feature-group/test-2",
+    "FeatureGroupName": "test-2",
+    "RecordIdentifierFeatureName": "some-feature-2",
+    "EventTimeFeatureName": "some-feature-3",
     "FeatureDefinitions": [
-        {"FeatureName": "feature_1", "FeatureType": "String"},
-        {"FeatureName": "feature_2", "FeatureType": "Integral"},
-        {"FeatureName": "feature_3", "FeatureType": "Fractional"},
+        {"FeatureName": "some-feature-1", "FeatureType": "String"},
+        {"FeatureName": "some-feature-2", "FeatureType": "Integral"},
+        {"FeatureName": "some-feature-3", "FeatureType": "Fractional"},
     ],
-    "CreationTime": datetime.datetime(
-        2021,
-        6,
-        14,
-        11,
-        3,
-        0,
-        803000,
-    ),
+    "CreationTime": datetime.datetime(2021, 6, 24, 9, 48, 37, 35000),
     "OnlineStoreConfig": {"EnableOnlineStore": True},
+    "OfflineStoreConfig": {
+        "S3StorageConfig": {
+            "S3Uri": "s3://datahub-sagemaker-outputs",
+            "ResolvedOutputS3Uri": "s3://datahub-sagemaker-outputs/123412341234/sagemaker/us-west-2/offline-store/test-2-123412341234/data",
+        },
+        "DisableGlueTableCreation": False,
+        "DataCatalogConfig": {
+            "TableName": "test-2-123412341234",
+            "Catalog": "AwsDataCatalog",
+            "Database": "sagemaker_featurestore",
+        },
+    },
+    "RoleArn": "arn:aws:iam::123412341234:role/service-role/AmazonSageMaker-ExecutionRole-20210614T104201",
     "FeatureGroupStatus": "Created",
+    "Description": "Yet another test feature group",
     "NextToken": "",
 }
 
@@ -59,3 +71,27 @@
     "Description": "First test feature group",
     "NextToken": "",
 }
+
+describe_feature_group_response_3 = {
+    "FeatureGroupArn": "arn:aws:sagemaker:us-west-2:123412341234:feature-group/test",
+    "FeatureGroupName": "test",
+    "RecordIdentifierFeatureName": "feature_1",
+    "EventTimeFeatureName": "feature_3",
+    "FeatureDefinitions": [
+        {"FeatureName": "feature_1", "FeatureType": "String"},
+        {"FeatureName": "feature_2", "FeatureType": "Integral"},
+        {"FeatureName": "feature_3", "FeatureType": "Fractional"},
+    ],
+    "CreationTime": datetime.datetime(
+        2021,
+        6,
+        14,
+        11,
+        3,
+        0,
+        803000,
+    ),
+    "OnlineStoreConfig": {"EnableOnlineStore": True},
+    "FeatureGroupStatus": "Created",
+    "NextToken": "",
+}