diff --git a/sqlmesh/core/engine_adapter/athena.py b/sqlmesh/core/engine_adapter/athena.py index 88ab9b2c5d..abaf7ba281 100644 --- a/sqlmesh/core/engine_adapter/athena.py +++ b/sqlmesh/core/engine_adapter/athena.py @@ -41,6 +41,10 @@ class AthenaEngineAdapter(PandasNativeFetchDFSupportMixin, RowDiffMixin): COMMENT_CREATION_VIEW = CommentCreationView.UNSUPPORTED SCHEMA_DIFFER = TrinoEngineAdapter.SCHEMA_DIFFER MAX_TIMESTAMP_PRECISION = 3 # copied from Trino + # Athena does not deal with comments well, e.g: + # >>> self._execute('/* test */ DESCRIBE foo') + # pyathena.error.OperationalError: FAILED: ParseException line 1:0 cannot recognize input near '/' '*' 'test' + ATTACH_CORRELATION_ID = False def __init__( self, *args: t.Any, s3_warehouse_location: t.Optional[str] = None, **kwargs: t.Any diff --git a/sqlmesh/core/engine_adapter/base.py b/sqlmesh/core/engine_adapter/base.py index 1d34ff1401..33ad4c398a 100644 --- a/sqlmesh/core/engine_adapter/base.py +++ b/sqlmesh/core/engine_adapter/base.py @@ -109,6 +109,7 @@ class EngineAdapter: DEFAULT_CATALOG_TYPE = DIALECT QUOTE_IDENTIFIERS_IN_VIEWS = True MAX_IDENTIFIER_LENGTH: t.Optional[int] = None + ATTACH_CORRELATION_ID = True def __init__( self, @@ -2219,8 +2220,7 @@ def execute( else: sql = t.cast(str, e) - if self.correlation_id: - sql = f"/* {self.correlation_id} */ {sql}" + sql = self._attach_correlation_id(sql) self._log_sql( sql, @@ -2229,6 +2229,11 @@ def execute( ) self._execute(sql, **kwargs) + def _attach_correlation_id(self, sql: str) -> str: + if self.ATTACH_CORRELATION_ID and self.correlation_id: + return f"/* {self.correlation_id} */ {sql}" + return sql + def _log_sql( self, sql: str,