PR Feedback 1

Author: VaggelisD

sqlmesh/core/context.py

@@ -116,7 +116,7 @@
     run_tests,
 )
 from sqlmesh.core.user import User
-from sqlmesh.utils import UniqueKeyDict, Verbosity
+from sqlmesh.utils import UniqueKeyDict, Verbosity, CorrelationId
 from sqlmesh.utils.concurrency import concurrent_apply_to_values
 from sqlmesh.utils.dag import DAG
 from sqlmesh.utils.date import (
@@ -418,7 +418,7 @@ def __init__(
             self.config.get_state_connection(self.gateway) or self.connection_config
         )
 
-        self._snapshot_evaluator: t.Optional[SnapshotEvaluator] = None
+        self._snapshot_evaluators: t.Dict[t.Optional[CorrelationId], SnapshotEvaluator] = {}
 
         self.console = get_console()
         setattr(self.console, "dialect", self.config.dialect)
@@ -446,20 +446,18 @@ def engine_adapter(self) -> EngineAdapter:
             self._engine_adapter = self.connection_config.create_engine_adapter()
         return self._engine_adapter
 
-    def snapshot_evaluator(self, job_id: t.Optional[str] = None) -> SnapshotEvaluator:
+    def snapshot_evaluator(self, job_id: t.Optional[CorrelationId] = None) -> SnapshotEvaluator:
         # Cache snapshot evaluators by job_id to avoid old job_ids being attached to future Context operations
-        if not self._snapshot_evaluator or any(
-            adapter._job_id != job_id for adapter in self._snapshot_evaluator.adapters.values()
-        ):
-            self._snapshot_evaluator = SnapshotEvaluator(
+        if job_id not in self._snapshot_evaluators:
+            self._snapshot_evaluators[job_id] = SnapshotEvaluator(
                 {
                     gateway: adapter.with_settings(level=logging.INFO, job_id=job_id)
                     for gateway, adapter in self.engine_adapters.items()
                 },
                 ddl_concurrent_tasks=self.concurrent_tasks,
                 selected_gateway=self.selected_gateway,
             )
-        return self._snapshot_evaluator
+        return self._snapshot_evaluators[job_id]
 
     def execution_context(
         self,
@@ -541,7 +539,7 @@ def scheduler(self, environment: t.Optional[str] = None) -> Scheduler:
         return self.create_scheduler(snapshots)
 
     def create_scheduler(
-        self, snapshots: t.Iterable[Snapshot], job_id: t.Optional[str] = None
+        self, snapshots: t.Iterable[Snapshot], job_id: t.Optional[CorrelationId] = None
     ) -> Scheduler:
         """Creates the built-in scheduler.
 
@@ -1594,7 +1592,7 @@ def apply(
             )
             explainer.evaluate(
                 plan.to_evaluatable(),
-                snapshot_evaluator=self.snapshot_evaluator(job_id=plan.plan_id),
+                snapshot_evaluator=self.snapshot_evaluator(job_id=CorrelationId.from_plan(plan)),
             )
             return
 
@@ -2341,8 +2339,8 @@ def print_environment_names(self) -> None:
 
     def close(self) -> None:
         """Releases all resources allocated by this context."""
-        if self._snapshot_evaluator:
-            self._snapshot_evaluator.close()
+        for evaluator in self._snapshot_evaluators.values():
+            evaluator.close()
         if self._state_sync:
             self._state_sync.close()
 
@@ -2397,7 +2395,7 @@ def _run(
     def _apply(self, plan: Plan, circuit_breaker: t.Optional[t.Callable[[], bool]]) -> None:
         self._scheduler.create_plan_evaluator(self).evaluate(
             plan.to_evaluatable(),
-            snapshot_evaluator=self.snapshot_evaluator(job_id=plan.plan_id),
+            snapshot_evaluator=self.snapshot_evaluator(job_id=CorrelationId.from_plan(plan)),
             circuit_breaker=circuit_breaker,
         )
 

sqlmesh/core/engine_adapter/base.py

@@ -39,7 +39,7 @@
 )
 from sqlmesh.core.model.kind import TimeColumn
 from sqlmesh.core.schema_diff import SchemaDiffer
-from sqlmesh.utils import columns_to_types_all_known, random_id
+from sqlmesh.utils import columns_to_types_all_known, random_id, CorrelationId
 from sqlmesh.utils.connection_pool import create_connection_pool, ConnectionPool
 from sqlmesh.utils.date import TimeLike, make_inclusive, to_time_column
 from sqlmesh.utils.errors import (
@@ -123,7 +123,7 @@ def __init__(
         pre_ping: bool = False,
         pretty_sql: bool = False,
         shared_connection: bool = False,
-        job_id: t.Optional[str] = None,
+        job_id: t.Optional[CorrelationId] = None,
         **kwargs: t.Any,
     ):
         self.dialect = dialect.lower() or self.DIALECT
@@ -2179,7 +2179,7 @@ def execute(
                     sql = t.cast(str, e)
 
                 if self._job_id:
-                    sql = f"/* sqlmesh_ref: {self._job_id} */ {sql}"
+                    sql = f"/* {self._job_id} */ {sql}"
 
                 self._log_sql(
                     sql,

sqlmesh/utils/__init__.py

@@ -13,6 +13,7 @@
 import types
 import typing as t
 import uuid
+from dataclasses import dataclass
 from collections import defaultdict
 from contextlib import contextmanager
 from copy import deepcopy
@@ -25,6 +26,9 @@
 
 logger = logging.getLogger(__name__)
 
+if t.TYPE_CHECKING:
+    from sqlmesh.core.plan import Plan
+
 T = t.TypeVar("T")
 KEY = t.TypeVar("KEY", bound=t.Hashable)
 VALUE = t.TypeVar("VALUE")
@@ -382,3 +386,23 @@ def to_snake_case(name: str) -> str:
     return "".join(
         f"_{c.lower()}" if c.isupper() and idx != 0 else c.lower() for idx, c in enumerate(name)
     )
+
+
+class JobType(Enum):
+    PLAN = "PLAN"
+    RUN = "RUN"
+
+
+@dataclass(frozen=True)
+class CorrelationId:
+    """ID that is added to each query in order to identify the job that created it."""
+
+    job_type: JobType
+    job_id: str
+
+    def __str__(self) -> str:
+        return f"{self.job_type.value}: {self.job_id}"
+
+    @classmethod
+    def from_plan(cls, plan: Plan) -> CorrelationId:
+        return CorrelationId(JobType.PLAN, plan.plan_id)

tests/conftest.py

@@ -42,7 +42,7 @@
     SnapshotDataVersion,
     SnapshotFingerprint,
 )
-from sqlmesh.utils import random_id
+from sqlmesh.utils import random_id, CorrelationId
 from sqlmesh.utils.date import TimeLike, to_date
 from sqlmesh.utils.windows import IS_WINDOWS, fix_windows_path
 from sqlmesh.core.engine_adapter.shared import CatalogSupport
@@ -269,7 +269,7 @@ def push_plan(context: Context, plan: Plan) -> None:
         context.create_scheduler,
         context.default_catalog,
     )
-    plan_evaluator.snapshot_evaluator = context.snapshot_evaluator(job_id=plan.plan_id)
+    plan_evaluator.snapshot_evaluator = context.snapshot_evaluator(CorrelationId.from_plan(plan))
     deployability_index = DeployabilityIndex.create(context.snapshots.values())
     evaluatable_plan = plan.to_evaluatable()
     stages = plan_stages.build_plan_stages(

tests/core/test_integration.py

@@ -67,6 +67,7 @@
     SnapshotInfoLike,
     SnapshotTableInfo,
 )
+from sqlmesh.utils import CorrelationId
 from sqlmesh.utils.date import TimeLike, now, to_date, to_datetime, to_timestamp
 from sqlmesh.utils.errors import NoChangesPlanError, SQLMeshError, PlanError, ConfigError
 from sqlmesh.utils.pydantic import validate_string
@@ -6469,23 +6470,27 @@ def plan_with_output(ctx: Context, environment: str):
         assert context_diff.environment == environment
 
 
-def test_plan_evaluator_job_id(tmp_path: Path, mocker: MockerFixture):
-    def _to_sqls(mock_logger):
-        return [call[0][0] for call in mock_logger.call_args_list]
+def test_plan_evaluator_correlation_id(tmp_path: Path):
+    def _correlation_id_in_sqls(correlation_id: CorrelationId, mock_logger):
+        sqls = [call[0][0] for call in mock_logger.call_args_list]
+        return any(f"/* {correlation_id} */" in sql for sql in sqls)
 
     create_temp_file(
         tmp_path, Path("models") / "test.sql", "MODEL (name test.a, kind FULL); SELECT 1 AS col"
     )
 
-    # Case 1: Ensure that the job id (plan_id) is included in the SQL
+    # Case 1: Ensure that the correlation id (plan_id) is included in the SQL
     with mock.patch("sqlmesh.core.engine_adapter.base.EngineAdapter._log_sql") as mock_logger:
         ctx = Context(paths=[tmp_path], config=Config())
         plan = ctx.plan(auto_apply=True, no_prompts=True)
 
-    assert any(f"/* sqlmesh_ref: {plan.plan_id} */" in sql for sql in _to_sqls(mock_logger))
+    correlation_id = CorrelationId.from_plan(plan)
+    assert str(correlation_id) == f"PLAN: {plan.plan_id}"
 
-    # Case 2: Ensure that the previous job id is not included in the SQL for other operations
+    assert _correlation_id_in_sqls(correlation_id, mock_logger)
+
+    # Case 2: Ensure that the previous correlation id is not included in the SQL for other operations
     with mock.patch("sqlmesh.core.engine_adapter.base.EngineAdapter._log_sql") as mock_logger:
         ctx.snapshot_evaluator().adapter.execute("SELECT 1")
 
-    assert not any(f"/* sqlmesh_ref: {plan.plan_id} */" in sql for sql in _to_sqls(mock_logger))
+    assert not _correlation_id_in_sqls(correlation_id, mock_logger)

tests/core/test_plan_evaluator.py

@@ -11,6 +11,7 @@
     stages as plan_stages,
 )
 from sqlmesh.core.snapshot import SnapshotChangeCategory
+from sqlmesh.utils import CorrelationId
 
 
 @pytest.fixture
@@ -63,7 +64,7 @@ def test_builtin_evaluator_push(sushi_context: Context, make_snapshot):
         sushi_context.default_catalog,
         console=sushi_context.console,
     )
-    evaluator.snapshot_evaluator = sushi_context.snapshot_evaluator(job_id=plan.plan_id)
+    evaluator.snapshot_evaluator = sushi_context.snapshot_evaluator(CorrelationId.from_plan(plan))
 
     evaluatable_plan = plan.to_evaluatable()
     stages = plan_stages.build_plan_stages(