gistfile1.py

import logging

import google.auth
from google.api_core.exceptions import AlreadyExists
from google.cloud import datacatalog
from rich.logging import RichHandler

logging.basicConfig(level=logging.INFO, handlers=[RichHandler()])

log = logging.getLogger(__name__)

credentials, PROJECT = google.auth.default()
client = datacatalog.DataCatalogClient()

LOCATION = "us-central1"
ENTRY_GROUP_ID = "important_data_jan_2019"
ENTRY_ID = "python_files"
TEMPLATE_ID = "template_id"
FIELD_NAME_1 = "first"
FIELD_NAME_2 = "second"
FIELD_NAME_3 = "first-rename"

LOCATION_PATH = f"projects/{PROJECT}/locations/{LOCATION}"
ENTRY_PATH = f"projects/{PROJECT}/locations/{LOCATION}/entryGroups/{ENTRY_GROUP_ID}/entries/{ENTRY_ID}"
ENTRY_GROUP_PATH = f"projects/{PROJECT}/locations/{LOCATION}/entryGroups/{ENTRY_GROUP_ID}"
TAG_TEMPLATE_PATH = f"projects/{PROJECT}/locations/{LOCATION}/tagTemplates/{TEMPLATE_ID}"


try:
    entry_group = client.create_entry_group(
        request=datacatalog.CreateEntryGroupRequest(
            parent=LOCATION_PATH,
            entry_group_id=ENTRY_GROUP_ID,
            entry_group=datacatalog.EntryGroup(
                name=ENTRY_GROUP_PATH,
                display_name="analytics data - jan 2011",
            ),
        )
    )
except AlreadyExists as e:
    log.info("%s", e)

    entry_group = client.get_entry_group(request=datacatalog.GetEntryGroupRequest(name=ENTRY_GROUP_PATH))

log.info("entry_group=%s", entry_group)

try:
    entry = client.create_entry(
        request=datacatalog.CreateEntryRequest(
            parent=ENTRY_GROUP_PATH,
            entry_id=ENTRY_ID,
            entry={
                "display_name": "Wizard",
                "type_": "FILESET",
                "gcs_fileset_spec": {"file_patterns": ["gs://test-datacatalog/**"]},
            },
        )
    )
except AlreadyExists as e:
    log.info("%s", e)
    entry = client.get_entry(request=datacatalog.GetEntryRequest(name=ENTRY_PATH))

log.info("entry=%s", entry)

try:
    # request_good = datacatalog.CreateTagTemplateRequest(
    #     parent=LOCATION_PATH,
    #     tag_template_id=TEMPLATE_ID,
    #     tag_template=datacatalog.TagTemplate(
    #         display_name="Awesome Tag template",
    #         fields={
    #             FIELD_NAME_1: datacatalog.TagTemplateField(
    #                 display_name="first-field", type_=dict(primitive_type="STRING")
    #             )
    #         },
    #     ),
    # )

    request = dict(
        parent=LOCATION_PATH,
        tag_template_id=TEMPLATE_ID,
        tag_template=dict(
            display_name="Awesome Tag template",
            fields={FIELD_NAME_1: dict(display_name="first-field", type_=dict(primitive_type="STRING"))},
        ),
    )
    # HACK: To create tag template with dict
    # request = datacatalog.CreateTagTemplateRequest(
    #     parent=request.get("parent"),
    #     tag_template_id=request.get("tag_template_id"),
    #     tag_template=(
    #         datacatalog.TagTemplate(
    #             display_name=request.get("tag_template", {}).get("display_name"),
    #             fields={
    #                 k: datacatalog.TagTemplateField(**v)
    #                 for k, v in request.get("tag_template", {}).get("fields", {}).items()
    #             },
    #         )
    #         if request.get("tag_template")
    #         else None
    #     ),
    # )

    log.info("create_tag_template(request=%s)", request)
    tag_template = client.create_tag_template(request=request)
except AlreadyExists as e:
    log.info("%s", e)
    tag_template = client.get_tag_template(request=datacatalog.GetTagTemplateRequest(name=TAG_TEMPLATE_PATH))

log.info("tag_template=%s", tag_template)
client.delete_tag_template(name=TAG_TEMPLATE_PATH, force=True)
client.delete_entry(name=ENTRY_PATH)
client.delete_entry_group(name=ENTRY_GROUP_PATH)