Add engine support flags

Author: treysp

sqlmesh/core/engine_adapter/athena.py

@@ -45,6 +45,7 @@ class AthenaEngineAdapter(PandasNativeFetchDFSupportMixin, RowDiffMixin):
     # >>> self._execute('/* test */ DESCRIBE foo')
     #     pyathena.error.OperationalError: FAILED: ParseException line 1:0 cannot recognize input near '/' '*' 'test'
     ATTACH_CORRELATION_ID = False
+    SUPPORTS_QUERY_EXECUTION_TRACKING = True
 
     def __init__(
         self, *args: t.Any, s3_warehouse_location: t.Optional[str] = None, **kwargs: t.Any

sqlmesh/core/engine_adapter/base.py

@@ -112,8 +112,7 @@ class EngineAdapter:
     QUOTE_IDENTIFIERS_IN_VIEWS = True
     MAX_IDENTIFIER_LENGTH: t.Optional[int] = None
     ATTACH_CORRELATION_ID = True
-    # TODO: change to False
-    SUPPORTS_QUERY_EXECUTION_TRACKING = True
+    SUPPORTS_QUERY_EXECUTION_TRACKING = False
 
     def __init__(
         self,

sqlmesh/core/engine_adapter/base_postgres.py

@@ -24,6 +24,7 @@ class BasePostgresEngineAdapter(EngineAdapter):
     DEFAULT_BATCH_SIZE = 400
     COMMENT_CREATION_TABLE = CommentCreationTable.COMMENT_COMMAND_ONLY
     COMMENT_CREATION_VIEW = CommentCreationView.COMMENT_COMMAND_ONLY
+    SUPPORTS_QUERY_EXECUTION_TRACKING = True
 
     def columns(
         self, table_name: TableName, include_pseudo_columns: bool = False

sqlmesh/core/engine_adapter/bigquery.py

@@ -69,6 +69,7 @@ class BigQueryEngineAdapter(InsertOverwriteWithMergeMixin, ClusteredByMixin, Row
     SUPPORTS_CLONING = True
     MAX_TABLE_COMMENT_LENGTH = 1024
     MAX_COLUMN_COMMENT_LENGTH = 1024
+    SUPPORTS_QUERY_EXECUTION_TRACKING = True
 
     SCHEMA_DIFFER = SchemaDiffer(
         compatible_types={
@@ -1084,7 +1085,13 @@ def _execute(
         self.cursor._set_description(query_results.schema)
 
         if track_row_count:
-            track_execution_record(sql, query_results.total_rows)
+            if query_job.statement_type == "CREATE_TABLE_AS_SELECT":
+                query_table = self.client.get_table(query_job.destination)
+                num_rows = query_table.num_rows
+            elif query_job.statement_type in ["INSERT", "DELETE", "MERGE", "UPDATE"]:
+                num_rows = query_job.num_dml_affected_rows
+
+            track_execution_record(sql, num_rows)
 
     def _get_data_objects(
         self, schema_name: SchemaName, object_names: t.Optional[t.Set[str]] = None

sqlmesh/core/engine_adapter/postgres.py

@@ -36,7 +36,6 @@ class PostgresEngineAdapter(
     CURRENT_CATALOG_EXPRESSION = exp.column("current_catalog")
     SUPPORTS_REPLACE_TABLE = False
     MAX_IDENTIFIER_LENGTH = 63
-    SUPPORTS_QUERY_EXECUTION_TRACKING = True
     SCHEMA_DIFFER = SchemaDiffer(
         parameterized_type_defaults={
             # DECIMAL without precision is "up to 131072 digits before the decimal point; up to 16383 digits after the decimal point"