Skip to content

Commit

Permalink
Merge community commits into local branch.
Browse files Browse the repository at this point in the history
  • Loading branch information
VersusFacit committed May 26, 2024
2 parents 4ff6bdb + 67c4bd4 commit fca76d7
Show file tree
Hide file tree
Showing 5 changed files with 34 additions and 16 deletions.
6 changes: 6 additions & 0 deletions .changes/unreleased/Under the Hood-20240327-001304.yaml
Original file line number Diff line number Diff line change
@@ -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"
11 changes: 7 additions & 4 deletions dbt/adapters/snowflake/connections.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -46,6 +44,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")

Expand Down Expand Up @@ -468,9 +469,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:
Expand Down
14 changes: 8 additions & 6 deletions dbt/adapters/snowflake/impl.py
Original file line number Diff line number Diff line change
@@ -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
Expand All @@ -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):
Expand Down Expand Up @@ -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])
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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:
Expand Down
11 changes: 8 additions & 3 deletions dbt/adapters/snowflake/relation_configs/base.py
Original file line number Diff line number Diff line change
@@ -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,
Expand All @@ -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):
Expand Down Expand Up @@ -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())
8 changes: 5 additions & 3 deletions dbt/adapters/snowflake/relation_configs/dynamic_table.py
Original file line number Diff line number Diff line change
@@ -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):
Expand Down Expand Up @@ -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"),
Expand Down

0 comments on commit fca76d7

Please sign in to comment.