{#- Snowcap dbt Governance Macros -#}
{#- Generated by: snowcap generate dbt-macros -#}
{#- Documentation: https://snowcap.datacoves.com -#}

{% macro snowcap_apply_policies() %}
  {{ snowcap_apply_masking_tags() }}
{% endmacro %}

{% macro snowcap_apply_masking_tags() %}
  {% set tag_db = var('snowcap_tag_database', target.database) %}
  {% set tag_schema = var('snowcap_tag_schema', 'TAGS') %}
  {% set columns = model.columns.values() if model.columns else [] %}

  {#- Build a map of column -> desired masking tag -#}
  {% set desired_tags = {} %}
  {% for column in columns %}
    {% set meta = column.meta if column.meta is mapping else {} %}
    {% set masking_tag = meta.get('masking_tag') %}
    {% do desired_tags.update({column.name | upper: masking_tag}) %}
  {% endfor %}

  {#- Query current tags on this table's columns from our tag schema -#}
  {% set current_tags_query %}
    SELECT COLUMN_NAME, TAG_NAME
    FROM SNOWFLAKE.ACCOUNT_USAGE.TAG_REFERENCES
    WHERE OBJECT_DATABASE = '{{ this.database | upper }}'
      AND OBJECT_SCHEMA = '{{ this.schema | upper }}'
      AND OBJECT_NAME = '{{ this.identifier | upper }}'
      AND DOMAIN = 'COLUMN'
      AND TAG_DATABASE = '{{ tag_db | upper }}'
      AND TAG_SCHEMA = '{{ tag_schema | upper }}'
  {% endset %}
  {% set current_tags_result = run_query(current_tags_query) %}

  {#- Unset tags that should no longer be present -#}
  {% for row in current_tags_result %}
    {% set col_name = row['COLUMN_NAME'] %}
    {% set current_tag = row['TAG_NAME'] %}
    {% set desired_tag = desired_tags.get(col_name) %}

    {#- If column has no desired tag, or has a different tag, unset current -#}
    {% if not desired_tag or desired_tag | upper != current_tag | upper %}
      ALTER TABLE {{ this }} MODIFY COLUMN {{ col_name }}
        UNSET TAG {{ tag_db }}.{{ tag_schema }}.{{ current_tag }};
    {% endif %}
  {% endfor %}

  {#- Set desired masking tags -#}
  {% for column in columns %}
    {% set meta = column.meta if column.meta is mapping else {} %}
    {% set masking_tag = meta.get('masking_tag') %}

    {% if masking_tag %}
      ALTER TABLE {{ this }} MODIFY COLUMN {{ column.name }}
        SET TAG {{ tag_db }}.{{ tag_schema }}.{{ masking_tag }} = 'true';
    {% endif %}
  {% endfor %}
{% endmacro %}

