def get_columns_and_policy_tags()

in helpers/bq-remote-functions/get-policy-tags/main.py [0:0]

• 15 lines of code
• 8 McCabe index (conditional complexity)

def get_columns_and_policy_tags(project_id, dataset_id, table_id):
    """Fetches all columns and their associated policy tags associated for a BigQuery table columns."""
    bq_client = bigquery.Client(project=project_id)
    table_ref = bq_client.dataset(dataset_id).table(table_id)

    try:
        table = bq_client.get_table(table_ref)
    except Exception as e:
        # return no results and the exception
        return None, e

    columns_and_policy_tags = {}
    for field in table.schema:
        if field.policy_tags:
            # field.policy_tags.names is an iterator, but it will contain max one policy tag
            for tag in field.policy_tags.names:
                columns_and_policy_tags[field.name] = tag
        else:
            # if no policy tag attached, still include the column in the final output
            columns_and_policy_tags[field.name] = None
    # return the results and no exception
    return columns_and_policy_tags, None