add TagColumnOperation to quicksight.create_athena_dataset #1340

Author: cnfait

awswrangler/quicksight/_create.py

@@ -98,7 +98,9 @@ def _generate_permissions(
 
 
 def _generate_transformations(
-    rename_columns: Optional[Dict[str, str]], cast_columns_types: Optional[Dict[str, str]]
+    rename_columns: Optional[Dict[str, str]],
+    cast_columns_types: Optional[Dict[str, str]],
+    tag_columns: Optional[Dict[str, str]],
 ) -> List[Dict[str, Dict[str, Any]]]:
     trans: List[Dict[str, Dict[str, Any]]] = []
     if rename_columns is not None:
@@ -107,6 +109,9 @@ def _generate_transformations(
     if cast_columns_types is not None:
         for k, v in cast_columns_types.items():
             trans.append({"CastColumnTypeOperation": {"ColumnName": k, "NewColumnType": v.upper()}})
+    if tag_columns is not None:
+        for k, v in tag_columns.items():
+            trans.append({"TagColumnOperation": {"ColumnName": k, "Tags": v}})
     return trans
 
 
@@ -204,6 +209,7 @@ def create_athena_dataset(
     logical_table_alias: str = "LogicalTable",
     rename_columns: Optional[Dict[str, str]] = None,
     cast_columns_types: Optional[Dict[str, str]] = None,
+    tag_columns: Optional[Dict[str, str]] = None,
     tags: Optional[Dict[str, str]] = None,
     account_id: Optional[str] = None,
     boto3_session: Optional[boto3.Session] = None,
@@ -259,6 +265,10 @@ def create_athena_dataset(
     cast_columns_types : Dict[str, str], optional
         Dictionary to map column casts. e.g. {"col_name": "STRING", "col_name2": "DECIMAL"}
         Valid types: 'STRING'|'INTEGER'|'DECIMAL'|'DATETIME'
+    tag_columns : Dict[str, str], optional
+        Dictionary to map column tags. e.g.
+        {"col_name": { "ColumnGeographicRole": "CITY" }, "col_name2": { "ColumnDescription": { "Text": "description" }}}
+        Valid geospatial roles: 'COUNTRY'|'STATE'|'COUNTY'|'CITY'|'POSTCODE'|'LONGITUDE'|'LATITUDE'
     account_id : str, optional
         If None, the account ID will be inferred from your boto3 session.
     boto3_session : boto3.Session(), optional
@@ -335,7 +345,7 @@ def create_athena_dataset(
         "LogicalTableMap": {table_uuid: {"Alias": logical_table_alias, "Source": {"PhysicalTableId": table_uuid}}},
     }
     trans: List[Dict[str, Dict[str, Any]]] = _generate_transformations(
-        rename_columns=rename_columns, cast_columns_types=cast_columns_types
+        rename_columns=rename_columns, cast_columns_types=cast_columns_types, tag_columns=tag_columns
     )
     if trans:
         args["LogicalTableMap"][table_uuid]["DataTransforms"] = trans

tests/_utils.py

@@ -28,7 +28,7 @@ def get_df(governed=False):
             "ddouble": [0.0, None, 1.1],
             "decimal": [Decimal((0, (1, 9, 9), -2)), None, Decimal((0, (1, 9, 0), -2))],
             "string_object": ["foo", None, "boo"],
-            "string": ["foo", None, "boo"],
+            "string": ["Seattle", None, "Washington"],
             "date": [dt("2020-01-01"), None, dt("2020-01-02")],
             "timestamp": [ts("2020-01-01 00:00:00.0"), None, ts("2020-01-02 00:00:01.0")],
             "bool": [True, None, False],

tests/test_quicksight.py

@@ -37,6 +37,7 @@ def test_quicksight(path, glue_database, glue_table):
         allowed_to_manage=[wr.sts.get_current_identity_name()],
         rename_columns={"iint16": "new_col"},
         cast_columns_types={"new_col": "STRING"},
+        tag_columns={"string": {"ColumnGeographicRole": "CITY"}},
     )
     assert wr.quicksight.describe_dataset("test-table")["Name"] == "test-table"
 
@@ -49,6 +50,7 @@ def test_quicksight(path, glue_database, glue_table):
         allowed_to_manage=[wr.sts.get_current_identity_name()],
         rename_columns={"iint16": "new_col"},
         cast_columns_types={"new_col": "STRING"},
+        tag_columns={"string": {"ColumnGeographicRole": "CITY"}},
         tags={"foo": "boo"},
     )
 

tutorials/018 - QuickSight.ipynb

@@ -1242,6 +1242,9 @@
     "    cast_columns_types={\n",
     "        \"Count_of_Cases\": \"INTEGER\"\n",
     "    },\n",
+    "    tag_columns={\n",
+    "        \"combined_key\": {\"ColumnGeographicRole\": \"COUNTY\"}\n",
+    "    },\n",
     "    allowed_to_manage=[\"dev\"]\n",
     ")"
    ]