chore: cache workgroup (#596)

dbt-labs · Mar 13, 2024 · d7a36e2 · d7a36e2
1 parent 289be4f
commit d7a36e2
Showing 1 changed file with 11 additions and 2 deletions.
diff --git a/dbt/adapters/athena/impl.py b/dbt/adapters/athena/impl.py
@@ -6,6 +6,7 @@
 import tempfile
 from dataclasses import dataclass
 from datetime import date, datetime
+from functools import lru_cache
 from itertools import chain
 from textwrap import dedent
 from threading import Lock
@@ -16,7 +17,8 @@
 import agate
 import mmh3
 from botocore.exceptions import ClientError
-from mypy_boto3_athena.type_defs import DataCatalogTypeDef
+from mypy_boto3_athena import AthenaClient
+from mypy_boto3_athena.type_defs import DataCatalogTypeDef, GetWorkGroupOutputTypeDef
 from mypy_boto3_glue.type_defs import (
     ColumnTypeDef,
     GetTableResponseTypeDef,
@@ -214,6 +216,13 @@ def apply_lf_grants(self, relation: AthenaRelation, lf_grants_config: Dict[str,
             lf_permissions.process_filters(lf_config)
             lf_permissions.process_permissions(lf_config)
 
+    @lru_cache()
+    def _get_work_group(self, client: AthenaClient, work_group: str) -> GetWorkGroupOutputTypeDef:
+        """
+        helper function to cache the result of the get_work_group to avoid APIs throttling
+        """
+        return client.get_work_group(WorkGroup=work_group)
+
     @available
     def is_work_group_output_location_enforced(self) -> bool:
         conn = self.connections.get_thread_connection()
@@ -228,7 +237,7 @@ def is_work_group_output_location_enforced(self) -> bool:
             )
 
         if creds.work_group:
-            work_group = athena_client.get_work_group(WorkGroup=creds.work_group)
+            work_group = self._get_work_group(athena_client, creds.work_group)
             output_location = (
                 work_group.get("WorkGroup", {})
                 .get("Configuration", {})