diff --git a/src/houseplant/houseplant.py b/src/houseplant/houseplant.py index b41a3bf..d403f6b 100644 --- a/src/houseplant/houseplant.py +++ b/src/houseplant/houseplant.py @@ -11,6 +11,13 @@ from .utils import MIGRATIONS_DIR, get_migration_files +class EnvVars: + """Helper class to access environment variables via attribute syntax.""" + + def __getattr__(self, name: str) -> str: + return os.getenv(name, "") + + class Houseplant: def __init__(self): self.console = Console() @@ -119,7 +126,7 @@ def migrate_up(self, version: str | None = None): table_definition = migration.get("table_definition", "").strip() table_settings = migration.get("table_settings", "").strip() - format_args = {"table": table} + format_args = {"table": table, "env": EnvVars()} if table_definition and table_settings: format_args.update( { @@ -214,7 +221,9 @@ def migrate_down(self, version: str | None = None): # Get migration SQL based on environment migration_env = migration.get(self.env, {}) migration_sql = ( - migration_env.get("down", {}).format(table=table).strip() + migration_env.get("down", {}) + .format(table=table, env=EnvVars()) + .strip() ) if migration_sql: diff --git a/tests/test_houseplant.py b/tests/test_houseplant.py index deb0d7b..76839ae 100644 --- a/tests/test_houseplant.py +++ b/tests/test_houseplant.py @@ -629,3 +629,78 @@ def test_check_migrations_dir_not_found(houseplant, tmp_path): # Verify SystemExit is raised when migrations dir not found with pytest.raises(SystemExit): houseplant._check_migrations_dir() + + +@pytest.fixture +def test_migration_with_env_vars(tmp_path, monkeypatch): + # Set environment variables for the test + monkeypatch.setenv("TEST_CLUSTER", "my_cluster") + monkeypatch.setenv("TEST_ENGINE", "ReplicatedMergeTree") + + migrations_dir = tmp_path / "ch/migrations" + migrations_dir.mkdir(parents=True) + migration_file = migrations_dir / "20240101000000_test_env_vars.yml" + + migration_content = """version: "20240101000000" +name: test_env_vars +table: events + +development: &development + up: | + CREATE TABLE {table} ON CLUSTER '{env.TEST_CLUSTER}' ( + id UInt32, + name String + ) ENGINE = {env.TEST_ENGINE}() + ORDER BY id + down: | + DROP TABLE {table} ON CLUSTER '{env.TEST_CLUSTER}' + +test: + <<: *development + +production: + <<: *development +""" + + migration_file.write_text(migration_content) + os.chdir(tmp_path) + return migration_content + + +def test_migrate_up_with_env_vars(houseplant, test_migration_with_env_vars, mocker): + # Mock database calls + houseplant.env = "development" + mock_execute = mocker.patch.object(houseplant.db, "execute_migration") + mocker.patch.object(houseplant.db, "mark_migration_applied") + mocker.patch.object(houseplant.db, "get_applied_migrations", return_value=[]) + + # Run migration + houseplant.migrate_up() + + # Verify env vars were expanded in the SQL + expected_sql = """CREATE TABLE events ON CLUSTER 'my_cluster' ( + id UInt32, + name String +) ENGINE = ReplicatedMergeTree() +ORDER BY id""" + + mock_execute.assert_called_once_with(expected_sql, None) + + +def test_migrate_down_with_env_vars(houseplant, test_migration_with_env_vars, mocker): + # Mock database calls + houseplant.env = "development" + mock_execute = mocker.patch.object(houseplant.db, "execute_migration") + mocker.patch.object(houseplant.db, "mark_migration_rolled_back") + mocker.patch.object( + houseplant.db, "get_applied_migrations", return_value=[("20240101000000",)] + ) + mocker.patch.object(houseplant, "update_schema") + + # Run migration down + houseplant.migrate_down() + + # Verify env vars were expanded in the SQL + expected_sql = "DROP TABLE events ON CLUSTER 'my_cluster'" + + mock_execute.assert_called_once_with(expected_sql, None)