From 17c2df0080fa005ea876f62815f07fd579f6d4d3 Mon Sep 17 00:00:00 2001 From: dwreeves Date: Wed, 27 Mar 2024 00:10:55 -0400 Subject: [PATCH] lazy load agate --- .../unreleased/Under the Hood-20240327-001304.yaml | 6 ++++++ dbt/adapters/snowflake/connections.py | 12 +++++++----- dbt/adapters/snowflake/impl.py | 12 +++++++----- dbt/adapters/snowflake/relation_configs/base.py | 10 +++++++--- .../snowflake/relation_configs/dynamic_table.py | 8 +++++--- 5 files changed, 32 insertions(+), 16 deletions(-) create mode 100644 .changes/unreleased/Under the Hood-20240327-001304.yaml diff --git a/.changes/unreleased/Under the Hood-20240327-001304.yaml b/.changes/unreleased/Under the Hood-20240327-001304.yaml new file mode 100644 index 000000000..3e823ec86 --- /dev/null +++ b/.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" diff --git a/dbt/adapters/snowflake/connections.py b/dbt/adapters/snowflake/connections.py index fcd16aafc..86f28c561 100644 --- a/dbt/adapters/snowflake/connections.py +++ b/dbt/adapters/snowflake/connections.py @@ -8,10 +8,7 @@ from dataclasses import dataclass from io import StringIO from time import sleep -from typing import Optional, Tuple, Union, Any, List - -import agate -from dbt_common.clients.agate_helper import empty_table +from typing import Optional, Tuple, Union, Any, List, TYPE_CHECKING from cryptography.hazmat.backends import default_backend from cryptography.hazmat.primitives import serialization @@ -46,6 +43,9 @@ from dbt.adapters.events.types import AdapterEventWarning from dbt_common.ui import line_wrap_message, warning_tag +if TYPE_CHECKING: + import agate + logger = AdapterLogger("Snowflake") @@ -465,9 +465,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: diff --git a/dbt/adapters/snowflake/impl.py b/dbt/adapters/snowflake/impl.py index 3583c888c..9c5539186 100644 --- a/dbt/adapters/snowflake/impl.py +++ b/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 # type: ignore 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): @@ -62,8 +64,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]) @@ -180,7 +182,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: diff --git a/dbt/adapters/snowflake/relation_configs/base.py b/dbt/adapters/snowflake/relation_configs/base.py index 7b4367e2d..4d5803a20 100644 --- a/dbt/adapters/snowflake/relation_configs/base.py +++ b/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,9 @@ SnowflakeQuotePolicy, ) +if TYPE_CHECKING: + import agate + @dataclass(frozen=True, eq=True, unsafe_hash=True) class SnowflakeRelationConfigBase(RelationConfigBase): @@ -62,8 +64,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/snowflake/relation_configs/dynamic_table.py b/dbt/adapters/snowflake/relation_configs/dynamic_table.py index cc1b9112d..d13c9353d 100644 --- a/dbt/adapters/snowflake/relation_configs/dynamic_table.py +++ b/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"),