Repository: dbt-athena/dbt-athena
Pull request #449 (2023-10-10 11:35:22)

feat(persist_docs): write dbt model and project metadata to properties
# Description

Hi team, I suggest improvement which should significantly improve data observability from my point of view.

If `persist_docs` is enabled we write additional metadata to Glue Table Properties:

- all fields from `model.config.meta` (i.e. owner, unique_key, other custom fields)
- `model.config.materialized` value (to mark incremental tables)
- dbt project name and version (to identify tables depoyed via dbt and find outdated tables by project version)

This improvement allows to treat Glue Catalog as a storage of tabel metadata and build interactive application around it without need to access dbt manifest.

## Models used to test - Optional
Tested on view, table, incremental table

## 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
- [X] You added functional testing when necessary


roslovets: > Overall the codes looks good, but it seems like a lot of new metadata, would it be possible to think of a way to activate or not this feature or to be at least able to customize which keys we want to push ?
> 
> Also please provide some unit tests :)

Well, I thought about it... Then I decided that Athena and Glue store lots of trash in Table Properties you will never use. Why should we be humble then?
In general I can add global switch for metadata persisting. But isn't it too complicated for such simple feature?
I will add some tests, thanks.
Jrmyy: Ok ok let's go for this then :)

> there is a max length in terms of string size for each property value.
I think there is something like 255 characters
roslovets: > > In general I can add global switch for metadata persisting. But isn't it too complicated for such simple feature?
> 
> I believe that is not necessary for this case.
> I'm wondering if we know if there are some sort of limitations in terms of the amount of properties that we can see and if there is a max length in terms of string size for each property value.

It's a very good point.
I've found limitations here: https://docs.aws.amazon.com/glue/latest/webapi/API_Table.html

I'm going to truncate long values and omit long keys.

Patch:
diff --git a/dbt/adapters/athena/impl.py b/dbt/adapters/athena/impl.py
index 9cac97f..9bd0a02 100755
--- a/dbt/adapters/athena/impl.py
+++ b/dbt/adapters/athena/impl.py
@@ -51,11 +51,14 @@ from dbt.adapters.athena.utils import (
     get_catalog_id,
     get_catalog_type,
     get_chunks,
+    is_valid_table_parameter_key,
+    stringify_table_parameter_value,
 )
 from dbt.adapters.base import ConstraintSupport, available
 from dbt.adapters.base.impl import AdapterConfig
 from dbt.adapters.base.relation import BaseRelation, InformationSchema
 from dbt.adapters.sql import SQLAdapter
+from dbt.config.runtime import RuntimeConfig
 from dbt.contracts.graph.manifest import Manifest
 from dbt.contracts.graph.nodes import CompiledNode, ConstraintType
 from dbt.exceptions import DbtRuntimeError
@@ -810,11 +813,12 @@ class AthenaAdapter(SQLAdapter):
             glue_client = client.session.client("glue", region_name=client.region_name, config=get_boto3_config())
 
         # By default, there is no need to update Glue Table
-        need_udpate_table = False
+        need_to_update_table = False
         # Get Table from Glue
         table = glue_client.get_table(CatalogId=catalog_id, DatabaseName=relation.schema, Name=relation.name)["Table"]
         # Prepare new version of Glue Table picking up significant fields
-        updated_table = self._get_table_input(table)
+        table_input = self._get_table_input(table)
+        table_parameters = table_input["Parameters"]
         # Update table description
         if persist_relation_docs:
             # Prepare dbt description
@@ -825,16 +829,40 @@ class AthenaAdapter(SQLAdapter):
             glue_table_comment = table["Parameters"].get("comment", "")
             # Update description if it's different
             if clean_table_description != glue_table_description or clean_table_description != glue_table_comment:
-                updated_table["Description"] = clean_table_description
-                updated_table_parameters: Dict[str, str] = dict(updated_table["Parameters"])
-                updated_table_parameters["comment"] = clean_table_description
-                updated_table["Parameters"] = updated_table_parameters
-                need_udpate_table = True
+                table_input["Description"] = clean_table_description
+                table_parameters["comment"] = clean_table_description
+                need_to_update_table = True
+
+            # Get dbt model meta if available
+            meta: Dict[str, Any] = model.get("config", {}).get("meta", {})
+            # Add some of dbt model config fields as table meta
+            meta["unique_id"] = model.get("unique_id")
+            meta["materialized"] = model.get("config", {}).get("materialized")
+            # Get dbt runtime config to be able to get dbt project metadata
+            runtime_config: RuntimeConfig = self.config
+            # Add dbt project metadata to table meta
+            meta["dbt_project_name"] = runtime_config.project_name
+            meta["dbt_project_version"] = runtime_config.version
+            # Prepare meta values for table properties and check if update is required
+            for meta_key, meta_value_raw in 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 that meta value is already attached to Glue table
+                        current_meta_value: Optional[str] = table_parameters.get(meta_key)
+                        if current_meta_value is None or current_meta_value != meta_value:
+                            # Update Glue table parameter only if needed
+                            table_parameters[meta_key] = meta_value
+                            need_to_update_table = True
+                    else:
+                        LOGGER.warning(f"Meta value for key '{meta_key}' is not supported and will be ignored")
+                else:
+                    LOGGER.warning(f"Meta key '{meta_key}' is not supported and will be ignored")
 
         # Update column comments
         if persist_column_docs:
             # Process every column
-            for col_obj in updated_table["StorageDescriptor"]["Columns"]:
+            for col_obj in table_input["StorageDescriptor"]["Columns"]:
                 # Get column description from dbt
                 col_name = col_obj["Name"]
                 if col_name in model["columns"]:
@@ -846,15 +874,16 @@ class AthenaAdapter(SQLAdapter):
                     # Update column description if it's different
                     if glue_col_comment != clean_col_comment:
                         col_obj["Comment"] = clean_col_comment
-                        need_udpate_table = True
+                        need_to_update_table = True
 
         # Update Glue Table only if table/column description is modified.
         # It prevents redundant schema version creating after incremental runs.
-        if need_udpate_table:
+        if need_to_update_table:
+            table_input["Parameters"] = table_parameters
             glue_client.update_table(
                 CatalogId=catalog_id,
                 DatabaseName=relation.schema,
-                TableInput=updated_table,
+                TableInput=table_input,
                 SkipArchive=skip_archive_table_version,
             )
 
diff --git a/dbt/adapters/athena/utils.py b/dbt/adapters/athena/utils.py
index dcd7491..0922f6d 100644
--- a/dbt/adapters/athena/utils.py
+++ b/dbt/adapters/athena/utils.py
@@ -1,14 +1,40 @@
+import json
+import re
 from enum import Enum
-from typing import Generator, List, Optional, TypeVar
+from typing import Any, Generator, List, Optional, TypeVar
 
 from mypy_boto3_athena.type_defs import DataCatalogTypeDef
 
+from dbt.adapters.athena.constants import LOGGER
+
 
 def clean_sql_comment(comment: str) -> str:
     split_and_strip = [line.strip() for line in comment.split("\n")]
     return " ".join(line for line in split_and_strip if line)
 
 
+def stringify_table_parameter_value(value: Any) -> Optional[str]:
+    """Convert any variable to string for Glue Table property."""
+    try:
+        if isinstance(value, (dict, list)):
+            value_str: str = json.dumps(value)
+        else:
+            value_str = str(value)
+        return value_str[:512000]
+    except (TypeError, ValueError) as e:
+        # Handle non-stringifiable objects and non-serializable objects
+        LOGGER.warning(f"Non-stringifiable object. Error: {str(e)}")
+        return None
+
+
+def is_valid_table_parameter_key(key: str) -> bool:
+    """Check if key is valid for Glue Table property according to official documentation."""
+    # Simplified version of key pattern which works with re
+    # Original pattern can be found here https://docs.aws.amazon.com/glue/latest/webapi/API_Table.html
+    key_pattern: str = r"^[\u0020-\uD7FF\uE000-\uFFFD\t]*$"
+    return len(key) <= 255 and bool(re.match(key_pattern, key))
+
+
 def get_catalog_id(catalog: Optional[DataCatalogTypeDef]) -> Optional[str]:
     return catalog["Parameters"]["catalog-id"] if catalog and catalog["Type"] == AthenaCatalogType.GLUE.value else None
 


Test patch:
diff --git a/tests/unit/test_utils.py b/tests/unit/test_utils.py
index 2385136..543558e 100644
--- a/tests/unit/test_utils.py
+++ b/tests/unit/test_utils.py
@@ -1,4 +1,9 @@
-from dbt.adapters.athena.utils import clean_sql_comment, get_chunks
+from dbt.adapters.athena.utils import (
+    clean_sql_comment,
+    get_chunks,
+    is_valid_table_parameter_key,
+    stringify_table_parameter_value,
+)
 
 
 def test_clean_comment():
@@ -14,6 +19,28 @@ def test_clean_comment():
     )
 
 
+def test_stringify_table_parameter_value():
+    class NonStringifiableObject:
+        def __str__(self):
+            raise ValueError("Non-stringifiable object")
+
+    assert stringify_table_parameter_value(True) == "True"
+    assert stringify_table_parameter_value(123) == "123"
+    assert stringify_table_parameter_value("dbt-athena") == "dbt-athena"
+    assert stringify_table_parameter_value(["a", "b", 3]) == '["a", "b", 3]'
+    assert stringify_table_parameter_value({"a": 1, "b": "c"}) == '{"a": 1, "b": "c"}'
+    assert len(stringify_table_parameter_value("a" * 512001)) == 512000
+    assert stringify_table_parameter_value(NonStringifiableObject()) is None
+    assert stringify_table_parameter_value([NonStringifiableObject()]) is None
+
+
+def test_is_valid_table_parameter_key():
+    assert is_valid_table_parameter_key("valid_key") is True
+    assert is_valid_table_parameter_key("Valid Key 123*!") is True
+    assert is_valid_table_parameter_key("invalid \n key") is False
+    assert is_valid_table_parameter_key("long_key" * 100) is False
+
+
 def test_get_chunks_empty():
     assert len(list(get_chunks([], 5))) == 0