diff --git a/.gitignore b/.gitignore index 20cb51d0..25ccbdc8 100644 --- a/.gitignore +++ b/.gitignore @@ -10,3 +10,4 @@ docker/dbt/.user.yml .DS_Store .vscode/ logs/ +venv diff --git a/dbt/adapters/trino/connections.py b/dbt/adapters/trino/connections.py index d4d46934..4350c5e2 100644 --- a/dbt/adapters/trino/connections.py +++ b/dbt/adapters/trino/connections.py @@ -1,6 +1,7 @@ import decimal import os import re +import time from abc import ABCMeta, abstractmethod from contextlib import contextmanager from dataclasses import dataclass, field @@ -505,11 +506,34 @@ def add_query(self, sql, auto_begin=True, bindings=None, abridge_sql_log=False): return connection, cursor - def execute(self, sql, auto_begin=False, fetch=False): - _, cursor = self.add_query(sql, auto_begin) - status = self.get_response(cursor) - table = self.get_result_from_cursor(cursor) - return status, table + def execute( + self, + sql: str, + auto_begin: bool = False, + fetch: bool = False, + max_retries=trino.constants.DEFAULT_MAX_ATTEMPTS, + retry_delay: int = 1, + ): + retry_count = 0 + while retry_count < max_retries: + try: + _, cursor = self.add_query(sql, auto_begin) + status = self.get_response(cursor) + table = self.get_result_from_cursor(cursor) + return status, table + + except DbtDatabaseError: + retry_count += 1 + if retry_count < max_retries: + time.sleep(retry_delay) + logger.info(f"Retrying query ({retry_count}/{max_retries})...") + else: + raise + + except Exception as exc: + raise exc + + raise DbtRuntimeError("Failed to execute the query after maximum retries.") @classmethod def data_type_code_to_name(cls, type_code) -> str: