Add flag for supporting row tracking to engine adapters

Author: treysp

.circleci/continue_config.yml

@@ -305,10 +305,11 @@ workflows:
                 - clickhouse-cloud
                 - athena
                 - gcp-postgres
-          filters:
-            branches:
-              only:
-                - main
+          # TODO: uncomment this
+          # filters:
+          #   branches:
+          #     only:
+          #       - main
       - ui_style
       - ui_test
       - vscode_test

sqlmesh/core/console.py

@@ -3995,6 +3995,7 @@ def show_table_diff_summary(self, table_diff: TableDiff) -> None:
         self._write(f"Join On: {keys}")
 
 
+# TODO: remove this
 # _CONSOLE: Console = NoopConsole()
 _CONSOLE: Console = TerminalConsole()
 

sqlmesh/core/engine_adapter/base.py

@@ -112,6 +112,8 @@ 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
 
     def __init__(
         self,
@@ -2282,7 +2284,7 @@ def _log_sql(
     def _execute(self, sql: str, track_row_count: bool = False, **kwargs: t.Any) -> None:
         self.cursor.execute(sql, **kwargs)
 
-        if track_row_count:
+        if track_row_count and self.SUPPORTS_QUERY_EXECUTION_TRACKING:
             rowcount_raw = getattr(self.cursor, "rowcount", None)
             rowcount = None
             if rowcount_raw is not None:

sqlmesh/core/engine_adapter/mssql.py

@@ -53,6 +53,7 @@ class MSSQLEngineAdapter(
     COMMENT_CREATION_TABLE = CommentCreationTable.UNSUPPORTED
     COMMENT_CREATION_VIEW = CommentCreationView.UNSUPPORTED
     SUPPORTS_REPLACE_TABLE = False
+    SUPPORTS_QUERY_EXECUTION_TRACKING = True
     SCHEMA_DIFFER = SchemaDiffer(
         parameterized_type_defaults={
             exp.DataType.build("DECIMAL", dialect=DIALECT).this: [(18, 0), (0,)],

sqlmesh/core/engine_adapter/mysql.py

@@ -40,6 +40,7 @@ class MySQLEngineAdapter(
     MAX_COLUMN_COMMENT_LENGTH = 1024
     SUPPORTS_REPLACE_TABLE = False
     MAX_IDENTIFIER_LENGTH = 64
+    SUPPORTS_QUERY_EXECUTION_TRACKING = True
     SCHEMA_DIFFER = SchemaDiffer(
         parameterized_type_defaults={
             exp.DataType.build("BIT", dialect=DIALECT).this: [(1,)],

sqlmesh/core/engine_adapter/postgres.py

@@ -36,6 +36,7 @@ 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"

sqlmesh/core/engine_adapter/trino.py

@@ -54,6 +54,7 @@ class TrinoEngineAdapter(
     SUPPORTS_REPLACE_TABLE = False
     DEFAULT_CATALOG_TYPE = "hive"
     QUOTE_IDENTIFIERS_IN_VIEWS = False
+    SUPPORTS_QUERY_EXECUTION_TRACKING = True
     SCHEMA_DIFFER = SchemaDiffer(
         parameterized_type_defaults={
             # default decimal precision varies across backends

sqlmesh/core/execution_tracker.py

@@ -31,6 +31,7 @@ def add_execution(self, sql: str, row_count: t.Optional[int]) -> None:
         if row_count is not None and row_count >= 0:
             self.total_rows_processed += row_count
         self.query_count += 1
+        # TODO: remove this
         # for debugging
         self.queries_executed.append((sql[:300], row_count, time.time()))
 

tests/core/engine_adapter/integration/test_integration.py

@@ -1657,10 +1657,11 @@ def capture_row_counts(
     assert len(physical_layer_results.materialized_views) == 0
     assert len(physical_layer_results.tables) == len(physical_layer_results.non_temp_tables) == 3
 
-    assert len(row_counts) == 3
-    assert row_counts["seed_model"] == 7
-    assert row_counts["incremental_model"] == 7
-    assert row_counts["full_model"] == 3
+    if ctx.engine_adapter.SUPPORTS_QUERY_EXECUTION_TRACKING:
+        assert len(row_counts) == 3
+        assert row_counts["seed_model"] == 7
+        assert row_counts["incremental_model"] == 7
+        assert row_counts["full_model"] == 3
 
     # make and validate unmodified dev environment
     no_change_plan: Plan = context.plan_builder(