SageMaker: Add exception when resource already exists (#7239)

This commit is contained in:
Bogdan Girman 2024-01-25 23:34:04 +01:00 committed by GitHub
parent 494eac54aa
commit 33024619eb
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
3 changed files with 54 additions and 3 deletions

View File

@ -41,6 +41,10 @@ class AWSValidationException(AWSError):
TYPE = "ValidationException" TYPE = "ValidationException"
class ResourceInUseException(AWSError):
TYPE = "ResourceInUse"
class ResourceNotFound(JsonRESTError): class ResourceNotFound(JsonRESTError):
def __init__(self, message: str): def __init__(self, message: str):
super().__init__(__class__.__name__, message) # type: ignore super().__init__(__class__.__name__, message) # type: ignore

View File

@ -14,6 +14,7 @@ from moto.utilities.paginator import paginate
from .exceptions import ( from .exceptions import (
AWSValidationException, AWSValidationException,
MissingModel, MissingModel,
ResourceInUseException,
ResourceNotFound, ResourceNotFound,
ValidationError, ValidationError,
) )
@ -1006,6 +1007,9 @@ class FeatureGroup(BaseObject):
"Catalog": "AwsDataCatalog", "Catalog": "AwsDataCatalog",
"Database": "sagemaker_featurestore", "Database": "sagemaker_featurestore",
} }
offline_store_config["S3StorageConfig"][
"ResolvedOutputS3Uri"
] = f'{offline_store_config["S3StorageConfig"]["S3Uri"]}/{account_id}/{region_name}/offline-store/{feature_group_name}-{int(datetime.now().timestamp())}/data'
self.offline_store_config = offline_store_config self.offline_store_config = offline_store_config
self.role_arn = role_arn self.role_arn = role_arn
@ -3529,6 +3533,17 @@ class SageMakerModelBackend(BaseBackend):
role_arn: str, role_arn: str,
tags: Any, tags: Any,
) -> str: ) -> str:
feature_group_arn = arn_formatter(
region_name=self.region_name,
account_id=self.account_id,
_type="feature-group",
_id=f"{feature_group_name.lower()}",
)
if feature_group_arn in self.feature_groups:
raise ResourceInUseException(
message=f"An error occurred (ResourceInUse) when calling the CreateFeatureGroup operation: Resource Already Exists: FeatureGroup with name {feature_group_name} already exists. Choose a different name.\nInfo: Feature Group '{feature_group_name}' already exists."
)
feature_group = FeatureGroup( feature_group = FeatureGroup(
feature_group_name=feature_group_name, feature_group_name=feature_group_name,
record_identifier_feature_name=record_identifier_feature_name, record_identifier_feature_name=record_identifier_feature_name,

View File

@ -3,6 +3,8 @@ import re
from datetime import datetime from datetime import datetime
import boto3 import boto3
import pytest
from botocore.exceptions import ClientError
from moto import mock_sagemaker from moto import mock_sagemaker
@ -34,6 +36,29 @@ def test_create_feature_group():
== "arn:aws:sagemaker:us-east-2:123456789012:feature-group/some-feature-group-name" == "arn:aws:sagemaker:us-east-2:123456789012:feature-group/some-feature-group-name"
) )
with pytest.raises(ClientError) as raised_exception:
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 raised_exception.value.response["Error"]["Code"] == "ResourceInUse"
assert (
raised_exception.value.response["Error"]["Message"]
== "An error occurred (ResourceInUse) when calling the CreateFeatureGroup operation: Resource Already Exists: FeatureGroup with name some-feature-group-name already exists. Choose a different name.\nInfo: Feature Group 'some-feature-group-name' already exists."
)
@mock_sagemaker @mock_sagemaker
def test_describe_feature_group(): def test_describe_feature_group():
@ -55,7 +80,7 @@ def test_describe_feature_group():
RoleArn=role_arn, RoleArn=role_arn,
OfflineStoreConfig={ OfflineStoreConfig={
"DisableGlueTableCreation": False, "DisableGlueTableCreation": False,
"S3StorageConfig": {"S3Uri": "s3://mybucket"}, "S3StorageConfig": {"S3Uri": "s3://mybucket/some-folder/some-subfolder"},
}, },
) )
resp = client.describe_feature_group(FeatureGroupName=feature_group_name) resp = client.describe_feature_group(FeatureGroupName=feature_group_name)
@ -70,7 +95,7 @@ def test_describe_feature_group():
assert resp["FeatureDefinitions"] == feature_definitions assert resp["FeatureDefinitions"] == feature_definitions
assert resp["RoleArn"] == role_arn assert resp["RoleArn"] == role_arn
assert re.match( assert re.match(
r"^some_feature_group_name_[0-9]+$", f"^{feature_group_name.replace('-', '_')}_[0-9]+$",
resp["OfflineStoreConfig"]["DataCatalogConfig"]["TableName"], resp["OfflineStoreConfig"]["DataCatalogConfig"]["TableName"],
) )
assert ( assert (
@ -80,6 +105,13 @@ def test_describe_feature_group():
resp["OfflineStoreConfig"]["DataCatalogConfig"]["Database"] resp["OfflineStoreConfig"]["DataCatalogConfig"]["Database"]
== "sagemaker_featurestore" == "sagemaker_featurestore"
) )
assert resp["OfflineStoreConfig"]["S3StorageConfig"]["S3Uri"] == "s3://mybucket" assert (
resp["OfflineStoreConfig"]["S3StorageConfig"]["S3Uri"]
== "s3://mybucket/some-folder/some-subfolder"
)
assert re.match(
f"^s3://mybucket/some-folder/some-subfolder/123456789012/us-east-2/offline-store/{feature_group_name}-[0-9]+/data$",
resp["OfflineStoreConfig"]["S3StorageConfig"]["ResolvedOutputS3Uri"],
)
assert isinstance(resp["CreationTime"], datetime) assert isinstance(resp["CreationTime"], datetime)
assert resp["FeatureGroupStatus"] == "Created" assert resp["FeatureGroupStatus"] == "Created"