Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
36 changes: 34 additions & 2 deletions sqlmesh/core/context.py
Original file line number Diff line number Diff line change
Expand Up @@ -108,6 +108,7 @@
CachingStateSync,
StateReader,
StateSync,
Versions,
)
from sqlmesh.core.janitor import cleanup_expired_views, delete_expired_snapshots
from sqlmesh.core.table_diff import TableDiff
Expand Down Expand Up @@ -2610,15 +2611,18 @@ def migrate(self) -> None:
"""
self.notification_target_manager.notify(NotificationEvent.MIGRATION_START)
self._load_materializations()
state_sync = self._new_state_sync()
previous_versions = self._state_versions(state_sync)
try:
self._new_state_sync().migrate(
state_sync.migrate(
promoted_snapshots_only=self.config.migration.promoted_snapshots_only,
)
except Exception as e:
self.notification_target_manager.notify(
NotificationEvent.MIGRATION_FAILURE, traceback.format_exc()
)
raise e
self._print_state_versions(self._state_versions(state_sync), previous_versions)
self.notification_target_manager.notify(NotificationEvent.MIGRATION_END)

@python_api_analytics
Expand All @@ -2627,7 +2631,10 @@ def rollback(self) -> None:

Please contact your SQLMesh administrator before doing this. This action cannot be undone.
"""
self._new_state_sync().rollback()
state_sync = self._new_state_sync()
previous_versions = self._state_versions(state_sync)
state_sync.rollback()
self._print_state_versions(self._state_versions(state_sync), previous_versions)

@python_api_analytics
def create_external_models(self, strict: bool = False) -> None:
Expand Down Expand Up @@ -2701,6 +2708,12 @@ def print_info(
if state_connection:
self._try_connection("state backend", state_connection.connection_validator())

if verbosity >= Verbosity.VERBOSE:
try:
self._print_state_versions(self._state_versions())
except Exception as ex:
self.console.log_error(f"Failed to fetch the state backend versions. {ex}")

@python_api_analytics
def print_environment_names(self) -> None:
"""Prints all environment names along with expiry datetime."""
Expand Down Expand Up @@ -3290,6 +3303,25 @@ def _try_connection(self, connection_name: str, validator: t.Callable[[], None])
except Exception as ex:
self.console.log_error(f"{connection_name} connection failed. {ex}")

def _state_versions(self, state_sync: t.Optional[StateSync] = None) -> Versions:
"""Returns the versions recorded in the state backend without validating them."""
return (state_sync or self._new_state_sync()).get_versions(validate=False)

def _print_state_versions(
self, versions: Versions, previous_versions: t.Optional[Versions] = None
) -> None:
"""Prints the state backend versions, optionally alongside the ones they replaced."""
self.console.log_status_update("\nState backend versions:")
for label, attribute in (
("Schema version", "schema_version"),
("SQLGlot version", "sqlglot_version"),
("SQLMesh version", "sqlmesh_version"),
):
version = getattr(versions, attribute)
if previous_versions is not None:
version = f"{getattr(previous_versions, attribute)} -> {version}"
self.console.log_status_update(f"{label}: {version}")

def _new_state_sync(self) -> StateSync:
return self._provided_state_sync or self._scheduler.create_state_sync(self)

Expand Down
9 changes: 9 additions & 0 deletions sqlmesh/core/state_sync/db/migrator.py
Original file line number Diff line number Diff line change
Expand Up @@ -96,6 +96,15 @@ def migrate(
migrate_rows = self._apply_migrations(schema, skip_backup)

if not migrate_rows and major_minor(SQLMESH_VERSION) == versions.minor_sqlmesh_version:
# Nothing to migrate, but a patch-level bump still leaves the recorded versions
# behind what is actually running, so they are brought up to date here. The
# schema version is carried over rather than defaulted, since no migration ran
# and moving it could hide one that is genuinely needed later.
if (
versions.sqlmesh_version != SQLMESH_VERSION
or versions.sqlglot_version != SQLGLOT_VERSION
):
self.version_state.update_versions(schema_version=versions.schema_version)
return

if migrate_rows:
Expand Down
105 changes: 105 additions & 0 deletions tests/cli/test_cli.py
Original file line number Diff line number Diff line change
Expand Up @@ -10,11 +10,15 @@

from click import ClickException
from click.testing import CliRunner
from sqlglot import __version__ as SQLGLOT_VERSION
from sqlmesh import RuntimeEnv
from sqlmesh._version import __version__ as SQLMESH_VERSION
from sqlmesh.cli.project_init import ProjectTemplate, init_example_project
from sqlmesh.cli.main import cli
from sqlmesh.core.context import Context
from sqlmesh.core.state_sync.base import SCHEMA_VERSION
from sqlmesh.integrations.dlt import generate_dlt_models
from sqlmesh.utils import major_minor
from sqlmesh.utils.date import now_ds, time_like_to_str, timedelta, to_datetime, yesterday_ds
from sqlmesh.core.config.connection import DIALECT_TO_TYPE

Expand Down Expand Up @@ -1024,6 +1028,107 @@ def test_info_on_new_project_does_not_create_state_sync(runner, tmp_path):
assert not context.engine_adapter.table_exists("sqlmesh._versions")


def test_info_state_versions(runner, tmp_path):
create_example_project(tmp_path)
init_prod_and_backfill(runner, tmp_path)

result = runner.invoke(cli, ["--log-file-dir", tmp_path, "--paths", tmp_path, "info"])
assert result.exit_code == 0
assert "State backend versions" not in result.output

result = runner.invoke(cli, ["--log-file-dir", tmp_path, "--paths", tmp_path, "info", "-v"])
assert result.exit_code == 0
assert "State backend versions" in result.output
assert f"Schema version: {SCHEMA_VERSION}" in result.output
assert f"SQLGlot version: {SQLGLOT_VERSION}" in result.output
assert f"SQLMesh version: {SQLMESH_VERSION}" in result.output


def test_rollback_state_versions(runner, tmp_path):
create_example_project(tmp_path)
init_prod_and_backfill(runner, tmp_path)

context = Context(paths=tmp_path)
state_sync = context._new_state_sync()
# Back up the current state, then pretend the state was migrated by a newer SQLMesh.
state_sync.migrator._backup_state()
state_sync.version_state.update_versions(
schema_version=SCHEMA_VERSION + 1,
sqlglot_version="9999.0.0",
sqlmesh_version="9999.0.0",
)
context.close()

result = runner.invoke(cli, ["--log-file-dir", tmp_path, "--paths", tmp_path, "rollback"])
assert result.exit_code == 0
assert "State backend versions" in result.output
assert f"Schema version: {SCHEMA_VERSION + 1} -> {SCHEMA_VERSION}" in result.output
assert f"SQLGlot version: 9999.0.0 -> {SQLGLOT_VERSION}" in result.output
assert f"SQLMesh version: 9999.0.0 -> {SQLMESH_VERSION}" in result.output


def test_migrate_state_versions(runner, tmp_path):
create_example_project(tmp_path)
init_prod_and_backfill(runner, tmp_path)

context = Context(paths=tmp_path)
# Pretend the state was written by an older patch release of the same minor version, which
# is the case `migrate` used to leave untouched.
major, minor = major_minor(SQLMESH_VERSION)
older_sqlmesh = f"{major}.{minor}.dev0"
context._new_state_sync().version_state.update_versions(
sqlglot_version="0.0.1",
sqlmesh_version=older_sqlmesh,
)
context.close()

result = runner.invoke(cli, ["--log-file-dir", tmp_path, "--paths", tmp_path, "migrate"])
assert result.exit_code == 0
assert "State backend versions" in result.output
assert f"SQLGlot version: 0.0.1 -> {SQLGLOT_VERSION}" in result.output
assert f"SQLMesh version: {older_sqlmesh} -> {SQLMESH_VERSION}" in result.output


def test_migrate_updates_versions_after_a_patch_bump(runner, tmp_path):
"""A patch bump leaves the minor version equal, but the recorded versions must still move.

Both minor versions have to match the installed ones, otherwise `_apply_migrations` reports
rows to migrate and the early return this covers is never reached.
"""
create_example_project(tmp_path)
init_prod_and_backfill(runner, tmp_path)

sqlmesh_major, sqlmesh_minor = major_minor(SQLMESH_VERSION)
sqlglot_major, sqlglot_minor = major_minor(SQLGLOT_VERSION)
context = Context(paths=tmp_path)
context._new_state_sync().version_state.update_versions(
sqlglot_version=f"{sqlglot_major}.{sqlglot_minor}.dev0",
sqlmesh_version=f"{sqlmesh_major}.{sqlmesh_minor}.dev0",
)
context.close()

assert (
runner.invoke(cli, ["--log-file-dir", tmp_path, "--paths", tmp_path, "migrate"]).exit_code
== 0
)

context = Context(paths=tmp_path)
versions = context._new_state_sync().get_versions(validate=False)
context.close()
assert versions.sqlmesh_version == SQLMESH_VERSION
assert versions.sqlglot_version == SQLGLOT_VERSION


def test_rollback_without_backup_does_not_print_state_versions(runner, tmp_path):
create_example_project(tmp_path)
init_prod_and_backfill(runner, tmp_path)

result = runner.invoke(cli, ["--log-file-dir", tmp_path, "--paths", tmp_path, "rollback"])
assert result.exit_code == 1
assert "There are no prior migrations to roll back to." in result.output
assert "State backend versions" not in result.output


def test_dlt_pipeline_errors(runner, tmp_path):
# Error if no pipeline is provided
result = runner.invoke(cli, ["--paths", tmp_path, "init", "-t", "dlt", "duckdb"])
Expand Down
32 changes: 32 additions & 0 deletions tests/core/state_sync/test_state_sync.py
Original file line number Diff line number Diff line change
Expand Up @@ -4254,3 +4254,35 @@ def test_state_version_is_too_old(
match="The current state belongs to an old version of SQLMesh that is no longer supported. Please upgrade to 0.134.0 first before upgrading to.*",
):
state_sync.migrate(skip_backup=True)


def test_migrate_patch_bump_preserves_schema_version(
state_sync: EngineAdapterStateSync, mocker: MockerFixture
) -> None:
"""A run with nothing to migrate must not move the recorded schema version.

`_apply_migrations` is forced to report no rows so the patch-bump branch is reached with a
schema version that differs from the current one. A real state can't get into that shape,
which is the point: the schema version is carried over rather than defaulted so that a
migration which is genuinely still needed can't be masked.
"""
from sqlmesh import __version__ as SQLMESH_VERSION
from sqlmesh.core.state_sync.base import SCHEMA_VERSION

stale_schema_version = SCHEMA_VERSION - 1
state_sync.version_state.update_versions(
schema_version=stale_schema_version,
sqlglot_version="0.0.1",
sqlmesh_version=SQLMESH_VERSION,
)

mocker.patch(
"sqlmesh.core.state_sync.db.migrator.StateMigrator._apply_migrations",
return_value=False,
)

state_sync.migrate()

versions = state_sync.get_versions(validate=False)
assert versions.schema_version == stale_schema_version
assert versions.sqlglot_version == SQLGLOT_VERSION
51 changes: 37 additions & 14 deletions tests/integrations/jupyter/test_magics.py
Original file line number Diff line number Diff line change
Expand Up @@ -12,8 +12,11 @@
from IPython.utils.capture import CapturedIO, capture_output
from pytest_mock.plugin import MockerFixture
from rich.console import Console as RichConsole
from sqlglot import __version__ as SQLGLOT_VERSION

from sqlmesh import Context, RuntimeEnv
from sqlmesh._version import __version__ as SQLMESH_VERSION
from sqlmesh.core.state_sync.base import SCHEMA_VERSION, Versions
from sqlmesh.magics import register_magics
from pathlib import Path

Expand Down Expand Up @@ -740,14 +743,19 @@ def test_info(notebook, sushi_context, convert_all_html_output_to_text, get_all_

assert not output.stdout
assert not output.stderr
assert len(output.outputs) == 6
assert len(output.outputs) == 10
# No plan has been applied, so the state backend is still empty and reports the defaults.
assert convert_all_html_output_to_text(output) == [
"Models: 20",
"Macros: 8",
"",
"Connection:\n type: duckdb\n concurrent_tasks: 1\n register_comments: true\n pre_ping: false\n pretty_sql: false\n extensions: []\n connector_config: {}\n secrets: None\n filesystems: []\n shared_connection: true",
"Test Connection:\n type: duckdb\n concurrent_tasks: 1\n register_comments: true\n pre_ping: false\n pretty_sql: false\n extensions: []\n connector_config: {}\n secrets: None\n filesystems: []\n shared_connection: true",
"Data warehouse connection succeeded",
"State backend versions:",
"Schema version: 0",
"SQLGlot version: 0.0.0",
"SQLMesh version: 0.0.0",
]
assert get_all_html_output(output) == [
"<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\">Models: <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">20</span></pre>",
Expand All @@ -756,6 +764,10 @@ def test_info(notebook, sushi_context, convert_all_html_output_to_text, get_all_
'<pre style="white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,\'DejaVu Sans Mono\',consolas,\'Courier New\',monospace">Connection: type: duckdb concurrent_tasks: <span style="color: #008080; text-decoration-color: #008080; font-weight: bold">1</span> register_comments: true pre_ping: false pretty_sql: false extensions: <span style="font-weight: bold">[]</span> connector_config: <span style="font-weight: bold">{}</span> secrets: <span style="color: #800080; text-decoration-color: #800080; font-style: italic">None</span> filesystems: <span style="font-weight: bold">[]</span> shared_connection: true</pre>',
'<pre style="white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,\'DejaVu Sans Mono\',consolas,\'Courier New\',monospace">Test Connection: type: duckdb concurrent_tasks: <span style="color: #008080; text-decoration-color: #008080; font-weight: bold">1</span> register_comments: true pre_ping: false pretty_sql: false extensions: <span style="font-weight: bold">[]</span> connector_config: <span style="font-weight: bold">{}</span> secrets: <span style="color: #800080; text-decoration-color: #800080; font-style: italic">None</span> filesystems: <span style="font-weight: bold">[]</span> shared_connection: true</pre>',
"<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\">Data warehouse connection <span style=\"color: #008000; text-decoration-color: #008000\">succeeded</span></pre>",
"<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\">State backend versions:</pre>",
"<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\">Schema version: <span style=\"color: #008080; text-decoration-color: #008080; font-weight: bold\">0</span></pre>",
'<pre style="white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,\'DejaVu Sans Mono\',consolas,\'Courier New\',monospace">SQLGlot version: <span style="color: #008080; text-decoration-color: #008080; font-weight: bold">0.0</span>.<span style="color: #008080; text-decoration-color: #008080; font-weight: bold">0</span></pre>',
'<pre style="white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,\'DejaVu Sans Mono\',consolas,\'Courier New\',monospace">SQLMesh version: <span style="color: #008080; text-decoration-color: #008080; font-weight: bold">0.0</span>.<span style="color: #008080; text-decoration-color: #008080; font-weight: bold">0</span></pre>',
]


Expand All @@ -768,25 +780,36 @@ def test_migrate(

assert not output.stdout
assert not output.stderr
assert len(output.outputs) == 1
assert len(output.outputs) == 5
# The sushi state lives in an in-memory DuckDB database, so the state sync that `migrate`
# opens starts empty and the versions move from the defaults to the running ones.
empty = Versions()
assert convert_all_html_output_to_text(output) == [
"State backend versions:",
f"Schema version: {empty.schema_version} -> {SCHEMA_VERSION}",
f"SQLGlot version: {empty.sqlglot_version} -> {SQLGLOT_VERSION}",
f"SQLMesh version: {empty.sqlmesh_version} -> {SQLMESH_VERSION}",
"Migration complete",
]
assert get_all_html_output(output) == [
str(
# Rich highlights the numbers inside the version lines, and that markup depends on the
# running versions, so only the fixed lines are compared as HTML.
html_output = get_all_html_output(output)
assert html_output[0] == str(
h("pre", {"style": RICH_PRE_STYLE}, "State backend versions:", autoescape=False)
)
assert html_output[-1] == str(
h(
"pre",
{"style": RICH_PRE_STYLE},
h(
"pre",
{"style": RICH_PRE_STYLE},
h(
"span",
{"style": SUCCESS_STYLE},
"Migration complete",
autoescape=False,
),
"span",
{"style": SUCCESS_STYLE},
"Migration complete",
autoescape=False,
)
),
autoescape=False,
)
]
)


# TODO: Add test for rollback
Expand Down
Loading