From 0c3f514978814ec302d7f31cf3b7bf8e3fe0a98c Mon Sep 17 00:00:00 2001 From: Mila Page <67295367+VersusFacit@users.noreply.github.com> Date: Thu, 13 Jun 2024 00:45:43 +0000 Subject: [PATCH] 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 Co-authored-by: Mila Page --- .../unreleased/Under the Hood-20240331-103115.yaml | 6 ++++++ dbt/adapters/redshift/connections.py | 13 +++++++++---- dbt/adapters/redshift/impl.py | 9 ++++++--- dbt/adapters/redshift/relation_configs/base.py | 11 ++++++++--- dbt/adapters/redshift/relation_configs/dist.py | 8 +++++--- .../redshift/relation_configs/materialized_view.py | 10 ++++++---- dbt/adapters/redshift/relation_configs/sort.py | 8 +++++--- 7 files changed, 45 insertions(+), 20 deletions(-) create mode 100644 .changes/unreleased/Under the Hood-20240331-103115.yaml diff --git a/.changes/unreleased/Under the Hood-20240331-103115.yaml b/.changes/unreleased/Under the Hood-20240331-103115.yaml new file mode 100644 index 000000000..ffb5a6489 --- /dev/null +++ b/.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" diff --git a/dbt/adapters/redshift/connections.py b/dbt/adapters/redshift/connections.py index 752c81e32..dfd7ab30b 100644 --- a/dbt/adapters/redshift/connections.py +++ b/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 diff --git a/dbt/adapters/redshift/impl.py b/dbt/adapters/redshift/impl.py index 18faee48c..d498685ed 100644 --- a/dbt/adapters/redshift/impl.py +++ b/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 diff --git a/dbt/adapters/redshift/relation_configs/base.py b/dbt/adapters/redshift/relation_configs/base.py index c4faab664..6f1409659 100644 --- a/dbt/adapters/redshift/relation_configs/base.py +++ b/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()) diff --git a/dbt/adapters/redshift/relation_configs/dist.py b/dbt/adapters/redshift/relation_configs/dist.py index 0104d8db4..2bcdb9566 100644 --- a/dbt/adapters/redshift/relation_configs/dist.py +++ b/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. diff --git a/dbt/adapters/redshift/relation_configs/materialized_view.py b/dbt/adapters/redshift/relation_configs/materialized_view.py index f6d93754e..48c04b554 100644 --- a/dbt/adapters/redshift/relation_configs/materialized_view.py +++ b/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"), diff --git a/dbt/adapters/redshift/relation_configs/sort.py b/dbt/adapters/redshift/relation_configs/sort.py index 91152615e..be5b3627d 100644 --- a/dbt/adapters/redshift/relation_configs/sort.py +++ b/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.