feat: add python library versions to the environment (#3275)

Author: tobymao

sqlmesh/core/environment.py

@@ -101,6 +101,7 @@ class Environment(EnvironmentNamingInfo):
         promoted_snapshot_ids: The IDs of the snapshots that are promoted in this environment
             (i.e. for which the views are created). If not specified, all snapshots are promoted.
         previous_finalized_snapshots: Snapshots that were part of this environment last time it was finalized.
+        requirements: A mapping of library versions for all the snapshots in this environment.
     """
 
     snapshots_: t.List[t.Any] = Field(alias="snapshots")
@@ -116,6 +117,7 @@ class Environment(EnvironmentNamingInfo):
     previous_finalized_snapshots_: t.Optional[t.List[t.Any]] = Field(
         default=None, alias="previous_finalized_snapshots"
     )
+    requirements: t.Dict[str, str] = {}
 
     @field_validator("snapshots_", "previous_finalized_snapshots_", mode="before")
     @classmethod
@@ -135,6 +137,12 @@ def _load_snapshot_ids(cls, v: str | t.List[t.Any] | None) -> t.List[t.Any] | No
             raise ValueError("Must be a list of SnapshotId dicts or objects")
         return v
 
+    @field_validator("requirements", mode="before")
+    def _load_requirements(cls, v: t.Any) -> t.Any:
+        if isinstance(v, str):
+            v = json.loads(v)
+        return v or {}
+
     @property
     def snapshots(self) -> t.List[SnapshotTableInfo]:
         return self._convert_list_to_models_and_store("snapshots_", SnapshotTableInfo)

sqlmesh/core/plan/definition.py

@@ -1,5 +1,7 @@
 from __future__ import annotations
 
+import logging
+import sys
 import typing as t
 from dataclasses import dataclass
 from datetime import datetime
@@ -26,9 +28,20 @@
 from sqlmesh.utils.date import TimeLike, now, to_datetime, to_timestamp
 from sqlmesh.utils.pydantic import PydanticModel
 
+logger = logging.getLogger(__name__)
+
 SnapshotMapping = t.Dict[SnapshotId, t.Set[SnapshotId]]
 
 
+if sys.version_info >= (3, 12):
+    from importlib import metadata
+else:
+    import importlib_metadata as metadata  # type: ignore
+
+
+IGNORED_PACKAGES = {"sqlmesh", "sqlglot"}
+
+
 class Plan(PydanticModel, frozen=True):
     context_diff: ContextDiff
     plan_id: str
@@ -209,6 +222,23 @@ def environment(self) -> Environment:
             else self.context_diff.previous_finalized_snapshots
         )
 
+        requirements = {}
+        distributions = metadata.packages_distributions()
+
+        for snapshot in self.context_diff.snapshots.values():
+            if snapshot.is_model:
+                for executable in snapshot.model.python_env.values():
+                    if executable.kind == "import":
+                        try:
+                            start = "from " if executable.payload.startswith("from ") else "import "
+                            lib = executable.payload.split(start)[1].split()[0].split(".")[0]
+                            if lib in distributions:
+                                for dist in distributions[lib]:
+                                    if dist not in requirements and dist not in IGNORED_PACKAGES:
+                                        requirements[dist] = metadata.version(dist)
+                        except metadata.PackageNotFoundError:
+                            logger.warning("Failed to find package for %s", lib)
+
         return Environment(
             snapshots=snapshots,
             start_at=self.provided_start or self._earliest_interval_start,
@@ -218,6 +248,7 @@ def environment(self) -> Environment:
             expiration_ts=expiration_ts,
             promoted_snapshot_ids=promoted_snapshot_ids,
             previous_finalized_snapshots=previous_finalized_snapshots,
+            requirements=requirements,
             **self.environment_naming_info.dict(),
         )
 

sqlmesh/core/state_sync/engine_adapter.py

@@ -152,6 +152,7 @@ def __init__(
             "catalog_name_override": exp.DataType.build("text"),
             "previous_finalized_snapshots": exp.DataType.build("text"),
             "normalize_name": exp.DataType.build("boolean"),
+            "requirements": exp.DataType.build("text"),
         }
 
         self._interval_columns_to_types = {
@@ -1711,6 +1712,7 @@ def _environment_to_df(environment: Environment) -> pd.DataFrame:
                     else None
                 ),
                 "normalize_name": environment.normalize_name,
+                "requirements": json.dumps(environment.requirements),
             }
         ]
     )

sqlmesh/migrations/v0058_add_requirements.py

@@ -0,0 +1,23 @@
+"""Add requirements to environments table"""
+
+from sqlglot import exp
+
+
+def migrate(state_sync, **kwargs):  # type: ignore
+    engine_adapter = state_sync.engine_adapter
+    environments_table = "_environments"
+    if state_sync.schema:
+        environments_table = f"{state_sync.schema}.{environments_table}"
+
+    alter_table_exp = exp.Alter(
+        this=exp.to_table(environments_table),
+        kind="TABLE",
+        actions=[
+            exp.ColumnDef(
+                this=exp.to_column("requirements"),
+                kind=exp.DataType.build("text"),
+            )
+        ],
+    )
+
+    engine_adapter.execute(alter_table_exp)

tests/core/test_plan.py

@@ -41,6 +41,7 @@
     yesterday_ds,
 )
 from sqlmesh.utils.errors import PlanError
+from sqlmesh.utils.metaprogramming import Executable
 
 
 def test_forward_only_plan_sets_version(make_snapshot, mocker: MockerFixture):
@@ -2562,3 +2563,17 @@ def test_interval_end_per_model(make_snapshot):
         is_dev=True,
     )
     assert plan_builder.build().interval_end_per_model is None
+
+
+def test_plan_requirements():
+    context = Context(paths="examples/sushi")
+    model = context.get_model("sushi.items")
+    model.python_env["ruamel"] = Executable(payload="import ruamel", kind="import")
+    model.python_env["Image"] = Executable(
+        payload="from ipywidgets.widgets.widget_media import Image", kind="import"
+    )
+
+    plan = context.plan(
+        "dev", no_prompts=True, skip_tests=True, skip_backfill=True
+    ).environment.requirements
+    assert set(plan) == {"ipywidgets", "numpy", "pandas", "ruamel.yaml", "ruamel.yaml.clib"}

tests/schedulers/airflow/test_client.py

@@ -172,6 +172,7 @@ def test_apply_plan(mocker: MockerFixture, snapshot: Snapshot):
                 ],
                 "suffix_target": "schema",
                 "normalize_name": True,
+                "requirements": {},
             },
             "no_gaps": False,
             "skip_backfill": False,