feat(config): add is_patroni_postgresql flag to database configurations

VERSION

@@ -1 +1 @@
-6.0.0
+6.1.0

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: ""
+    additional_parameters: ""
+    is_patroni_postgresql: true

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
+    allow_migration_for_empty_database: false
+    is_patroni_postgresql: true

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
 

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:

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
         )

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)

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})