Repository: dbt-athena/dbt-athena
Pull request #644 (2024-05-09 08:29:08)

feat: Add column meta to glue column parameters
# Description

Modify `persist_docs_to_glue()` so that it can update [glue column parameters](https://docs.aws.amazon.com/glue/latest/webapi/API_Column.html) based on dbt column meta.

For example:
```
version: 2

models:
  - name: test_deduplicate
    description: another value
    config:
      persist_docs:
        relation: true
        columns: true
      meta: 
        test: banana
    columns:
      - name: id
        meta: 
          primary_key: true
```

See [slack](https://getdbt.slack.com/archives/C013MLFR7BQ/p1715167910378029) for more details

## Models used to test - Optional
Wasn't sure how to run `persist_docs_to_glue()` against my models

## Checklist

- [x] You followed [contributing section](https://github.com/dbt-athena/dbt-athena#contributing)
- [x] You kept your Pull Request small and focused on a single feature or bug fix.
- [x] You added unit testing when necessary
- [ ] You added functional testing when necessary


SoumayaMauthoorMOJ: @nicor88 updated the code :-) How do I run persist_docs_to_glue() against my models? Can I run via a post hook?
nicor88: @SoumayaMauthoorMOJ thanks :) 

The method that you modified is called inside [persist_docs](https://github.com/dbt-athena/dbt-athena/blob/c5b89912c901ab9ba1aec39db2223ac3cbfa1a0b/dbt/include/athena/macros/materializations/models/view/view.sql#L14) macro, and that is called in different places based on the type of materialisation:
* for table [here](https://github.com/dbt-athena/dbt-athena/blob/c5b89912c901ab9ba1aec39db2223ac3cbfa1a0b/dbt/include/athena/macros/materializations/models/table/table.sql#L171)
* for incremental runs [here](https://github.com/dbt-athena/dbt-athena/blob/c5b89912c901ab9ba1aec39db2223ac3cbfa1a0b/dbt/include/athena/macros/materializations/models/incremental/incremental.sql#L156)
SoumayaMauthoorMOJ: I just tested locally (I hadn't realised you needed to add `persist_docs` to config) and it failed. It passed after I removed the commit. I need the [if condition](https://github.com/dbt-athena/dbt-athena/pull/644#discussion_r1595563563) because you can't add a dict to a missing dictionary e.g.:

```
>>> my_dict = {}
>>> my_dict['a']['b']=1
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
KeyError: 'a'
```
so I need to add `Parameters` as a key if it doesn't exit, but before I start looping through the different parameters
nicor88: @SoumayaMauthoorMOJ yeah let's remove the check on model description, in case for example we have the meta field it should work. Good catch. Regarding the functional tests I can only deep dive next week. 
nicor88: @SoumayaMauthoorMOJ the issue with functional tests was a not related issue. You right, the check on Parameters being empty needs to stay, if you see my comments I simply suggested to move it to an another part of the code, that imho is more appropriate 
SoumayaMauthoorMOJ: @nicor88 done :-)

Patch:
diff --git a/README.md b/README.md
index 4daad68..e345c03 100644
--- a/README.md
+++ b/README.md
@@ -33,6 +33,7 @@
     - [Iceberg](#iceberg)
     - [Highly available table (HA)](#highly-available-table-ha)
       - [HA Known issues](#ha-known-issues)
+    - [Update glue data catalog](#update-glue-data-catalog)
   - [Snapshots](#snapshots)
     - [Timestamp strategy](#timestamp-strategy)
     - [Check strategy](#check-strategy)
@@ -505,6 +506,34 @@ By default, the materialization keeps the last 4 table versions, you can change
 - By default, Glue "duplicates" the versions internally, so the last two versions of a table point to the same location
 - It's recommended to have `versions_to_keep` >= 4, as this will avoid having the older location removed
 
+### Update glue data catalog
+
+Optionally persist resource descriptions as column and relation comments to the glue data catalog, and meta as
+[glue table properties](https://docs.aws.amazon.com/glue/latest/dg/tables-described.html#table-properties)
+and [column parameters](https://docs.aws.amazon.com/glue/latest/webapi/API_Column.html).
+By default, documentation persistence is disabled, but it can be enabled for specific resources or
+groups of resources as needed.
+
+For example:
+
+```yaml
+models:
+  - name: test_deduplicate
+    description: another value
+    config:
+      persist_docs:
+        relation: true
+        columns: true
+      meta:
+        test: value
+    columns:
+      - name: id
+        meta:
+          primary_key: true
+```
+
+See [persist docs](https://docs.getdbt.com/reference/resource-configs/persist_docs) for more details.
+
 ## Snapshots
 
 The adapter supports snapshot materialization. It supports both timestamp and check strategy. To create a snapshot
diff --git a/dbt/adapters/athena/impl.py b/dbt/adapters/athena/impl.py
index 2986569..76266f5 100755
--- a/dbt/adapters/athena/impl.py
+++ b/dbt/adapters/athena/impl.py
@@ -712,7 +712,7 @@ class AthenaAdapter(SQLAdapter):
             else:
                 raise e
 
-        relations: list[BaseRelation] = []
+        relations: List[BaseRelation] = []
         quote_policy = {"database": True, "schema": True, "identifier": True}
         for table in tables:
             if "TableType" not in table:
@@ -956,6 +956,7 @@ class AthenaAdapter(SQLAdapter):
         # Prepare new version of Glue Table picking up significant fields
         table_input = self._get_table_input(table)
         table_parameters = table_input["Parameters"]
+
         # Update table description
         if persist_relation_docs:
             # Prepare dbt description
@@ -1013,6 +1014,29 @@ class AthenaAdapter(SQLAdapter):
                     # Save column description from dbt
                     col_obj["Comment"] = clean_col_comment
 
+                    # Get dbt model column meta if available
+                    col_meta: Dict[str, Any] = model["columns"][col_name].get("meta", {})
+                    # Add empty Parameters dictionary if not present
+                    if col_meta and "Parameters" not in col_obj.keys():
+                        col_obj["Parameters"] = {}
+                    # Prepare meta values for column properties and check if update is required
+                    for meta_key, meta_value_raw in col_meta.items():
+                        if is_valid_table_parameter_key(meta_key):
+                            meta_value = stringify_table_parameter_value(meta_value_raw)
+                            if meta_value is not None:
+                                # Check if meta value is already attached to Glue column
+                                col_current_meta_value: Optional[str] = col_obj["Parameters"].get(meta_key)
+                                if col_current_meta_value is None or col_current_meta_value != meta_value:
+                                    need_to_update_table = True
+                                # Save Glue column parameter
+                                col_obj["Parameters"][meta_key] = meta_value
+                            else:
+                                LOGGER.warning(
+                                    f"Column meta value for key '{meta_key}' is not supported and will be ignored"
+                                )
+                        else:
+                            LOGGER.warning(f"Column meta key '{meta_key}' is not supported and will be ignored")
+
         # Update Glue Table only if table/column description is modified.
         # It prevents redundant schema version creating after incremental runs.
         if need_to_update_table:
diff --git a/dbt/include/athena/macros/adapters/persist_docs.sql b/dbt/include/athena/macros/adapters/persist_docs.sql
index 6d9182b..7b95bd7 100644
--- a/dbt/include/athena/macros/adapters/persist_docs.sql
+++ b/dbt/include/athena/macros/adapters/persist_docs.sql
@@ -1,5 +1,5 @@
 {% macro athena__persist_docs(relation, model, for_relation=true, for_columns=true) -%}
-  {% set persist_relation_docs = for_relation and config.persist_relation_docs() and model.description %}
+  {% set persist_relation_docs = for_relation and config.persist_relation_docs()%}
   {% set persist_column_docs = for_columns and config.persist_column_docs() and model.columns %}
   {% if persist_relation_docs or persist_column_docs %}
     {% do adapter.persist_docs_to_glue(


Test patch:
diff --git a/tests/unit/test_adapter.py b/tests/unit/test_adapter.py
index 4abbbb2..fa4aa23 100644
--- a/tests/unit/test_adapter.py
+++ b/tests/unit/test_adapter.py
@@ -970,6 +970,7 @@ class TestAthenaAdapter:
                     """,
                 "columns": {
                     "id": {
+                        "meta": {"primary_key": "true"},
                         "description": """
                         A column with str, 123, &^% \" and '
 
@@ -986,6 +987,7 @@ class TestAthenaAdapter:
         assert not table.get("Description", "")
         assert not table["Parameters"].get("comment")
         assert all(not col.get("Comment") for col in table["StorageDescriptor"]["Columns"])
+        assert all(not col.get("Parameters") for col in table["StorageDescriptor"]["Columns"])
 
     @mock_aws
     def test_persist_docs_to_glue_comment(self, mock_aws_service):
@@ -1009,6 +1011,7 @@ class TestAthenaAdapter:
                     """,
                 "columns": {
                     "id": {
+                        "meta": {"primary_key": True},
                         "description": """
                         A column with str, 123, &^% \" and '
 
@@ -1026,6 +1029,7 @@ class TestAthenaAdapter:
         assert table["Parameters"]["comment"] == "A table with str, 123, &^% \" and ' and an other paragraph."
         col_id = [col for col in table["StorageDescriptor"]["Columns"] if col["Name"] == "id"][0]
         assert col_id["Comment"] == "A column with str, 123, &^% \" and ' and an other paragraph."
+        assert col_id["Parameters"] == {"primary_key": "True"}
 
     @mock_aws
     def test_list_schemas(self, mock_aws_service):