Repository: getmoto/moto
Pull request #7227 (2024-01-18 14:46:27)

Add SageMaker Feature Group

codecov[bot]: ## [Codecov](https://app.codecov.io/gh/getmoto/moto/pull/7227?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=getmoto) Report
All modified and coverable lines are covered by tests :white_check_mark:
> Comparison is base [(`e3f60e2`)](https://app.codecov.io/gh/getmoto/moto/commit/e3f60e2b1f35336dccb393a8cdd2abf5741f7dad?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=getmoto) 95.90% compared to head [(`786c4cd`)](https://app.codecov.io/gh/getmoto/moto/pull/7227?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=getmoto) 95.90%.
> Report is 7 commits behind head on master.


<details><summary>Additional details and impacted files</summary>


```diff
@@           Coverage Diff           @@
##           master    #7227   +/-   ##
=======================================
  Coverage   95.90%   95.90%           
=======================================
  Files         840      840           
  Lines       82653    82750   +97     
=======================================
+ Hits        79268    79363   +95     
- Misses       3385     3387    +2     
```

| [Flag](https://app.codecov.io/gh/getmoto/moto/pull/7227/flags?src=pr&el=flags&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=getmoto) | Coverage Δ | |
|---|---|---|
| [servertests](https://app.codecov.io/gh/getmoto/moto/pull/7227/flags?src=pr&el=flag&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=getmoto) | `35.47% <20.00%> (-0.02%)` | :arrow_down: |
| [unittests](https://app.codecov.io/gh/getmoto/moto/pull/7227/flags?src=pr&el=flag&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=getmoto) | `95.85% <100.00%> (+<0.01%)` | :arrow_up: |

Flags with carried forward coverage won't be shown. [Click here](https://docs.codecov.io/docs/carryforward-flags?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=getmoto#carryforward-flags-in-the-pull-request-comment) to find out more.


</details>

[:umbrella: View full report in Codecov by Sentry](https://app.codecov.io/gh/getmoto/moto/pull/7227?src=pr&el=continue&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=getmoto).   
:loudspeaker: Have feedback on the report? [Share it here](https://about.codecov.io/codecov-pr-comment-feedback/?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=getmoto).


Patch:
diff --git a/moto/sagemaker/models.py b/moto/sagemaker/models.py
index 06f5896cc..c284c65ec 100644
--- a/moto/sagemaker/models.py
+++ b/moto/sagemaker/models.py
@@ -980,6 +980,60 @@ class ModelPackageGroup(BaseObject):
         return {k: v for k, v in response_object.items() if k in response_values}
 
 
+class FeatureGroup(BaseObject):
+    def __init__(
+        self,
+        region_name: str,
+        account_id: str,
+        feature_group_name: str,
+        record_identifier_feature_name: str,
+        event_time_feature_name: str,
+        feature_definitions: List[Dict[str, str]],
+        offline_store_config: Dict[str, Any],
+        role_arn: str,
+        tags: Optional[List[Dict[str, str]]] = None,
+    ) -> None:
+        self.feature_group_name = feature_group_name
+        self.record_identifier_feature_name = record_identifier_feature_name
+        self.event_time_feature_name = event_time_feature_name
+        self.feature_definitions = feature_definitions
+
+        table_name = (
+            f"{feature_group_name.replace('-','_')}_{int(datetime.now().timestamp())}"
+        )
+        offline_store_config["DataCatalogConfig"] = {
+            "TableName": table_name,
+            "Catalog": "AwsDataCatalog",
+            "Database": "sagemaker_featurestore",
+        }
+
+        self.offline_store_config = offline_store_config
+        self.role_arn = role_arn
+
+        self.creation_time = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
+        self.feature_group_arn = arn_formatter(
+            region_name=region_name,
+            account_id=account_id,
+            _type="feature-group",
+            _id=f"{self.feature_group_name.lower()}",
+        )
+        self.tags = tags
+
+    def describe(self) -> Dict[str, Any]:
+        return {
+            "FeatureGroupArn": self.feature_group_arn,
+            "FeatureGroupName": self.feature_group_name,
+            "RecordIdentifierFeatureName": self.record_identifier_feature_name,
+            "EventTimeFeatureName": self.event_time_feature_name,
+            "FeatureDefinitions": self.feature_definitions,
+            "CreationTime": self.creation_time,
+            "OfflineStoreConfig": self.offline_store_config,
+            "RoleArn": self.role_arn,
+            "ThroughputConfig": {"ThroughputMode": "OnDemand"},
+            "FeatureGroupStatus": "Created",
+        }
+
+
 class ModelPackage(BaseObject):
     def __init__(
         self,
@@ -1768,6 +1822,7 @@ class SageMakerModelBackend(BaseBackend):
         self.model_package_groups: Dict[str, ModelPackageGroup] = {}
         self.model_packages: Dict[str, ModelPackage] = {}
         self.model_package_name_mapping: Dict[str, str] = {}
+        self.feature_groups: Dict[str, FeatureGroup] = {}
 
     @staticmethod
     def default_vpc_endpoint_service(
@@ -3464,6 +3519,44 @@ class SageMakerModelBackend(BaseBackend):
         self.model_packages[model_package.model_package_arn] = model_package
         return model_package.model_package_arn
 
+    def create_feature_group(
+        self,
+        feature_group_name: str,
+        record_identifier_feature_name: str,
+        event_time_feature_name: str,
+        feature_definitions: List[Dict[str, str]],
+        offline_store_config: Dict[str, Any],
+        role_arn: str,
+        tags: Any,
+    ) -> str:
+        feature_group = FeatureGroup(
+            feature_group_name=feature_group_name,
+            record_identifier_feature_name=record_identifier_feature_name,
+            event_time_feature_name=event_time_feature_name,
+            feature_definitions=feature_definitions,
+            offline_store_config=offline_store_config,
+            role_arn=role_arn,
+            region_name=self.region_name,
+            account_id=self.account_id,
+            tags=tags,
+        )
+        self.feature_groups[feature_group.feature_group_arn] = feature_group
+        return feature_group.feature_group_arn
+
+    def describe_feature_group(
+        self,
+        feature_group_name: str,
+    ) -> Dict[str, Any]:
+        feature_group_arn = arn_formatter(
+            region_name=self.region_name,
+            account_id=self.account_id,
+            _type="feature-group",
+            _id=f"{feature_group_name.lower()}",
+        )
+
+        feature_group = self.feature_groups[feature_group_arn]
+        return feature_group.describe()
+
 
 class FakeExperiment(BaseObject):
     def __init__(
diff --git a/moto/sagemaker/responses.py b/moto/sagemaker/responses.py
index f260130d1..fa63920b0 100644
--- a/moto/sagemaker/responses.py
+++ b/moto/sagemaker/responses.py
@@ -959,3 +959,23 @@ class SageMakerResponse(BaseResponse):
             tags=tags,
         )
         return json.dumps(dict(ModelPackageGroupArn=model_package_group_arn))
+
+    def create_feature_group(self) -> str:
+        feature_group_arn = self.sagemaker_backend.create_feature_group(
+            feature_group_name=self._get_param("FeatureGroupName"),
+            record_identifier_feature_name=self._get_param(
+                "RecordIdentifierFeatureName"
+            ),
+            event_time_feature_name=self._get_param("EventTimeFeatureName"),
+            feature_definitions=self._get_param("FeatureDefinitions"),
+            offline_store_config=self._get_param("OfflineStoreConfig"),
+            role_arn=self._get_param("RoleArn"),
+            tags=self._get_param("Tags"),
+        )
+        return json.dumps(dict(FeatureGroupArn=feature_group_arn))
+
+    def describe_feature_group(self) -> str:
+        resp = self.sagemaker_backend.describe_feature_group(
+            feature_group_name=self._get_param("FeatureGroupName"),
+        )
+        return json.dumps(resp)


Test patch:
diff --git a/tests/test_sagemaker/test_sagemaker_feature_groups.py b/tests/test_sagemaker/test_sagemaker_feature_groups.py
new file mode 100644
index 000000000..c3f4a1475
--- /dev/null
+++ b/tests/test_sagemaker/test_sagemaker_feature_groups.py
@@ -0,0 +1,85 @@
+"""Unit tests for sagemaker-supported APIs."""
+import re
+from datetime import datetime
+
+import boto3
+
+from moto import mock_sagemaker
+
+# See our Development Tips on writing tests for hints on how to write good tests:
+# http://docs.getmoto.org/en/latest/docs/contributing/development_tips/tests.html
+
+
+@mock_sagemaker
+def test_create_feature_group():
+    client = boto3.client("sagemaker", region_name="us-east-2")
+    resp = client.create_feature_group(
+        FeatureGroupName="some-feature-group-name",
+        RecordIdentifierFeatureName="some_record_identifier",
+        EventTimeFeatureName="EventTime",
+        FeatureDefinitions=[
+            {"FeatureName": "some_feature", "FeatureType": "String"},
+            {"FeatureName": "EventTime", "FeatureType": "Fractional"},
+            {"FeatureName": "some_record_identifier", "FeatureType": "String"},
+        ],
+        RoleArn="arn:aws:iam::123456789012:role/AWSFeatureStoreAccess",
+        OfflineStoreConfig={
+            "DisableGlueTableCreation": False,
+            "S3StorageConfig": {"S3Uri": "s3://mybucket"},
+        },
+    )
+
+    assert (
+        resp["FeatureGroupArn"]
+        == "arn:aws:sagemaker:us-east-2:123456789012:feature-group/some-feature-group-name"
+    )
+
+
+@mock_sagemaker
+def test_describe_feature_group():
+    client = boto3.client("sagemaker", region_name="us-east-2")
+    feature_group_name = "some-feature-group-name"
+    record_identifier_feature_name = "some_record_identifier"
+    event_time_feature_name = "EventTime"
+    role_arn = "arn:aws:iam::123456789012:role/AWSFeatureStoreAccess"
+    feature_definitions = [
+        {"FeatureName": "some_feature", "FeatureType": "String"},
+        {"FeatureName": event_time_feature_name, "FeatureType": "Fractional"},
+        {"FeatureName": record_identifier_feature_name, "FeatureType": "String"},
+    ]
+    client.create_feature_group(
+        FeatureGroupName=feature_group_name,
+        RecordIdentifierFeatureName=record_identifier_feature_name,
+        EventTimeFeatureName=event_time_feature_name,
+        FeatureDefinitions=feature_definitions,
+        RoleArn=role_arn,
+        OfflineStoreConfig={
+            "DisableGlueTableCreation": False,
+            "S3StorageConfig": {"S3Uri": "s3://mybucket"},
+        },
+    )
+    resp = client.describe_feature_group(FeatureGroupName=feature_group_name)
+
+    assert resp["FeatureGroupName"] == feature_group_name
+    assert (
+        resp["FeatureGroupArn"]
+        == "arn:aws:sagemaker:us-east-2:123456789012:feature-group/some-feature-group-name"
+    )
+    assert resp["RecordIdentifierFeatureName"] == record_identifier_feature_name
+    assert resp["EventTimeFeatureName"] == event_time_feature_name
+    assert resp["FeatureDefinitions"] == feature_definitions
+    assert resp["RoleArn"] == role_arn
+    assert re.match(
+        r"^some_feature_group_name_[0-9]+$",
+        resp["OfflineStoreConfig"]["DataCatalogConfig"]["TableName"],
+    )
+    assert (
+        resp["OfflineStoreConfig"]["DataCatalogConfig"]["Catalog"] == "AwsDataCatalog"
+    )
+    assert (
+        resp["OfflineStoreConfig"]["DataCatalogConfig"]["Database"]
+        == "sagemaker_featurestore"
+    )
+    assert resp["OfflineStoreConfig"]["S3StorageConfig"]["S3Uri"] == "s3://mybucket"
+    assert isinstance(resp["CreationTime"], datetime)
+    assert resp["FeatureGroupStatus"] == "Created"