(enhancement): Extend get and update ruleset DQ methods (#1882)

* (enhancement): Extend get and update ruleset DQ methods

Author: jaidisido

awswrangler/data_quality/_create.py

@@ -10,6 +10,7 @@
 
 from awswrangler import _utils, exceptions
 from awswrangler._config import apply_configs
+from awswrangler.data_quality._get import get_ruleset
 from awswrangler.data_quality._utils import (
     _create_datasource,
     _get_data_quality_results,
@@ -27,7 +28,7 @@ def _create_dqdl(
     """Create DQDL from pandas data frame."""
     rules = []
     for rule_type, parameter, expression in df_rules.itertuples(index=False):
-        parameter_str = f' "{parameter}" ' if parameter else " "
+        parameter_str = f" {parameter} " if parameter else " "
         expression_str = expression if expression else ""
         rules.append(f"{rule_type}{parameter_str}{expression_str}")
     return "Rules = [ " + ", ".join(rules) + " ]"
@@ -85,7 +86,7 @@ def create_ruleset(
     >>> df = pd.DataFrame({"c0": [0, 1, 2], "c1": [0, 1, 2], "c2": [0, 0, 1]})
     >>> df_rules = pd.DataFrame({
     >>>        "rule_type": ["RowCount", "IsComplete", "Uniqueness"],
-    >>>        "parameter": [None, "c0", "c0"],
+    >>>        "parameter": [None, '"c0"', '"c0"'],
     >>>        "expression": ["between 1 and 6", None, "> 0.95"],
     >>> })
     >>> wr.s3.to_parquet(df, path, dataset=True, database="database", table="table")
@@ -121,6 +122,7 @@ def create_ruleset(
 def update_ruleset(
     name: str,
     updated_name: Optional[str] = None,
+    mode: str = "overwrite",
     df_rules: Optional[pd.DataFrame] = None,
     dqdl_rules: Optional[str] = None,
     description: str = "",
@@ -134,6 +136,8 @@ def update_ruleset(
         Ruleset name.
     updated_name : str
         New ruleset name if renaming an existing ruleset.
+    mode : str
+        overwrite (default) or upsert.
     df_rules : str, optional
         Data frame with `rule_type`, `parameter`, and `expression` columns.
     dqdl_rules : str, optional
@@ -145,25 +149,46 @@ def update_ruleset(
 
     Examples
     --------
+    Overwrite rules in the existing ruleset.
     >>> wr.data_quality.update_ruleset(
     >>>     name="ruleset",
     >>>     new_name="my_ruleset",
     >>>     dqdl_rules="Rules = [ RowCount between 1 and 3 ]",
     >>>)
+
+    Update or insert rules in the existing ruleset.
+    >>> wr.data_quality.update_ruleset(
+    >>>     name="ruleset",
+    >>>     mode="insert",
+    >>>     dqdl_rules="Rules = [ RowCount between 1 and 3 ]",
+    >>>)
     """
     if (df_rules is not None and dqdl_rules) or (df_rules is None and not dqdl_rules):
         raise exceptions.InvalidArgumentCombination("You must pass either ruleset `df_rules` or `dqdl_rules`.")
+    if mode not in ["overwrite", "upsert"]:
+        raise exceptions.InvalidArgumentValue("`mode` must be one of 'overwrite' or 'upsert'.")
+
+    if mode == "upsert":
+        df_existing = get_ruleset(name=name, boto3_session=boto3_session)
+        df_existing = df_existing.set_index(keys=["rule_type", "parameter"], drop=False, verify_integrity=True)
+        df_updated = _rules_to_df(dqdl_rules) if dqdl_rules is not None else df_rules
+        df_updated = df_updated.set_index(keys=["rule_type", "parameter"], drop=False, verify_integrity=True)
+        merged_df = pd.concat([df_existing[~df_existing.index.isin(df_updated.index)], df_updated])
+        dqdl_rules = _create_dqdl(merged_df.reset_index(drop=True))
+    else:
+        dqdl_rules = _create_dqdl(df_rules) if df_rules is not None else dqdl_rules
+
+    args = {
+        "Name": name,
+        "Description": description,
+        "Ruleset": dqdl_rules,
+    }
+    if updated_name:
+        args["UpdatedName"] = updated_name
 
     client_glue: boto3.client = _utils.client(service_name="glue", session=boto3_session)
-    dqdl_rules = _create_dqdl(df_rules) if df_rules is not None else dqdl_rules
-
     try:
-        client_glue.update_data_quality_ruleset(
-            Name=name,
-            UpdatedName=updated_name,
-            Description=description,
-            Ruleset=dqdl_rules,
-        )
+        client_glue.update_data_quality_ruleset(**args)
     except client_glue.exceptions.EntityNotFoundException as not_found:
         raise exceptions.ResourceDoesNotExist(f"Ruleset {name} does not exist.") from not_found
 
@@ -327,7 +352,7 @@ def evaluate_ruleset(
     >>>     dqdl_rules="Rules = [ RowCount between 1 and 3 ]",
     >>>)
     >>> df_ruleset_results = wr.data_quality.evaluate_ruleset(
-    >>>     name=["ruleset1", "rulseset2"],
+    >>>     name="ruleset",
     >>>     iam_role_arn=glue_data_quality_role,
     >>> )
     """

awswrangler/data_quality/_get.py

@@ -1,6 +1,6 @@
 """AWS Glue Data Quality Get Module."""
 
-from typing import Optional, cast
+from typing import List, Optional, Union, cast
 
 import boto3
 import pandas as pd
@@ -9,28 +9,39 @@
 
 
 def get_ruleset(
-    name: str,
+    name: Union[str, List[str]],
     boto3_session: Optional[boto3.Session] = None,
 ) -> pd.DataFrame:
     """Get a Data Quality ruleset.
 
     Parameters
     ----------
-    name : str
-        Ruleset name.
+    name : str or list[str]
+        Ruleset name or list of names.
     boto3_session : boto3.Session, optional
         Boto3 Session. If none, the default boto3 session is used.
 
     Returns
     -------
     pd.DataFrame
-        Data frame with ruleset details.
+        Data frame with ruleset(s) details.
 
     Examples
     --------
+    Get single ruleset
     >>> import awswrangler as wr
 
     >>> df_ruleset = wr.data_quality.get_ruleset(name="my_ruleset")
+
+    Get multiple rulesets. A column with the ruleset name is added to the data frame
+    >>> df_rulesets = wr.data_quality.get_ruleset(name=["ruleset_1", "ruleset_2"])
     """
-    rules = cast(str, _get_ruleset(ruleset_name=name, boto3_session=boto3_session)["Ruleset"])
-    return _rules_to_df(rules=rules)
+    ruleset_names: List[str] = name if isinstance(name, list) else [name]
+    dfs: List[pd.DataFrame] = []
+    for ruleset_name in ruleset_names:
+        rules = cast(str, _get_ruleset(ruleset_name=ruleset_name, boto3_session=boto3_session)["Ruleset"])
+        df = _rules_to_df(rules=rules)
+        if len(ruleset_names) > 1:
+            df["ruleset"] = ruleset_name
+        dfs.append(df)
+    return pd.concat(dfs)

awswrangler/data_quality/_utils.py

@@ -8,6 +8,7 @@
 from typing import Any, Dict, List, Optional, Tuple, Union, cast
 
 import boto3
+import botocore.exceptions
 import pandas as pd
 
 from awswrangler import _utils, exceptions
@@ -23,12 +24,13 @@ def _parse_rules(rules: List[str]) -> List[Tuple[str, Optional[str], Optional[st
     for rule in rules:
         rule_type, remainder = tuple(rule.split(maxsplit=1))
         if remainder.startswith('"'):
-            remainder_split = remainder.split(maxsplit=1)
-            parameter = remainder_split[0].strip('"')
-            expression = None if len(remainder_split) == 1 else remainder_split[1]
+            expression_regex = r"\s+(?:[=><]|between\s+.+\s+and\s+|in\s+\[.+\]|matches\s+).*"
+            expression_matches = re.findall(expression_regex, remainder)
+            expression = None if len(expression_matches) == 0 else expression_matches[0].strip()
+            parameter = remainder.split(expression)[0].strip() if expression else remainder
         else:
-            parameter = None
             expression = remainder
+            parameter = None
         parsed_rules.append((rule_type, parameter, expression))
     return parsed_rules
 
@@ -115,10 +117,18 @@ def _get_ruleset_run(
 ) -> Dict[str, Any]:
     session: boto3.Session = _utils.ensure_session(session=boto3_session)
     client_glue: boto3.client = _utils.client(service_name="glue", session=session)
-    if run_type == "recommendation":
-        response = client_glue.get_data_quality_rule_recommendation_run(RunId=run_id)
-    elif run_type == "evaluation":
-        response = client_glue.get_data_quality_ruleset_evaluation_run(RunId=run_id)
+    f = (
+        client_glue.get_data_quality_rule_recommendation_run
+        if run_type == "recommendation"
+        else client_glue.get_data_quality_ruleset_evaluation_run
+    )
+    response = _utils.try_it(
+        f=f,
+        ex=botocore.exceptions.ClientError,
+        ex_code="ThrottlingException",
+        max_num_tries=5,
+        RunId=run_id,
+    )
     return cast(Dict[str, Any], response)
 
 
@@ -148,7 +158,14 @@ def _get_ruleset(
 ) -> Dict[str, Any]:
     boto3_session = _utils.ensure_session(session=boto3_session)
     client_glue: boto3.client = _utils.client(service_name="glue", session=boto3_session)
-    return cast(Dict[str, Any], client_glue.get_data_quality_ruleset(Name=ruleset_name))
+    response = _utils.try_it(
+        f=client_glue.get_data_quality_ruleset,
+        ex=botocore.exceptions.ClientError,
+        ex_code="ThrottlingException",
+        max_num_tries=5,
+        Name=ruleset_name,
+    )
+    return cast(Dict[str, Any], response)
 
 
 def _get_data_quality_results(