Adap 953/merge agate lazy load (#1063)

* lazy load agate

* Add test and documentation.

* Fix test.

* Fix test.

* Don't need a test for this.

---------

Co-authored-by: dwreeves <[email protected]>
Co-authored-by: Mila Page <[email protected]>

.changes/unreleased/Under the Hood-20240327-001304.yaml

@@ -0,0 +1,6 @@
+kind: Under the Hood
+body: Lazy load agate
+time: 2024-03-27T00:13:04.246062-04:00
+custom:
+  Author: dwreeves
+  Issue: "953"

dbt/adapters/snowflake/connections.py

@@ -8,10 +8,8 @@
 from dataclasses import dataclass
 from io import StringIO
 from time import sleep
-from typing import Any, List, Iterable, Optional, Tuple, Union
 
-import agate
-from dbt_common.clients.agate_helper import empty_table
+from typing import Optional, Tuple, Union, Any, List, Iterable, TYPE_CHECKING
 
 from cryptography.hazmat.backends import default_backend
 from cryptography.hazmat.primitives import serialization
@@ -46,6 +44,9 @@
 from dbt.adapters.events.types import AdapterEventWarning, AdapterEventError
 from dbt_common.ui import line_wrap_message, warning_tag
 
+if TYPE_CHECKING:
+    import agate
+
 
 logger = AdapterLogger("Snowflake")
 
@@ -494,9 +495,11 @@ def process_results(cls, column_names, rows):
 
     def execute(
         self, sql: str, auto_begin: bool = False, fetch: bool = False, limit: Optional[int] = None
-    ) -> Tuple[AdapterResponse, agate.Table]:
+    ) -> Tuple[AdapterResponse, "agate.Table"]:
         # don't apply the query comment here
         # it will be applied after ';' queries are split
+        from dbt_common.clients.agate_helper import empty_table
+
         _, cursor = self.add_query(sql, auto_begin)
         response = self.get_response(cursor)
         if fetch:

dbt/adapters/snowflake/impl.py

@@ -1,7 +1,6 @@
 from dataclasses import dataclass
-from typing import Mapping, Any, Optional, List, Union, Dict, FrozenSet, Tuple
+from typing import Mapping, Any, Optional, List, Union, Dict, FrozenSet, Tuple, TYPE_CHECKING
 
-import agate
 
 from dbt.adapters.base.impl import AdapterConfig, ConstraintSupport
 from dbt.adapters.base.meta import available
@@ -19,6 +18,9 @@
 from dbt_common.exceptions import CompilationError, DbtDatabaseError, DbtRuntimeError
 from dbt_common.utils import filter_null_values
 
+if TYPE_CHECKING:
+    import agate
+
 
 @dataclass
 class SnowflakeConfig(AdapterConfig):
@@ -63,8 +65,8 @@ def date_function(cls):
 
     @classmethod
     def _catalog_filter_table(
-        cls, table: agate.Table, used_schemas: FrozenSet[Tuple[str, str]]
-    ) -> agate.Table:
+        cls, table: "agate.Table", used_schemas: FrozenSet[Tuple[str, str]]
+    ) -> "agate.Table":
         # On snowflake, users can set QUOTED_IDENTIFIERS_IGNORE_CASE, so force
         # the column names to their lowercased forms.
         lowered = table.rename(column_names=[c.lower() for c in table.column_names])
@@ -150,7 +152,7 @@ def list_relations_without_caching(
 
         return [self._parse_list_relations_result(result) for result in results.select(columns)]
 
-    def _parse_list_relations_result(self, result: agate.Row) -> SnowflakeRelation:
+    def _parse_list_relations_result(self, result: "agate.Row") -> SnowflakeRelation:
         # this can be reduced to always including `is_dynamic` once bundle `2024_03` is mandatory
         try:
             database, schema, identifier, relation_type, is_dynamic = result
@@ -194,7 +196,7 @@ def quote_seed_column(self, column: str, quote_config: Optional[bool]) -> str:
             return column
 
     @available
-    def standardize_grants_dict(self, grants_table: agate.Table) -> dict:
+    def standardize_grants_dict(self, grants_table: "agate.Table") -> dict:
         grants_dict: Dict[str, Any] = {}
 
         for row in grants_table:

dbt/adapters/snowflake/relation_configs/base.py

@@ -1,6 +1,5 @@
 from dataclasses import dataclass
-from typing import Any, Dict, Optional
-import agate
+from typing import Any, Dict, Optional, TYPE_CHECKING
 from dbt.adapters.base.relation import Policy
 from dbt.adapters.relation_configs import (
     RelationConfigBase,
@@ -14,6 +13,10 @@
     SnowflakeQuotePolicy,
 )
 
+if TYPE_CHECKING:
+    # Imported downfile for specific row gathering function.
+    import agate
+
 
 @dataclass(frozen=True, eq=True, unsafe_hash=True)
 class SnowflakeRelationConfigBase(RelationConfigBase):
@@ -62,8 +65,10 @@ def _render_part(cls, component: ComponentName, value: Optional[str]) -> Optiona
         return None
 
     @classmethod
-    def _get_first_row(cls, results: agate.Table) -> agate.Row:
+    def _get_first_row(cls, results: "agate.Table") -> "agate.Row":
         try:
             return results.rows[0]
         except IndexError:
+            import agate
+
             return agate.Row(values=set())

dbt/adapters/snowflake/relation_configs/dynamic_table.py

@@ -1,13 +1,15 @@
 from dataclasses import dataclass
-from typing import Optional, Dict, Any
+from typing import Optional, Dict, Any, TYPE_CHECKING
 
-import agate
 from dbt.adapters.relation_configs import RelationConfigChange, RelationResults
 from dbt.adapters.contracts.relation import RelationConfig
 from dbt.adapters.contracts.relation import ComponentName
 
 from dbt.adapters.snowflake.relation_configs.base import SnowflakeRelationConfigBase
 
+if TYPE_CHECKING:
+    import agate
+
 
 @dataclass(frozen=True, eq=True, unsafe_hash=True)
 class SnowflakeDynamicTableConfig(SnowflakeRelationConfigBase):
@@ -62,7 +64,7 @@ def parse_relation_config(cls, relation_config: RelationConfig) -> Dict[str, Any
 
     @classmethod
     def parse_relation_results(cls, relation_results: RelationResults) -> Dict:
-        dynamic_table: agate.Row = relation_results["dynamic_table"].rows[0]
+        dynamic_table: "agate.Row" = relation_results["dynamic_table"].rows[0]
 
         config_dict = {
             "name": dynamic_table.get("name"),