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

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)

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

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)
+        if request.method == "PUT":
-        return json.dumps(recipe.as_dict())
+            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,
+            }
+        )

moto/databrew/urls.py

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

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

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