From 2a6ba0ddd1c90b4cc95e4f25873f904f7c010072 Mon Sep 17 00:00:00 2001 From: John-Paul Stanford Date: Thu, 7 Apr 2022 10:45:17 +0100 Subject: [PATCH] New methods to the databrew service for recipes and rulesets. (#4996) --- moto/databrew/exceptions.py | 10 + moto/databrew/models.py | 98 ++++++++++ moto/databrew/responses.py | 92 +++++++++- moto/databrew/urls.py | 4 +- ...t_databrew.py => test_databrew_recipes.py} | 40 ++++ tests/test_databrew/test_databrew_rulesets.py | 171 ++++++++++++++++++ 6 files changed, 411 insertions(+), 4 deletions(-) rename tests/test_databrew/{test_databrew.py => test_databrew_recipes.py} (70%) create mode 100644 tests/test_databrew/test_databrew_rulesets.py diff --git a/moto/databrew/exceptions.py b/moto/databrew/exceptions.py index bf2b726e4..bcc208502 100644 --- a/moto/databrew/exceptions.py +++ b/moto/databrew/exceptions.py @@ -15,6 +15,11 @@ class RecipeAlreadyExistsException(AlreadyExistsException): super().__init__("Recipe") +class RulesetAlreadyExistsException(AlreadyExistsException): + def __init__(self): + super().__init__("Ruleset") + + class EntityNotFoundException(DataBrewClientError): def __init__(self, msg): super().__init__("EntityNotFoundException", msg) @@ -23,3 +28,8 @@ class EntityNotFoundException(DataBrewClientError): class RecipeNotFoundException(EntityNotFoundException): def __init__(self, recipe_name): super().__init__("Recipe %s not found." % recipe_name) + + +class RulesetNotFoundException(EntityNotFoundException): + def __init__(self, recipe_name): + super().__init__("Ruleset %s not found." % recipe_name) diff --git a/moto/databrew/models.py b/moto/databrew/models.py index 6f085669b..cc27bcc97 100644 --- a/moto/databrew/models.py +++ b/moto/databrew/models.py @@ -5,6 +5,7 @@ from moto.core import BaseBackend, BaseModel from moto.core.utils import BackendDict from moto.utilities.paginator import paginate from .exceptions import RecipeAlreadyExistsException, RecipeNotFoundException +from .exceptions import RulesetAlreadyExistsException, RulesetNotFoundException class DataBrewBackend(BaseBackend): @@ -15,11 +16,18 @@ class DataBrewBackend(BaseBackend): "limit_default": 100, "unique_attribute": "name", }, + "list_rulesets": { + "input_token": "next_token", + "limit_key": "max_results", + "limit_default": 100, + "unique_attribute": "name", + }, } def __init__(self, region_name): self.region_name = region_name self.recipes = OrderedDict() + self.rulesets = OrderedDict() def reset(self): """Re-initialize all attributes for this instance.""" @@ -37,6 +45,20 @@ class DataBrewBackend(BaseBackend): self.recipes[recipe_name] = recipe return recipe + def update_recipe(self, recipe_name, recipe_description, recipe_steps, tags): + if recipe_name not in self.recipes: + raise RecipeNotFoundException(recipe_name) + + recipe = self.recipes[recipe_name] + if recipe_description is not None: + recipe.description = recipe_description + if recipe_steps is not None: + recipe.steps = recipe_steps + if tags is not None: + recipe.tags = tags + + return recipe + @paginate(pagination_model=PAGINATION_MODEL) def list_recipes(self): return [self.recipes[key] for key in self.recipes] if self.recipes else [] @@ -49,6 +71,52 @@ class DataBrewBackend(BaseBackend): raise RecipeNotFoundException(recipe_name) return self.recipes[recipe_name] + def create_ruleset( + self, ruleset_name, ruleset_description, ruleset_rules, ruleset_target_arn, tags + ): + if ruleset_name in self.rulesets: + raise RulesetAlreadyExistsException() + + ruleset = FakeRuleset( + self.region_name, + ruleset_name, + ruleset_description, + ruleset_rules, + ruleset_target_arn, + tags, + ) + self.rulesets[ruleset_name] = ruleset + return ruleset + + def update_ruleset(self, ruleset_name, ruleset_description, ruleset_rules, tags): + if ruleset_name not in self.rulesets: + raise RulesetNotFoundException(ruleset_name) + + ruleset = self.rulesets[ruleset_name] + if ruleset_description is not None: + ruleset.description = ruleset_description + if ruleset_rules is not None: + ruleset.rules = ruleset_rules + if tags is not None: + ruleset.tags = tags + + return ruleset + + def get_ruleset(self, ruleset_name): + if ruleset_name not in self.rulesets: + raise RulesetNotFoundException(ruleset_name) + return self.rulesets[ruleset_name] + + @paginate(pagination_model=PAGINATION_MODEL) + def list_rulesets(self): + return list(self.rulesets.values()) + + def delete_ruleset(self, ruleset_name): + if ruleset_name not in self.rulesets: + raise RulesetNotFoundException(ruleset_name) + + del self.rulesets[ruleset_name] + class FakeRecipe(BaseModel): def __init__( @@ -71,4 +139,34 @@ class FakeRecipe(BaseModel): } +class FakeRuleset(BaseModel): + def __init__( + self, + region_name, + ruleset_name, + ruleset_description, + ruleset_rules, + ruleset_target_arn, + tags, + ): + self.region_name = region_name + self.name = ruleset_name + self.description = ruleset_description + self.rules = ruleset_rules + self.target_arn = ruleset_target_arn + self.created_time = datetime.now() + + self.tags = tags + + def as_dict(self): + return { + "Name": self.name, + "Rules": self.rules, + "Description": self.description, + "TargetArn": self.target_arn, + "CreateTime": self.created_time.isoformat(), + "Tags": self.tags or dict(), + } + + databrew_backends = BackendDict(DataBrewBackend, "databrew") diff --git a/moto/databrew/responses.py b/moto/databrew/responses.py index 7955a8601..216d63ea2 100644 --- a/moto/databrew/responses.py +++ b/moto/databrew/responses.py @@ -50,12 +50,98 @@ class DataBrewResponse(BaseResponse): } ) + def put_recipe_response(self, recipe_name): + recipe_description = self.parameters.get("Description") + recipe_steps = self.parameters.get("Steps") + tags = self.parameters.get("Tags") + + recipe = self.databrew_backend.update_recipe( + recipe_name, recipe_description, recipe_steps, tags + ) + return 200, {}, json.dumps(recipe.as_dict()) + + def get_recipe_response(self, recipe_name): + recipe = self.databrew_backend.get_recipe(recipe_name) + return 201, {}, json.dumps(recipe.as_dict()) + @amzn_request_id - def describe_recipe_response(self, request, full_url, headers): + def recipe_response(self, request, full_url, headers): self.setup_class(request, full_url, headers) parsed_url = urlparse(full_url) recipe_name = parsed_url.path.rstrip("/").rsplit("/", 1)[1] - recipe = self.databrew_backend.get_recipe(recipe_name) - return json.dumps(recipe.as_dict()) + if request.method == "PUT": + return self.put_recipe_response(recipe_name) + elif request.method == "GET": + return self.get_recipe_response(recipe_name) + + @amzn_request_id + def create_ruleset(self): + ruleset_description = self.parameters.get("Description") + ruleset_rules = self.parameters.get("Rules") + ruleset_name = self.parameters.get("Name") + ruleset_target_arn = self.parameters.get("TargetArn") + tags = self.parameters.get("Tags") + + return json.dumps( + self.databrew_backend.create_ruleset( + ruleset_name, + ruleset_description, + ruleset_rules, + ruleset_target_arn, + tags, + ).as_dict() + ) + + def put_ruleset_response(self, ruleset_name): + ruleset_description = self.parameters.get("Description") + ruleset_rules = self.parameters.get("Rules") + tags = self.parameters.get("Tags") + + ruleset = self.databrew_backend.update_ruleset( + ruleset_name, ruleset_description, ruleset_rules, tags + ) + return 200, {}, json.dumps(ruleset.as_dict()) + + def get_ruleset_response(self, ruleset_name): + ruleset = self.databrew_backend.get_ruleset(ruleset_name) + return 201, {}, json.dumps(ruleset.as_dict()) + + def delete_ruleset_response(self, ruleset_name): + self.databrew_backend.delete_ruleset(ruleset_name) + return 204, {}, "" + + @amzn_request_id + def ruleset_response(self, request, full_url, headers): + self.setup_class(request, full_url, headers) + parsed_url = urlparse(full_url) + + ruleset_name = parsed_url.path.split("/")[-1] + + if request.method == "PUT": + response = self.put_ruleset_response(ruleset_name) + return response + elif request.method == "GET": + return self.get_ruleset_response(ruleset_name) + elif request.method == "DELETE": + return self.delete_ruleset_response(ruleset_name) + + @amzn_request_id + def list_rulesets(self): + # https://docs.aws.amazon.com/databrew/latest/dg/API_ListRulesets.html + next_token = self._get_param("NextToken", self._get_param("nextToken")) + max_results = self._get_int_param( + "MaxResults", self._get_int_param("maxResults") + ) + + # pylint: disable=unexpected-keyword-arg, unbalanced-tuple-unpacking + ruleset_list, next_token = self.databrew_backend.list_rulesets( + next_token=next_token, max_results=max_results + ) + return json.dumps( + { + "Rulesets": [ruleset.as_dict() for ruleset in ruleset_list], + "NextToken": next_token, + } + ) diff --git a/moto/databrew/urls.py b/moto/databrew/urls.py index 0394fa8ce..b3130c603 100644 --- a/moto/databrew/urls.py +++ b/moto/databrew/urls.py @@ -4,5 +4,7 @@ url_bases = [r"https?://databrew\.(.+)\.amazonaws.com"] url_paths = { "{0}/recipes$": DataBrewResponse.dispatch, - "{0}/recipes/(?P[^/]+)$": DataBrewResponse().describe_recipe_response, + "{0}/recipes/(?P[^/]+)$": DataBrewResponse().recipe_response, + "{0}/rulesets$": DataBrewResponse.dispatch, + "{0}/rulesets/(?P[^/]+)$": DataBrewResponse().ruleset_response, } diff --git a/tests/test_databrew/test_databrew.py b/tests/test_databrew/test_databrew_recipes.py similarity index 70% rename from tests/test_databrew/test_databrew.py rename to tests/test_databrew/test_databrew_recipes.py index b32c26b38..f0590cfc1 100644 --- a/tests/test_databrew/test_databrew.py +++ b/tests/test_databrew/test_databrew_recipes.py @@ -96,6 +96,46 @@ def test_describe_recipe(): recipe["Steps"].should.have.length_of(1) +@mock_databrew +def test_update_recipe(): + client = _create_databrew_client() + response = _create_test_recipe(client) + + recipe = client.update_recipe( + Name=response["Name"], + Steps=[ + { + "Action": { + "Operation": "REMOVE_COMBINED", + "Parameters": { + "collapseConsecutiveWhitespace": "false", + "removeAllPunctuation": "false", + "removeAllQuotes": "false", + "removeAllWhitespace": "false", + "removeCustomCharacters": "true", + "removeCustomValue": "true", + "removeLeadingAndTrailingPunctuation": "false", + "removeLeadingAndTrailingQuotes": "false", + "removeLeadingAndTrailingWhitespace": "false", + "removeLetters": "false", + "removeNumbers": "false", + "removeSpecialCharacters": "true", + "sourceColumn": "FakeColumn", + }, + } + } + ], + ) + + recipe["Name"].should.equal(response["Name"]) + + # Describe the recipe and change the changes + recipe = client.describe_recipe(Name=response["Name"]) + recipe["Name"].should.equal(response["Name"]) + recipe["Steps"].should.have.length_of(1) + recipe["Steps"][0]["Action"]["Parameters"]["removeCustomValue"].should.equal("true") + + @mock_databrew def test_describe_recipe_that_does_not_exist(): client = _create_databrew_client() diff --git a/tests/test_databrew/test_databrew_rulesets.py b/tests/test_databrew/test_databrew_rulesets.py new file mode 100644 index 000000000..c95de802b --- /dev/null +++ b/tests/test_databrew/test_databrew_rulesets.py @@ -0,0 +1,171 @@ +import uuid + +import boto3 +import pytest +from botocore.exceptions import ClientError + +from moto import mock_databrew + + +def _create_databrew_client(): + client = boto3.client("databrew", region_name="us-west-1") + return client + + +def _create_test_ruleset(client, tags=None, ruleset_name=None): + if ruleset_name is None: + ruleset_name = str(uuid.uuid4()) + + return client.create_ruleset( + Name=ruleset_name, + TargetArn="arn:aws:databrew:eu-west-1:000000000000:dataset/fake-dataset", + Rules=[ + { + "Name": "Assert values > 0", + "Disabled": False, + "CheckExpression": ":col1 > :val1", + "SubstitutionMap": {":col1": "`Value`", ":val1": "0"}, + "Threshold": { + "Value": 100, + "Type": "GREATER_THAN_OR_EQUAL", + "Unit": "PERCENTAGE", + }, + } + ], + Tags=tags or {}, + ) + + +def _create_test_rulesets(client, count): + for _ in range(count): + _create_test_ruleset(client) + + +@mock_databrew +def test_ruleset_list_when_empty(): + client = _create_databrew_client() + + response = client.list_rulesets() + response.should.have.key("Rulesets") + response["Rulesets"].should.have.length_of(0) + + +@mock_databrew +def test_list_ruleset_with_max_results(): + client = _create_databrew_client() + + _create_test_rulesets(client, 4) + response = client.list_rulesets(MaxResults=2) + response["Rulesets"].should.have.length_of(2) + response.should.have.key("NextToken") + + +@mock_databrew +def test_list_rulesets_from_next_token(): + client = _create_databrew_client() + _create_test_rulesets(client, 10) + first_response = client.list_rulesets(MaxResults=3) + response = client.list_rulesets(NextToken=first_response["NextToken"]) + response["Rulesets"].should.have.length_of(7) + + +@mock_databrew +def test_list_rulesets_with_max_results_greater_than_actual_results(): + client = _create_databrew_client() + _create_test_rulesets(client, 4) + response = client.list_rulesets(MaxResults=10) + response["Rulesets"].should.have.length_of(4) + + +@mock_databrew +def test_describe_ruleset(): + client = _create_databrew_client() + response = _create_test_ruleset(client) + + ruleset = client.describe_ruleset(Name=response["Name"]) + + ruleset["Name"].should.equal(response["Name"]) + ruleset["Rules"].should.have.length_of(1) + + +@mock_databrew +def test_describe_ruleset_that_does_not_exist(): + client = _create_databrew_client() + + with pytest.raises(ClientError) as exc: + client.describe_ruleset(Name="DoseNotExist") + err = exc.value.response["Error"] + err["Code"].should.equal("EntityNotFoundException") + err["Message"].should.equal("Ruleset DoseNotExist not found.") + + +@mock_databrew +def test_create_ruleset_that_already_exists(): + client = _create_databrew_client() + + response = _create_test_ruleset(client) + + with pytest.raises(ClientError) as exc: + _create_test_ruleset(client, ruleset_name=response["Name"]) + err = exc.value.response["Error"] + err["Code"].should.equal("AlreadyExistsException") + err["Message"].should.equal("Ruleset already exists.") + + +@mock_databrew +def test_delete_ruleset(): + client = _create_databrew_client() + response = _create_test_ruleset(client) + + # Check ruleset exists + ruleset = client.describe_ruleset(Name=response["Name"]) + ruleset["Name"].should.equal(response["Name"]) + + # Delete the ruleset + client.delete_ruleset(Name=response["Name"]) + + # Check it does not exist anymore + with pytest.raises(ClientError) as exc: + client.describe_ruleset(Name=response["Name"]) + + err = exc.value.response["Error"] + err["Code"].should.equal("EntityNotFoundException") + err["Message"].should.equal(f"Ruleset {response['Name']} not found.") + + # Check that a ruleset that does not exist errors + with pytest.raises(ClientError) as exc: + client.delete_ruleset(Name=response["Name"]) + err = exc.value.response["Error"] + err["Code"].should.equal("EntityNotFoundException") + err["Message"].should.equal(f"Ruleset {response['Name']} not found.") + + +@mock_databrew +def test_update_ruleset(): + client = _create_databrew_client() + response = _create_test_ruleset(client) + + # Update the ruleset and check response + ruleset = client.update_ruleset( + Name=response["Name"], + Rules=[ + { + "Name": "Assert values > 0", + "Disabled": False, + "CheckExpression": ":col1 > :val1", + "SubstitutionMap": {":col1": "`Value`", ":val1": "10"}, + "Threshold": { + "Value": 100, + "Type": "GREATER_THAN_OR_EQUAL", + "Unit": "PERCENTAGE", + }, + } + ], + ) + ruleset["Name"].should.equal(response["Name"]) + + # Describe the ruleset and check the changes + ruleset = client.describe_ruleset(Name=response["Name"]) + ruleset["Name"].should.equal(response["Name"]) + ruleset["Rules"].should.have.length_of(1) + ruleset["Rules"][0]["SubstitutionMap"][":val1"].should.equal("10")