diff --git a/VERSION b/VERSION index 09b254e..dfda3e0 100644 --- a/VERSION +++ b/VERSION @@ -1 +1 @@ -6.0.0 +6.1.0 diff --git a/example/multi-database-config.yaml b/example/multi-database-config.yaml index 73bd247..2c03ba1 100644 --- a/example/multi-database-config.yaml +++ b/example/multi-database-config.yaml @@ -17,6 +17,8 @@ databases: database: app_main allow_migration_for_empty_database: true additional_parameters: "" + is_patroni_postgresql: true + # Analytics database analytics: @@ -30,6 +32,8 @@ databases: database: analytics allow_migration_for_empty_database: false additional_parameters: "" + is_patroni_postgresql: true + # Audit database with custom parameters audit: @@ -43,6 +47,8 @@ databases: database: audit_logs allow_migration_for_empty_database: true additional_parameters: "-x audit_mode=yes" + is_patroni_postgresql: true + # Example with environment variable substitution (if supported by your deployment) reports: @@ -55,4 +61,5 @@ databases: port: 5432 database: reports allow_migration_for_empty_database: false - additional_parameters: "" \ No newline at end of file + additional_parameters: "" + is_patroni_postgresql: true diff --git a/example/simple-config.yaml b/example/simple-config.yaml index 9d59389..bf516d6 100644 --- a/example/simple-config.yaml +++ b/example/simple-config.yaml @@ -11,6 +11,7 @@ databases: port: 5432 database: myapp allow_migration_for_empty_database: true + is_patroni_postgresql: false # Cache database cache: @@ -22,4 +23,5 @@ databases: host: redis-db port: 5432 database: cache_db - allow_migration_for_empty_database: false \ No newline at end of file + allow_migration_for_empty_database: false + is_patroni_postgresql: true \ No newline at end of file diff --git a/src/chartreuse/chartreuse.py b/src/chartreuse/chartreuse.py index 7582986..7e588d7 100644 --- a/src/chartreuse/chartreuse.py +++ b/src/chartreuse/chartreuse.py @@ -49,6 +49,7 @@ def __init__( allow_migration_for_empty_database=db_config.allow_migration_for_empty_database, additional_parameters=additional_params, alembic_section_name=db_name, + is_patroni_postgresql=db_config.is_patroni_postgresql, ) self.migration_helpers[db_name] = helper diff --git a/src/chartreuse/config_loader.py b/src/chartreuse/config_loader.py index 85fc50d..2ec571f 100644 --- a/src/chartreuse/config_loader.py +++ b/src/chartreuse/config_loader.py @@ -28,6 +28,7 @@ class DatabaseConfig(BaseModel): # Optional migration settings allow_migration_for_empty_database: bool = Field(default=True, description="Allow migrations on empty database") additional_parameters: str = Field(default="", description="Additional Alembic parameters") + is_patroni_postgresql: bool = Field(default=False, description="Whether this is a Patroni PostgreSQL database") @computed_field @property @@ -39,7 +40,7 @@ def url(self) -> str: @classmethod def validate_dialect(cls, v: str) -> str: """Validate database dialect.""" - supported_dialects = ["postgresql", "mysql", "sqlite", "oracle", "mssql", "clickhouse"] + supported_dialects = ["postgresql", "clickhouse"] if v.lower() not in supported_dialects: logger.warning("Dialect '%s' might not be supported. Supported dialects: %s", v, supported_dialects) return v @@ -96,6 +97,7 @@ def load_multi_database_config(config_path: str) -> dict[str, DatabaseConfig]: alembic_config_file_path: alembic.ini allow_migration_for_empty_database: true additional_parameters: "" + is_patroni_postgresql: false # For single database setups, just include one database: # analytics: @@ -109,6 +111,7 @@ def load_multi_database_config(config_path: str) -> dict[str, DatabaseConfig]: # alembic_config_file_path: alembic.ini # allow_migration_for_empty_database: false # additional_parameters: "" + # is_patroni_postgresql: true ``` """ try: diff --git a/src/chartreuse/tests/unit_tests/test_chartreuse.py b/src/chartreuse/tests/unit_tests/test_chartreuse.py index d30e074..d69b687 100644 --- a/src/chartreuse/tests/unit_tests/test_chartreuse.py +++ b/src/chartreuse/tests/unit_tests/test_chartreuse.py @@ -72,6 +72,7 @@ def test_chartreuse_init_with_kubernetes_helper(self, mocker: MockerFixture) -> allow_migration_for_empty_database=True, additional_parameters="--verbose -n test-db", alembic_section_name="test-db", + is_patroni_postgresql=False, ) # Verify kubernetes helper is set @@ -516,4 +517,5 @@ def test_multi_chartreuse_default_values(self, mocker: MockerFixture) -> None: allow_migration_for_empty_database=True, # Default value from DatabaseConfig additional_parameters="-n test", # Section name parameter added (no leading space when original is empty) alembic_section_name="test", # New parameter for multi-database support + is_patroni_postgresql=False, # Default value from DatabaseConfig ) diff --git a/src/chartreuse/tests/unit_tests/test_patroni_postgresql.py b/src/chartreuse/tests/unit_tests/test_patroni_postgresql.py new file mode 100644 index 0000000..3fc666d --- /dev/null +++ b/src/chartreuse/tests/unit_tests/test_patroni_postgresql.py @@ -0,0 +1,116 @@ +from pytest_mock import MockerFixture + +from chartreuse.utils.alembic_migration_helper import AlembicMigrationHelper + + +class TestPatroniPostgreSQLFeature: + """Test cases for the is_patroni_postgresql feature.""" + + def test_patroni_postgresql_flag_false_by_default(self, mocker: MockerFixture) -> None: + """Test that is_patroni_postgresql defaults to False.""" + mocker.patch("chartreuse.utils.AlembicMigrationHelper._configure") + mocker.patch("chartreuse.utils.AlembicMigrationHelper._check_migration_needed", return_value=False) + + helper = AlembicMigrationHelper( + database_url="postgresql://user:pass@localhost:5432/db", + alembic_section_name="test", + skip_db_checks=True, # Skip database checks for testing + ) + + assert helper.is_patroni_postgresql is False + + def test_patroni_postgresql_flag_set_to_true(self, mocker: MockerFixture) -> None: + """Test that is_patroni_postgresql can be set to True.""" + mocker.patch("chartreuse.utils.AlembicMigrationHelper._configure") + mocker.patch("chartreuse.utils.AlembicMigrationHelper._check_migration_needed", return_value=False) + + helper = AlembicMigrationHelper( + database_url="postgresql://user:pass@localhost:5432/db", + alembic_section_name="test", + is_patroni_postgresql=True, + skip_db_checks=True, # Skip database checks for testing + ) + + assert helper.is_patroni_postgresql is True + + def test_patroni_postgresql_adds_parameter(self, mocker: MockerFixture) -> None: + """Test that when is_patroni_postgresql=True, it adds the correct parameter.""" + mocker.patch("chartreuse.utils.AlembicMigrationHelper._configure") + mocker.patch("chartreuse.utils.AlembicMigrationHelper._wait_postgres_is_configured") + mocker.patch("chartreuse.utils.AlembicMigrationHelper._check_migration_needed", return_value=False) + + helper = AlembicMigrationHelper( + database_url="postgresql://user:pass@localhost:5432/db", + alembic_section_name="test", + is_patroni_postgresql=True, + additional_parameters="--verbose", + ) + + # Should have added the patroni_postgresql parameter + assert " -x patroni_postgresql=yes" in helper.additional_parameters + assert "--verbose" in helper.additional_parameters + + def test_patroni_postgresql_does_not_add_parameter_when_false(self, mocker: MockerFixture) -> None: + """Test that when is_patroni_postgresql=False, it doesn't add the parameter.""" + mocker.patch("chartreuse.utils.AlembicMigrationHelper._configure") + mocker.patch("chartreuse.utils.AlembicMigrationHelper._check_migration_needed", return_value=False) + + helper = AlembicMigrationHelper( + database_url="postgresql://user:pass@localhost:5432/db", + alembic_section_name="test", + is_patroni_postgresql=False, + additional_parameters="--verbose", + ) + + # Should not have added the patroni_postgresql parameter + assert "patroni_postgresql=yes" not in helper.additional_parameters + assert "--verbose" in helper.additional_parameters + + def test_patroni_postgresql_calls_wait_postgres_configured(self, mocker: MockerFixture) -> None: + """Test that when is_patroni_postgresql=True, it calls _wait_postgres_is_configured.""" + mocker.patch("chartreuse.utils.AlembicMigrationHelper._configure") + mock_wait = mocker.patch("chartreuse.utils.AlembicMigrationHelper._wait_postgres_is_configured") + mocker.patch("chartreuse.utils.AlembicMigrationHelper._check_migration_needed", return_value=False) + + AlembicMigrationHelper( + database_url="postgresql://user:pass@localhost:5432/db", + alembic_section_name="test", + is_patroni_postgresql=True, + ) + + mock_wait.assert_called_once() + + def test_patroni_postgresql_does_not_call_wait_postgres_configured_when_false(self, mocker: MockerFixture) -> None: + """Test that when is_patroni_postgresql=False, it doesn't call _wait_postgres_is_configured.""" + mocker.patch("chartreuse.utils.AlembicMigrationHelper._configure") + mock_wait = mocker.patch("chartreuse.utils.AlembicMigrationHelper._wait_postgres_is_configured") + mocker.patch("chartreuse.utils.AlembicMigrationHelper._check_migration_needed", return_value=False) + + AlembicMigrationHelper( + database_url="postgresql://user:pass@localhost:5432/db", + alembic_section_name="test", + is_patroni_postgresql=False, + ) + + mock_wait.assert_not_called() + + def test_patroni_postgresql_skips_wait_when_skip_db_checks_true(self, mocker: MockerFixture) -> None: + """Test that when skip_db_checks=True, it doesn't call _wait_postgres_is_configured.""" + mocker.patch("chartreuse.utils.AlembicMigrationHelper._configure") + mock_wait = mocker.patch("chartreuse.utils.AlembicMigrationHelper._wait_postgres_is_configured") + mocker.patch("chartreuse.utils.AlembicMigrationHelper._check_migration_needed", return_value=False) + + AlembicMigrationHelper( + database_url="postgresql://user:pass@localhost:5432/db", + alembic_section_name="test", + is_patroni_postgresql=True, + skip_db_checks=True, + ) + + mock_wait.assert_not_called() + + def test_wait_postgres_is_configured_method_exists(self) -> None: + """Test that the _wait_postgres_is_configured method exists and is callable.""" + # This is a basic sanity check that the method exists + assert hasattr(AlembicMigrationHelper, "_wait_postgres_is_configured") + assert callable(AlembicMigrationHelper._wait_postgres_is_configured) diff --git a/src/chartreuse/utils/alembic_migration_helper.py b/src/chartreuse/utils/alembic_migration_helper.py index 27c8c31..1294617 100755 --- a/src/chartreuse/utils/alembic_migration_helper.py +++ b/src/chartreuse/utils/alembic_migration_helper.py @@ -26,6 +26,7 @@ def __init__( configure: bool = True, # skip_db_checks is used for testing purposes only skip_db_checks: bool = False, + is_patroni_postgresql: bool = False, ): if not database_url: raise ValueError("database_url not set, not upgrading database.") @@ -37,9 +38,9 @@ def __init__( self.alembic_config_file_path = alembic_config_file_path self.alembic_section_name = alembic_section_name self.skip_db_checks = skip_db_checks + self.is_patroni_postgresql = is_patroni_postgresql # Chartreuse will upgrade a PG managed/configured by postgres-operator - self.is_patroni_postgresql: bool = "CHARTREUSE_PATRONI_POSTGRESQL" in os.environ if self.is_patroni_postgresql and not skip_db_checks: self.additional_parameters += " -x patroni_postgresql=yes" self._wait_postgres_is_configured() @@ -77,6 +78,7 @@ def _wait_postgres_is_configured(self) -> None: Make sure the user `wiremind_owner_user` was created by the postgres-operator and that default privileges were configured. # TODO: Maybe make this a readinessProbe on Patroni PG Pods + # We need to remove this, it's an open source tool... """ wait_timeout = int(os.getenv("CHARTREUSE_ALEMBIC_POSTGRES_WAIT_CONFIGURED_TIMEOUT", 60)) engine = sqlalchemy.create_engine(self.database_url, poolclass=NullPool, connect_args={"connect_timeout": 1})