Adap 746/merge agate lazy load (#831)

* lazy load agate

* Add test and documentation.

* Alter doc strings.

* Lint.

* 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-20240331-103115.yaml

@@ -0,0 +1,6 @@
+kind: Under the Hood
+body: Lazy load agate
+time: 2024-03-31T10:31:15.65006-04:00
+custom:
+  Author: dwreeves
+  Issue: "745"

dbt/adapters/redshift/connections.py

@@ -1,14 +1,12 @@
 import re
 from multiprocessing import Lock
 from contextlib import contextmanager
-from typing import Any, Callable, Dict, Tuple, Union, Optional, List
+from typing import Any, Callable, Dict, Tuple, Union, Optional, List, TYPE_CHECKING
 from dataclasses import dataclass, field
 
-import agate
 import sqlparse
 import redshift_connector
 from dbt.adapters.exceptions import FailedToConnectError
-from dbt_common.clients import agate_helper
 from redshift_connector.utils.oids import get_datatype_name
 
 from dbt.adapters.sql import SQLConnectionManager
@@ -19,6 +17,11 @@
 from dbt_common.helper_types import Port
 from dbt_common.exceptions import DbtRuntimeError, CompilationError, DbtDatabaseError
 
+if TYPE_CHECKING:
+    # Indirectly imported via agate_helper, which is lazy loaded further downfile.
+    # Used by mypy for earlier type hints.
+    import agate
+
 
 class SSLConfigError(CompilationError):
     def __init__(self, exc: ValidationError):
@@ -393,13 +396,15 @@ def execute(
         auto_begin: bool = False,
         fetch: bool = False,
         limit: Optional[int] = None,
-    ) -> Tuple[AdapterResponse, agate.Table]:
+    ) -> Tuple[AdapterResponse, "agate.Table"]:
         sql = self._add_query_comment(sql)
         _, cursor = self.add_query(sql, auto_begin)
         response = self.get_response(cursor)
         if fetch:
             table = self.get_result_from_cursor(cursor, limit)
         else:
+            from dbt_common.clients import agate_helper
+
             table = agate_helper.empty_table()
         return response, table
 

dbt/adapters/redshift/impl.py

@@ -1,7 +1,7 @@
 import os
 from dataclasses import dataclass
 from dbt_common.contracts.constraints import ConstraintType
-from typing import Optional, Set, Any, Dict, Type
+from typing import Optional, Set, Any, Dict, Type, TYPE_CHECKING
 from collections import namedtuple
 from dbt.adapters.base import PythonJobHelper
 from dbt.adapters.base.impl import AdapterConfig, ConstraintSupport
@@ -28,6 +28,9 @@
 
 GET_RELATIONS_MACRO_NAME = "redshift__get_relations"
 
+if TYPE_CHECKING:
+    import agate
+
 
 @dataclass
 class RedshiftConfig(AdapterConfig):
@@ -85,7 +88,7 @@ def drop_relation(self, relation):
             return super().drop_relation(relation)
 
     @classmethod
-    def convert_text_type(cls, agate_table, col_idx):
+    def convert_text_type(cls, agate_table: "agate.Table", col_idx):
         column = agate_table.columns[col_idx]
         # `lens` must be a list, so this can't be a generator expression,
         # because max() raises ane exception if its argument has no members.
@@ -94,7 +97,7 @@ def convert_text_type(cls, agate_table, col_idx):
         return "varchar({})".format(max_len)
 
     @classmethod
-    def convert_time_type(cls, agate_table, col_idx):
+    def convert_time_type(cls, agate_table: "agate.Table", col_idx):
         return "varchar(24)"
 
     @available

dbt/adapters/redshift/relation_configs/base.py

@@ -1,7 +1,6 @@
 from dataclasses import dataclass
-from typing import Optional, Dict
+from typing import Optional, Dict, TYPE_CHECKING
 
-import agate
 from dbt.adapters.base.relation import Policy
 from dbt.adapters.contracts.relation import ComponentName, RelationConfig
 from dbt.adapters.relation_configs import (
@@ -15,6 +14,10 @@
     RedshiftQuotePolicy,
 )
 
+if TYPE_CHECKING:
+    # Imported downfile for specific row gathering function.
+    import agate
+
 
 @dataclass(frozen=True, eq=True, unsafe_hash=True)
 class RedshiftRelationConfigBase(RelationConfigBase):
@@ -63,8 +66,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/redshift/relation_configs/dist.py

@@ -1,8 +1,7 @@
 from dataclasses import dataclass
 from dbt.adapters.contracts.relation import RelationConfig
-from typing import Optional, Set, Dict
+from typing import Optional, Set, Dict, TYPE_CHECKING
 
-import agate
 from dbt.adapters.relation_configs import (
     RelationConfigChange,
     RelationConfigChangeAction,
@@ -15,6 +14,9 @@
 
 from dbt.adapters.redshift.relation_configs.base import RedshiftRelationConfigBase
 
+if TYPE_CHECKING:
+    import agate
+
 
 class RedshiftDistStyle(StrEnum):
     auto = "auto"
@@ -108,7 +110,7 @@ def parse_relation_config(cls, relation_config: RelationConfig) -> dict:
         return config
 
     @classmethod
-    def parse_relation_results(cls, relation_results_entry: agate.Row) -> Dict:
+    def parse_relation_results(cls, relation_results_entry: "agate.Row") -> Dict:
         """
         Translate agate objects from the database into a standard dictionary.
 

dbt/adapters/redshift/relation_configs/materialized_view.py

@@ -1,7 +1,6 @@
 from dataclasses import dataclass, field
-from typing import Optional, Set, Dict, Any
+from typing import Optional, Set, Dict, Any, TYPE_CHECKING
 
-import agate
 from dbt.adapters.relation_configs import (
     RelationResults,
     RelationConfigChange,
@@ -25,6 +24,9 @@
 )
 from dbt.adapters.redshift.utility import evaluate_bool
 
+if TYPE_CHECKING:
+    import agate
+
 
 @dataclass(frozen=True, eq=True, unsafe_hash=True)
 class RedshiftMaterializedViewConfig(RedshiftRelationConfigBase, RelationConfigValidationMixin):
@@ -173,10 +175,10 @@ def parse_relation_results(cls, relation_results: RelationResults) -> Dict:
 
         Returns: a standard dictionary describing this `RedshiftMaterializedViewConfig` instance
         """
-        materialized_view: agate.Row = cls._get_first_row(
+        materialized_view: "agate.Row" = cls._get_first_row(
             relation_results.get("materialized_view")
         )
-        query: agate.Row = cls._get_first_row(relation_results.get("query"))
+        query: "agate.Row" = cls._get_first_row(relation_results.get("query"))
 
         config_dict = {
             "mv_name": materialized_view.get("table"),

dbt/adapters/redshift/relation_configs/sort.py

@@ -1,8 +1,7 @@
 from dataclasses import dataclass
 from dbt.adapters.contracts.relation import RelationConfig
-from typing import Optional, FrozenSet, Set, Dict, Any
+from typing import Optional, FrozenSet, Set, Dict, Any, TYPE_CHECKING
 
-import agate
 from dbt.adapters.relation_configs import (
     RelationConfigChange,
     RelationConfigChangeAction,
@@ -15,6 +14,9 @@
 
 from dbt.adapters.redshift.relation_configs.base import RedshiftRelationConfigBase
 
+if TYPE_CHECKING:
+    import agate
+
 
 class RedshiftSortStyle(StrEnum):
     auto = "auto"
@@ -136,7 +138,7 @@ def parse_relation_config(cls, relation_config: RelationConfig) -> Dict[str, Any
         return config_dict
 
     @classmethod
-    def parse_relation_results(cls, relation_results_entry: agate.Row) -> dict:
+    def parse_relation_results(cls, relation_results_entry: "agate.Row") -> dict:
         """
         Translate agate objects from the database into a standard dictionary.