Yasuke/moto
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

New methods to the databrew service for recipes and rulesets. (#4996)

Browse Source
This commit is contained in:
John-Paul Stanford 2022-04-07 10:45:17 +01:00 committed by GitHub
parent 552385881c
commit 2a6ba0ddd1
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
6 changed files with 411 additions and 4 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

10
moto/databrew/exceptions.py
View File

@ -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)

98
moto/databrew/models.py
View File

@ -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")

92
moto/databrew/responses.py
View File

@ -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,
}
)

4
moto/databrew/urls.py
View File

@ -4,5 +4,7 @@ url_bases = [r"https?://databrew\.(.+)\.amazonaws.com"]
url_paths = {
"{0}/recipes$": DataBrewResponse.dispatch,
"{0}/recipes/(?P<recipe_name>[^/]+)$": DataBrewResponse().describe_recipe_response,
"{0}/recipes/(?P<recipe_name>[^/]+)$": DataBrewResponse().recipe_response,
"{0}/rulesets$": DataBrewResponse.dispatch,
"{0}/rulesets/(?P<ruleset_name>[^/]+)$": DataBrewResponse().ruleset_response,
}

40
tests/test_databrew/test_databrew.py → tests/test_databrew/test_databrew_recipes.py
View File

@ -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()

171
tests/test_databrew/test_databrew_rulesets.py Normal file
View File

@ -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")