From 1947ebbef5f77c24c63748d1f2e96bd88b8a363b Mon Sep 17 00:00:00 2001 From: Serhii Dimchenko Date: Mon, 16 Dec 2024 15:07:56 +0100 Subject: [PATCH 1/7] Implement drop_glue_database macro and method --- dbt-athena/src/dbt/adapters/athena/impl.py | 20 ++++++++++++++++++- .../include/athena/macros/adapters/schema.sql | 5 +++++ 2 files changed, 24 insertions(+), 1 deletion(-) diff --git a/dbt-athena/src/dbt/adapters/athena/impl.py b/dbt-athena/src/dbt/adapters/athena/impl.py index e70cca3c..69e0e328 100755 --- a/dbt-athena/src/dbt/adapters/athena/impl.py +++ b/dbt-athena/src/dbt/adapters/athena/impl.py @@ -9,7 +9,7 @@ from functools import lru_cache from textwrap import dedent from threading import Lock -from typing import Any, Dict, FrozenSet, Iterable, List, Optional, Set, Tuple, Type +from typing import TYPE_CHECKING, Any, Dict, FrozenSet, Iterable, List, Optional, Set, Tuple, Type from urllib.parse import urlparse from uuid import uuid4 @@ -69,6 +69,10 @@ from dbt.adapters.contracts.relation import RelationConfig from dbt.adapters.sql import SQLAdapter + +if TYPE_CHECKING: + from mypy_boto3_glue.client import GlueClient + boto3_client_lock = Lock() @@ -1173,6 +1177,20 @@ def delete_from_glue_catalog(self, relation: AthenaRelation) -> None: LOGGER.error(e) raise e + @available + def drop_glue_database(self, database_name: str, catalog_id: Optional[str] = None) -> None: + conn = self.connections.get_thread_connection() + creds = conn.credentials + client = conn.handle + + with boto3_client_lock: + glue_client: GlueClient = client.session.client( + "glue", + region_name=client.region_name, + config=get_boto3_config(num_retries=creds.effective_num_retries), + ) + glue_client.delete_database(Name=database_name, CatalogId=catalog_id) + @available.parse_none def valid_snapshot_target(self, relation: BaseRelation) -> None: """Log an error to help developers migrate to the new snapshot logic""" diff --git a/dbt-athena/src/dbt/include/athena/macros/adapters/schema.sql b/dbt-athena/src/dbt/include/athena/macros/adapters/schema.sql index 2750a7f9..c98d7963 100644 --- a/dbt-athena/src/dbt/include/athena/macros/adapters/schema.sql +++ b/dbt-athena/src/dbt/include/athena/macros/adapters/schema.sql @@ -13,3 +13,8 @@ drop schema if exists {{ relation.without_identifier().render_hive() }} cascade {% endcall %} {% endmacro %} + + +{% macro drop_glue_database(database_name, catalog_id) -%} + {{ adapter.drop_glue_database(database_name, catalog_id) }} +{% endmacro %} From 0f23cda27b27db8b90e13a9131d39bd9925dd083 Mon Sep 17 00:00:00 2001 From: Serhii Dimchenko Date: Mon, 16 Dec 2024 15:56:30 +0100 Subject: [PATCH 2/7] Add logs --- dbt-athena/src/dbt/adapters/athena/impl.py | 5 ++++- dbt-athena/src/dbt/include/athena/macros/adapters/schema.sql | 2 +- 2 files changed, 5 insertions(+), 2 deletions(-) diff --git a/dbt-athena/src/dbt/adapters/athena/impl.py b/dbt-athena/src/dbt/adapters/athena/impl.py index 69e0e328..2d3aadba 100755 --- a/dbt-athena/src/dbt/adapters/athena/impl.py +++ b/dbt-athena/src/dbt/adapters/athena/impl.py @@ -1178,10 +1178,12 @@ def delete_from_glue_catalog(self, relation: AthenaRelation) -> None: raise e @available - def drop_glue_database(self, database_name: str, catalog_id: Optional[str] = None) -> None: + def drop_glue_database(self, database_name: str, catalog_name: str = "awsdatacatalog") -> None: conn = self.connections.get_thread_connection() creds = conn.credentials client = conn.handle + catalog = self._get_data_catalog(catalog_name) + catalog_id = get_catalog_id(catalog) with boto3_client_lock: glue_client: GlueClient = client.session.client( @@ -1190,6 +1192,7 @@ def drop_glue_database(self, database_name: str, catalog_id: Optional[str] = Non config=get_boto3_config(num_retries=creds.effective_num_retries), ) glue_client.delete_database(Name=database_name, CatalogId=catalog_id) + LOGGER.info(f"Glue database successfully deleted: {catalog_name}.{database_name}") @available.parse_none def valid_snapshot_target(self, relation: BaseRelation) -> None: diff --git a/dbt-athena/src/dbt/include/athena/macros/adapters/schema.sql b/dbt-athena/src/dbt/include/athena/macros/adapters/schema.sql index c98d7963..094acd5c 100644 --- a/dbt-athena/src/dbt/include/athena/macros/adapters/schema.sql +++ b/dbt-athena/src/dbt/include/athena/macros/adapters/schema.sql @@ -15,6 +15,6 @@ {% endmacro %} -{% macro drop_glue_database(database_name, catalog_id) -%} +{% macro drop_glue_database(database_name, catalog_id='awsdatacatalog') -%} {{ adapter.drop_glue_database(database_name, catalog_id) }} {% endmacro %} From 40a16ffed81b33c775b1e9f0f20a40f3fe0b738c Mon Sep 17 00:00:00 2001 From: Serhii Dimchenko Date: Mon, 16 Dec 2024 16:00:27 +0100 Subject: [PATCH 3/7] Minor fix --- dbt-athena/src/dbt/include/athena/macros/adapters/schema.sql | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/dbt-athena/src/dbt/include/athena/macros/adapters/schema.sql b/dbt-athena/src/dbt/include/athena/macros/adapters/schema.sql index 094acd5c..a8a60fef 100644 --- a/dbt-athena/src/dbt/include/athena/macros/adapters/schema.sql +++ b/dbt-athena/src/dbt/include/athena/macros/adapters/schema.sql @@ -15,6 +15,6 @@ {% endmacro %} -{% macro drop_glue_database(database_name, catalog_id='awsdatacatalog') -%} - {{ adapter.drop_glue_database(database_name, catalog_id) }} +{% macro drop_glue_database(database_name, catalog_name='awsdatacatalog') -%} + {{ adapter.drop_glue_database(database_name, catalog_name) }} {% endmacro %} From b85c1c8c3aeb8e0aed08a9f5cd4c6153cfa14763 Mon Sep 17 00:00:00 2001 From: Serhii Dimchenko Date: Mon, 16 Dec 2024 18:11:57 +0100 Subject: [PATCH 4/7] Add unit test --- dbt-athena/tests/unit/test_adapter.py | 12 ++++++++++++ 1 file changed, 12 insertions(+) diff --git a/dbt-athena/tests/unit/test_adapter.py b/dbt-athena/tests/unit/test_adapter.py index e3c3b923..291979fd 100644 --- a/dbt-athena/tests/unit/test_adapter.py +++ b/dbt-athena/tests/unit/test_adapter.py @@ -12,6 +12,7 @@ from dbt_common.exceptions import ConnectionError, DbtRuntimeError from moto import mock_aws from moto.core import DEFAULT_ACCOUNT_ID +from mypy_boto3_glue.client import GlueClient from dbt.adapters.athena import AthenaAdapter from dbt.adapters.athena import Plugin as AthenaPlugin @@ -1252,6 +1253,17 @@ def test_format_unsupported_type(self): with pytest.raises(ValueError): self.adapter.format_value_for_partition("test", "unsupported_type") + @mock_aws + def test_drop_glue_database(self, mock_aws_service): + glue_client: GlueClient = boto3.client("glue", region_name=AWS_REGION) + test_input = {"Name": "test"} + glue_client.create_database(DatabaseInput=test_input) + database_list = glue_client.get_databases()["DatabaseList"] + assert [test_input["Name"]] == [db["Name"] for db in database_list] + self.adapter.acquire_connection("dummy") + self.adapter.drop_glue_database(database_name=test_input["Name"]) + assert glue_client.get_databases()["DatabaseList"] == [] + class TestAthenaFilterCatalog: def test__catalog_filter_table(self): From 91fd25a75bd904af21519dea90536713137b7441 Mon Sep 17 00:00:00 2001 From: Serhii Dimchenko Date: Mon, 16 Dec 2024 20:18:20 +0100 Subject: [PATCH 5/7] Fix log level --- dbt-athena/src/dbt/adapters/athena/impl.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/dbt-athena/src/dbt/adapters/athena/impl.py b/dbt-athena/src/dbt/adapters/athena/impl.py index 2d3aadba..f915de0c 100755 --- a/dbt-athena/src/dbt/adapters/athena/impl.py +++ b/dbt-athena/src/dbt/adapters/athena/impl.py @@ -1192,7 +1192,7 @@ def drop_glue_database(self, database_name: str, catalog_name: str = "awsdatacat config=get_boto3_config(num_retries=creds.effective_num_retries), ) glue_client.delete_database(Name=database_name, CatalogId=catalog_id) - LOGGER.info(f"Glue database successfully deleted: {catalog_name}.{database_name}") + LOGGER.debug(f"Glue database successfully deleted: {catalog_name}.{database_name}") @available.parse_none def valid_snapshot_target(self, relation: BaseRelation) -> None: From bf978d8fbf1d37a51e399585de6ced1c6457886d Mon Sep 17 00:00:00 2001 From: Serhii Dimchenko Date: Mon, 16 Dec 2024 20:21:04 +0100 Subject: [PATCH 6/7] Fix test args --- dbt-athena/tests/unit/test_adapter.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/dbt-athena/tests/unit/test_adapter.py b/dbt-athena/tests/unit/test_adapter.py index 291979fd..ad336f97 100644 --- a/dbt-athena/tests/unit/test_adapter.py +++ b/dbt-athena/tests/unit/test_adapter.py @@ -1254,7 +1254,7 @@ def test_format_unsupported_type(self): self.adapter.format_value_for_partition("test", "unsupported_type") @mock_aws - def test_drop_glue_database(self, mock_aws_service): + def test_drop_glue_database(self): glue_client: GlueClient = boto3.client("glue", region_name=AWS_REGION) test_input = {"Name": "test"} glue_client.create_database(DatabaseInput=test_input) From 69ced5ce7b7b4b69eeddc357159fb241b92baca1 Mon Sep 17 00:00:00 2001 From: Serhii Dimchenko Date: Tue, 17 Dec 2024 00:13:10 +0100 Subject: [PATCH 7/7] Add changelog file --- .../.changes/unreleased/Features-20241217-001238.yaml | 6 ++++++ 1 file changed, 6 insertions(+) create mode 100644 dbt-athena/.changes/unreleased/Features-20241217-001238.yaml diff --git a/dbt-athena/.changes/unreleased/Features-20241217-001238.yaml b/dbt-athena/.changes/unreleased/Features-20241217-001238.yaml new file mode 100644 index 00000000..e8da240b --- /dev/null +++ b/dbt-athena/.changes/unreleased/Features-20241217-001238.yaml @@ -0,0 +1,6 @@ +kind: Features +body: Implement drop_glue_database macro and method +time: 2024-12-17T00:12:38.943084+01:00 +custom: + Author: svdimchenko + Issue: "768"