Skip to content

Commit 40bb6c8

Browse files
authored
Fix: Do not attach correlation_id for Athena queries (#4935)
1 parent d1dc55c commit 40bb6c8

File tree

2 files changed

+11
-2
lines changed

2 files changed

+11
-2
lines changed

sqlmesh/core/engine_adapter/athena.py

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -41,6 +41,10 @@ class AthenaEngineAdapter(PandasNativeFetchDFSupportMixin, RowDiffMixin):
4141
COMMENT_CREATION_VIEW = CommentCreationView.UNSUPPORTED
4242
SCHEMA_DIFFER = TrinoEngineAdapter.SCHEMA_DIFFER
4343
MAX_TIMESTAMP_PRECISION = 3 # copied from Trino
44+
# Athena does not deal with comments well, e.g:
45+
# >>> self._execute('/* test */ DESCRIBE foo')
46+
# pyathena.error.OperationalError: FAILED: ParseException line 1:0 cannot recognize input near '/' '*' 'test'
47+
ATTACH_CORRELATION_ID = False
4448

4549
def __init__(
4650
self, *args: t.Any, s3_warehouse_location: t.Optional[str] = None, **kwargs: t.Any

sqlmesh/core/engine_adapter/base.py

Lines changed: 7 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -109,6 +109,7 @@ class EngineAdapter:
109109
DEFAULT_CATALOG_TYPE = DIALECT
110110
QUOTE_IDENTIFIERS_IN_VIEWS = True
111111
MAX_IDENTIFIER_LENGTH: t.Optional[int] = None
112+
ATTACH_CORRELATION_ID = True
112113

113114
def __init__(
114115
self,
@@ -2219,8 +2220,7 @@ def execute(
22192220
else:
22202221
sql = t.cast(str, e)
22212222

2222-
if self.correlation_id:
2223-
sql = f"/* {self.correlation_id} */ {sql}"
2223+
sql = self._attach_correlation_id(sql)
22242224

22252225
self._log_sql(
22262226
sql,
@@ -2229,6 +2229,11 @@ def execute(
22292229
)
22302230
self._execute(sql, **kwargs)
22312231

2232+
def _attach_correlation_id(self, sql: str) -> str:
2233+
if self.ATTACH_CORRELATION_ID and self.correlation_id:
2234+
return f"/* {self.correlation_id} */ {sql}"
2235+
return sql
2236+
22322237
def _log_sql(
22332238
self,
22342239
sql: str,

0 commit comments

Comments
 (0)