Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

lazy load agate #954

Closed
wants to merge 2 commits into from
Closed
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
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"
12 changes: 7 additions & 5 deletions dbt/adapters/snowflake/connections.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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")

Expand Down Expand Up @@ -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:
Expand Down
12 changes: 7 additions & 5 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 @@ -183,7 +185,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
10 changes: 7 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,9 @@
SnowflakeQuotePolicy,
)

if TYPE_CHECKING:
import agate


@dataclass(frozen=True, eq=True, unsafe_hash=True)
class SnowflakeRelationConfigBase(RelationConfigBase):
Expand Down Expand Up @@ -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())
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
Loading