diff --git a/.github/workflows/python-test.yml b/.github/workflows/python-test.yml index f824e03..d450fa9 100644 --- a/.github/workflows/python-test.yml +++ b/.github/workflows/python-test.yml @@ -34,17 +34,10 @@ jobs: python -m pip install -r requirements.txt python -m pip install -e .[dev] - - name: Lint with flake8 + - name: Lint and check with ruff run: | - flake8 . --count --show-source --statistics - - - name: Check typing with mypy - run: | - mypy --install-types --non-interactive . - - - name: Check syntax with pyupgrade - run: | - find . -type f -regex '.*\.py$' -exec pyupgrade --py311-plus {} \; + ruff check . --output-format=github + ruff format --check . - name: Set up Helm uses: azure/setup-helm@v1 @@ -76,4 +69,4 @@ jobs: with: github_token: ${{ secrets.github_token }} # Change reviewdog reporter if you need [github-pr-check, github-check]. - reporter: github-pr-check + reporter: github-pr-check \ No newline at end of file diff --git a/.gitignore b/.gitignore index c250529..b40fe5f 100644 --- a/.gitignore +++ b/.gitignore @@ -68,3 +68,7 @@ venv-python3.7 venv-python3.8 venv-python3.10 .python-version + +# Helm Chart +*.lock +example/helm-chart/my-example-chart/charts/ \ No newline at end of file diff --git a/CHANGELOG.md b/CHANGELOG.md index 5d44713..c69656f 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,5 +1,11 @@ # Changelog +## v6.0.0 (2025-10-07) + +### Feat +- add support for multi-database migration + + ## v5.0.0 (2025-05-05) ### BREAKING CHANGE - Change minimum supported version of python: 3.11 (drop 3.7, 3.8, 3.9, 3.10) diff --git a/VERSION b/VERSION index 0062ac9..09b254e 100644 --- a/VERSION +++ b/VERSION @@ -1 +1 @@ -5.0.0 +6.0.0 diff --git a/docs/multi-database.md b/docs/multi-database.md new file mode 100644 index 0000000..d715d1f --- /dev/null +++ b/docs/multi-database.md @@ -0,0 +1,166 @@ +# Multi-Database Migration Support + +Chartreuse now supports managing migrations for multiple databases simultaneously. This feature allows you to define multiple database configurations and run migrations across all of them. + +## Configuration + +### Single Database (Legacy) +The original single-database configuration using environment variables is still supported for backward compatibility. + +### Multiple Databases (New) +To use multiple databases, create a YAML configuration file and set the `CHARTREUSE_MULTI_CONFIG_PATH` environment variable to point to it. + +#### Configuration File Format + +```yaml +databases: + # Main application database + main: + alembic_directory_path: /app/alembic/main + alembic_config_file_path: alembic.ini + dialect: postgresql + user: app_user + password: app_password + host: postgres-main + port: 5432 + database: app_main + allow_migration_for_empty_database: true + additional_parameters: "" + + # Analytics database + analytics: + alembic_directory_path: /app/alembic/analytics + alembic_config_file_path: alembic.ini + dialect: postgresql + user: analytics_user + password: analytics_password + host: postgres-analytics + port: 5432 + database: analytics + allow_migration_for_empty_database: false + additional_parameters: "" +``` + +#### Configuration Fields + +**Required fields:** +- Database name (used as key in the YAML) +- `alembic_directory_path`: Path to the Alembic directory for this database +- `alembic_config_file_path`: Alembic configuration file name + +**Database connection components (all required):** +- `dialect`: Database dialect (e.g., postgresql, mysql, sqlite) +- `user`: Database username +- `password`: Database password +- `host`: Database host +- `port`: Database port +- `database`: Database name + +**Optional fields:** +- `allow_migration_for_empty_database`: Whether to allow migrations on empty databases (default: false) +- `additional_parameters`: Additional parameters to pass to Alembic commands (default: "") + +## Environment Variables + +For multi-database mode: +- `CHARTREUSE_MULTI_CONFIG_PATH`: Path to the YAML configuration file +- `CHARTREUSE_ENABLE_STOP_PODS`: Whether to stop pods during migration (optional, default: true) +- `CHARTREUSE_RELEASE_NAME`: Kubernetes release name +- `CHARTREUSE_UPGRADE_BEFORE_DEPLOYMENT`: Whether to upgrade before deployment (optional, default: false) +- `HELM_IS_INSTALL`: Whether this is a Helm install operation (optional, default: false) + +## Usage + +### Setting up the environment +```bash +export CHARTREUSE_MULTI_CONFIG_PATH="/path/to/multi-database-config.yaml" +export CHARTREUSE_RELEASE_NAME="my-app" +export CHARTREUSE_ENABLE_STOP_PODS="true" +``` + +### Directory Structure Example +``` +/app/ +├── alembic/ +│ ├── main/ +│ │ ├── alembic.ini +│ │ ├── env.py +│ │ └── versions/ +│ └── analytics/ +│ ├── alembic.ini +│ ├── env.py +│ └── versions/ +└── multi-database-config.yaml +``` + +### Configuration Validation +You can validate your configuration file before deployment: + +```bash +python3 scripts/validate_config.py /path/to/multi-database-config.yaml +``` + +## Migration Behavior + +When using multi-database configuration: + +1. **Initialization**: All databases are initialized with their respective Alembic configurations +2. **Migration Check**: Each database is checked individually for pending migrations +3. **Migration Execution**: Only databases that need migration will be upgraded +4. **Error Handling**: If any database migration fails, the entire process fails +5. **Logging**: Detailed logs show which databases are being processed + +## Backward Compatibility + +The original single-database configuration using environment variables continues to work unchanged. The multi-database feature is only activated when `CHARTREUSE_MULTI_CONFIG_PATH` is set. + +## Example Integration + +In your Helm chart, you can use a ConfigMap to store the multi-database configuration: + +```yaml +apiVersion: v1 +kind: ConfigMap +metadata: + name: chartreuse-config +data: + multi-database-config.yaml: | + databases: + main: + alembic_directory_path: /app/alembic/main + alembic_config_file_path: alembic.ini + dialect: postgresql + user: {{ .Values.database.main.user }} + password: {{ .Values.database.main.password }} + host: {{ .Values.database.main.host }} + port: {{ .Values.database.main.port }} + database: {{ .Values.database.main.name }} + allow_migration_for_empty_database: true + additional_parameters: "" + analytics: + alembic_directory_path: /app/alembic/analytics + alembic_config_file_path: alembic.ini + dialect: postgresql + user: {{ .Values.database.analytics.user }} + password: {{ .Values.database.analytics.password }} + host: {{ .Values.database.analytics.host }} + port: {{ .Values.database.analytics.port }} + database: {{ .Values.database.analytics.name }} + allow_migration_for_empty_database: false + additional_parameters: "" +``` + +Then mount it in your migration job: + +```yaml +env: +- name: CHARTREUSE_MULTI_CONFIG_PATH + value: "/config/multi-database-config.yaml" +volumeMounts: +- name: config + mountPath: /config +volumes: +- name: config + configMap: + name: chartreuse-config +``` \ No newline at end of file diff --git a/example/Dockerfile-dev b/example/Dockerfile-dev index efe32f5..588e753 100644 --- a/example/Dockerfile-dev +++ b/example/Dockerfile-dev @@ -1,12 +1,16 @@ # Only used for e2e tests. Do not use in production. - +# For Wiremind use only since it uses our private package wiremind-python FROM python:3.12 WORKDIR /app +# Install uv +RUN pip install uv + COPY requirements.txt requirements.txt -RUN pip install --no-cache-dir -r requirements.txt +RUN uv pip install --system --no-cache -r requirements.txt COPY . . -RUN pip install -e . -COPY example/alembic alembic +RUN uv pip install --system -e . +RUN uv pip install --system wiremind-python +COPY alembic alembic diff --git a/example/Dockerfile-dev-auth b/example/Dockerfile-dev-auth new file mode 100644 index 0000000..8c7cc61 --- /dev/null +++ b/example/Dockerfile-dev-auth @@ -0,0 +1,28 @@ +# Only used for e2e tests. Do not use in production. +# For Wiremind use only since it uses our private package wiremind-python + +FROM python:3.12 + +WORKDIR /app + +# Install uv +RUN pip install uv + +# Install dependencies if requirements.txt exists +COPY pyproject.toml . +COPY VERSION . +COPY README.md . +COPY requirements.txt* ./ +COPY src/ src/ +COPY example/alembic/ alembic/ + +RUN if [ -f requirements.txt ]; then uv pip install --system --no-cache -r requirements.txt; fi + +# Install the package itself +RUN uv pip install --system -e . + +# Install wiremind-python with authentication using build secrets +# The secret will be mounted at build time but not stored in the image +RUN --mount=type=secret,id=uv_config,target=/root/.config/uv/uv.toml \ + mkdir -p /root/.config/uv && \ + uv pip install --system wiremind-python cyrillic clickhouse-sqlalchemy diff --git a/example/alembic/alembic.ini b/example/alembic/alembic.ini index a649ff9..b1d468a 100644 --- a/example/alembic/alembic.ini +++ b/example/alembic/alembic.ini @@ -1,55 +1,17 @@ -# A generic, single database configuration. - -[alembic] -# path to migration scripts -script_location = . - +[postgresql] +script_location = postgresl +sqlalchemy.url = postgresql://wiremind_owner@localhost:5432/wiremind +prepend_sys_path = .. +file_template = %%(year)d%%(month).2d%%(day).2d-%%(slug)s + +[clickhouse] +script_location = clickhouse +sqlalchemy.url = clickhouse://default@localhost:8123/wiremind +prepend_sys_path = .. # template used to generate migration files -# file_template = %%(rev)s_%%(slug)s - -# timezone to use when rendering the date -# within the migration file as well as the filename. -# string value is passed to dateutil.tz.gettz() -# leave blank for localtime -# timezone = - -# max length of characters to apply to the -# "slug" field -# truncate_slug_length = 40 - -# set to 'true' to run the environment during -# the 'revision' command, regardless of autogenerate -# revision_environment = false - -# set to 'true' to allow .pyc and .pyo files without -# a source .py file to be detected as revisions in the -# versions/ directory -# sourceless = false - -# version location specification; this defaults -# to sample_alembic/versions. When using multiple version -# directories, initial revisions must be specified with --version-path -# version_locations = %(here)s/bar %(here)s/bat sample_alembic/versions - -# the output encoding used when revision files -# are written from script.py.mako -# output_encoding = utf-8 - -sqlalchemy.url=postgresql://foo:foo@localhost/foo?sslmode=prefer - - -[post_write_hooks] -# post_write_hooks defines scripts or Python functions that are run -# on newly generated revision scripts. See the documentation for further -# detail and examples - -# format using "black" - use the console_scripts runner, against the "black" entrypoint -# hooks=black -# black.type=console_scripts -# black.entrypoint=black -# black.options=-l 79 +file_template = %%(year)d%%(month).2d%%(day).2d-%%(slug)s -# Logging configuration +# Logging configuration - Required for fileConfig() in env.py [loggers] keys = root,sqlalchemy,alembic @@ -82,4 +44,4 @@ formatter = generic [formatter_generic] format = %(levelname)-5.5s [%(name)s] %(message)s -datefmt = %H:%M:%S +datefmt = %H:%M:%S \ No newline at end of file diff --git a/example/alembic/clickhouse/env.py b/example/alembic/clickhouse/env.py new file mode 100644 index 0000000..36d308a --- /dev/null +++ b/example/alembic/clickhouse/env.py @@ -0,0 +1,99 @@ +import logging + +from alembic import context +from sqlalchemy import engine_from_config, pool + +# inspired by https://github.com/xzkostyan/clickhouse-sqlalchemy-alembic-example/blob/main/simple/migrations/env.py + +config = context.config +logger = logging.getLogger(__name__) + +logging.basicConfig(level=logging.INFO) +for key in list(logging.root.manager.loggerDict.keys()): + if "alembic" in key: + logging.getLogger(key).disabled = False + +from cyrillic.models import ClickhouseBase # noqa: E402 + +target_metadata = ClickhouseBase.metadata + +# Currently there is an issue in clickhouse-sqlalchemy https://github.com/xzkostyan/clickhouse-sqlalchemy/pull/369 +# makes Alembic mimic what clickhouse_sqlalchemy.alembic is expecting (designed for Alembic 1.5.8) +# https://github.com/sqlalchemy/alembic/blob/rel_1_5_8/alembic/util/sqla_compat.py#L180 +from alembic.util import sqla_compat # noqa: E402 + +sqla_compat._reflect_table = lambda inspector, table, include_cols: inspector.reflect_table(table, None) +# patch before next import +from clickhouse_sqlalchemy.alembic.dialect import include_object, patch_alembic_version # noqa: E402 + +REPLICATION = True + + +def run_migrations_offline(): + """Run migrations in 'offline' mode. + + This configures the context with just a URL + and not an Engine, though an Engine is acceptable + here as well. By skipping the Engine creation + we don't even need a DBAPI to be available. + + Calls to context.execute() here emit the given string to the + script output. + + """ + url = config.get_main_option("sqlalchemy.url") + context.configure( + url=url, + target_metadata=target_metadata, + literal_binds=True, + dialect_opts={"paramstyle": "named"}, + ) + + if REPLICATION: + kwargs = { + "cluster": "wm", + "table_path": "/clickhouse/tables/wm/wiremind/alembic_version", + "replica_name": "{replica}", + } + else: + kwargs = {} + + with context.begin_transaction(): + patch_alembic_version(context, **kwargs) + context.run_migrations() + + +def run_migrations_online(): + """Run migrations in 'online' mode. + + In this scenario we need to create an Engine + and associate a connection with the context. + + """ + connectable = engine_from_config( + config.get_section(config.config_ini_section), + prefix="sqlalchemy.", + poolclass=pool.NullPool, + ) + + with connectable.connect() as connection: + context.configure(connection=connection, target_metadata=target_metadata, include_object=include_object) + + if REPLICATION: + kwargs = { + "cluster": "wm", + "table_path": "/clickhouse/tables/wm/wiremind/alembic_version", + "replica_name": "{replica}", + } + else: + kwargs = {} + + with context.begin_transaction(): + patch_alembic_version(context, **kwargs) + context.run_migrations() + + +if context.is_offline_mode(): + run_migrations_offline() +else: + run_migrations_online() diff --git a/example/alembic/clickhouse/script.py.mako b/example/alembic/clickhouse/script.py.mako new file mode 100644 index 0000000..9f9e8c2 --- /dev/null +++ b/example/alembic/clickhouse/script.py.mako @@ -0,0 +1,28 @@ +"""${message} + +Revision ID: ${up_revision} +Revises: ${down_revision | comma,n} +Create Date: ${create_date} + +Have a look to examples https://github.com/xzkostyan/clickhouse-sqlalchemy-alembic-example + +""" +from alembic import op +import clickhouse_sqlalchemy +from clickhouse_sqlalchemy import engines +import sqlalchemy as sa +${imports if imports else ""} + +# revision identifiers, used by Alembic. +revision = ${repr(up_revision)} +down_revision = ${repr(down_revision)} +branch_labels = ${repr(branch_labels)} +depends_on = ${repr(depends_on)} + + +def upgrade() -> None: + ${upgrades if upgrades else "pass"} + + +def downgrade() -> None: + ${downgrades if downgrades else "pass"} diff --git a/example/alembic/clickhouse/versions/20250606-inital.py b/example/alembic/clickhouse/versions/20250606-inital.py new file mode 100644 index 0000000..9173c8e --- /dev/null +++ b/example/alembic/clickhouse/versions/20250606-inital.py @@ -0,0 +1,343 @@ +"""inital + +Revision ID: 55622187799e +Revises: +Create Date: 2025-06-06 18:28:21.202381 + +Have a look to examples https://github.com/xzkostyan/clickhouse-sqlalchemy-alembic-example + +""" + +import sqlalchemy as sa +from alembic import op +from clickhouse_sqlalchemy import engines +from clickhouse_sqlalchemy.types import ( + Array, + Boolean, + Date, + DateTime64, + Decimal, + Int16, + Int32, + Int64, + LowCardinality, + Nullable, + String, +) + +# revision identifiers, used by Alembic. +revision = "55622187799e" +down_revision = None +branch_labels = None +depends_on = None + + +def upgrade() -> None: + op.create_table( + "fact_daily_leg_bucket", + sa.Column("time", DateTime64(3, timezone="UTC")), + sa.Column("day_x", Int16), + sa.Column("event_datetime", DateTime64(3, timezone="UTC")), + sa.Column("transporter_type", LowCardinality(String)), + sa.Column("entity_name", LowCardinality(String)), + sa.Column("route_name", LowCardinality(String)), + sa.Column("service_departure_date", Date), + sa.Column("service_number", String), + sa.Column("service_labels", Array(String)), + sa.Column("service_status", LowCardinality(String)), + sa.Column("service_origin_station_name", String), + sa.Column("service_destination_station_name", String), + sa.Column("service_departure_datetime", DateTime64(3, timezone="UTC")), + sa.Column("service_arrival_datetime", DateTime64(3, timezone="UTC")), + sa.Column("service_budget_objective", Nullable(Decimal(10, 2))), + sa.Column("service_yield_objective", Nullable(Decimal(10, 2))), + sa.Column("is_service_peak_leg", Boolean), + sa.Column("is_service_max_leg", Boolean), + sa.Column("leg_id", Int64), + sa.Column("leg_order", Int32), + sa.Column("leg_origin_station_name", String), + sa.Column("leg_destination_station_name", String), + sa.Column("leg_departure_datetime", DateTime64(3, timezone="UTC")), + sa.Column("leg_arrival_datetime", DateTime64(3, timezone="UTC")), + sa.Column("leg_list_alerts", Array(String)), + sa.Column("reference_service_number", Nullable(String)), + sa.Column("reference_service_departure_date", Nullable(Date)), + sa.Column("cabin_name", LowCardinality(String)), + sa.Column("family_name", LowCardinality(String)), + sa.Column("bucket_name", LowCardinality(String)), + sa.Column("cumulative_sum_net_revenue_vat_inc", Decimal(10, 2)), + sa.Column("cumulative_sum_net_revenue_vat_exc", Decimal(10, 2)), + sa.Column("cumulative_sum_net_ancillary_revenue_vat_inc", Decimal(10, 2)), + sa.Column("cumulative_sum_net_ancillary_revenue_vat_exc", Decimal(10, 2)), + sa.Column("sum_net_revenue_vat_inc", Decimal(10, 2)), + sa.Column("sum_net_revenue_vat_exc", Decimal(10, 2)), + sa.Column("sum_net_ancillary_revenue_vat_inc", Decimal(10, 2)), + sa.Column("sum_net_ancillary_revenue_vat_exc", Decimal(10, 2)), + sa.Column("availability_end_day", Nullable(Int32)), + sa.Column("cumul_availability_end_day", Int32), + sa.Column("cumulative_sum_net_bookings", Nullable(Int32)), + sa.Column("sum_net_bookings", Nullable(Int32)), + sa.Column("cabin_capacity", Nullable(Int32)), + sa.Column("cabin_lid", Nullable(Int32)), + sa.Column("bucket_authorization_end_day", Nullable(Int32)), + sa.Column("is_last", Boolean), + engines.ReplicatedMergeTree( + table_path="/clickhouse/tables/{shard}/fact_daily_leg_bucket/", + replica_name="{replica}", + primary_key=("time", "service_departure_date"), + order_by=( + "time", + "service_departure_date", + "day_x", + "service_number", + "leg_id", + "bucket_name", + ), + ), + sa.text("INDEX is_last_index is_last TYPE set GRANULARITY 1"), + clickhouse_cluster="wm", + ) + + op.create_table( + "fact_daily_leg_physical_inventory", + sa.Column("time", DateTime64(3, timezone="UTC")), + sa.Column("day_x", Int16), + sa.Column("event_datetime", DateTime64(3, timezone="UTC")), + sa.Column("transporter_type", LowCardinality(String)), + sa.Column("entity_name", LowCardinality(String)), + sa.Column("route_name", LowCardinality(String)), + sa.Column("service_number", String), + sa.Column("service_labels", Array(String)), + sa.Column("service_status", LowCardinality(String)), + sa.Column("service_departure_date", Date), + sa.Column("service_origin_station_name", String), + sa.Column("service_destination_station_name", String), + sa.Column("service_departure_datetime", DateTime64(3, timezone="UTC")), + sa.Column("service_arrival_datetime", DateTime64(3, timezone="UTC")), + sa.Column("is_service_peak_leg", Boolean), + sa.Column("is_service_max_leg", Boolean), + sa.Column("service_max_leg_origin_station_name", String), + sa.Column("service_max_leg_destination_station_name", String), + sa.Column("leg_id", Int64), + sa.Column("leg_status", String), + sa.Column("leg_origin_station_name", String), + sa.Column("leg_destination_station_name", String), + sa.Column("leg_departure_datetime", DateTime64(3, timezone="UTC")), + sa.Column("leg_arrival_datetime", DateTime64(3, timezone="UTC")), + sa.Column("leg_order", Int32), + sa.Column("leg_list_alerts", Array(String)), + sa.Column("cabin_name", LowCardinality(String)), + sa.Column("physical_inventory_name", LowCardinality(String)), + sa.Column("physical_inventory_capacity", Nullable(Int32)), + sa.Column("physical_inventory_lid", Nullable(Int32)), + sa.Column("physical_availability_start_day", Nullable(Int32)), + sa.Column("physical_availability_end_day", Nullable(Int32)), + sa.Column("cumulative_sum_net_revenue_vat_inc", Decimal(10, 2)), + sa.Column("cumulative_sum_net_revenue_vat_exc", Decimal(10, 2)), + sa.Column("cumulative_sum_net_ancillary_revenue_vat_inc", Decimal(10, 2)), + sa.Column("cumulative_sum_net_ancillary_revenue_vat_exc", Decimal(10, 2)), + sa.Column("sum_net_revenue_vat_inc", Decimal(10, 2)), + sa.Column("sum_net_revenue_vat_exc", Decimal(10, 2)), + sa.Column("sum_net_ancillary_revenue_vat_inc", Decimal(10, 2)), + sa.Column("sum_net_ancillary_revenue_vat_exc", Decimal(10, 2)), + sa.Column("cumulative_sum_net_bookings", Int32), + sa.Column("sum_net_bookings", Int32), + sa.Column("forecast_full_day_x", Nullable(Int32)), + sa.Column("optimization_full_day_x", Nullable(Int32)), + sa.Column("unconstrained_demand_bookings", Nullable(Decimal(10, 2))), + sa.Column("unconstrained_forecast_bookings", Nullable(Decimal(10, 2))), + sa.Column("is_last", Boolean), + engines.ReplicatedMergeTree( + table_path="/clickhouse/tables/{shard}/fact_daily_leg_physical_inventory/", + replica_name="{replica}", + primary_key=("time", "service_departure_date"), + order_by=( + "time", + "service_departure_date", + "day_x", + "service_number", + "leg_id", + "physical_inventory_name", + ), + ), + sa.text("INDEX is_last_index is_last TYPE set GRANULARITY 1"), + clickhouse_cluster="wm", + ) + + op.create_table( + "fact_daily_od_bucket", + sa.Column("time", DateTime64(3, timezone="UTC")), + sa.Column("day_x", Int32), + sa.Column("event_datetime", DateTime64(3, timezone="UTC")), + sa.Column("transporter_type", LowCardinality(String)), + sa.Column("entity_name", LowCardinality(String)), + sa.Column("route_name", LowCardinality(String)), + sa.Column("service_id", Int64), + sa.Column("service_departure_date", Date), + sa.Column("service_number", String), + sa.Column("service_labels", Array(String)), + sa.Column("service_status", String), + sa.Column("service_origin_station_name", String), + sa.Column("service_destination_station_name", String), + sa.Column("service_departure_datetime", DateTime64(3, timezone="UTC")), + sa.Column("service_arrival_datetime", DateTime64(3, timezone="UTC")), + sa.Column("service_timezone", String), + sa.Column("service_budget_objective", Nullable(Decimal(10, 2))), + sa.Column("service_yield_objective", Nullable(Decimal(10, 2))), + sa.Column("service_capacity", Nullable(Int32)), + sa.Column("service_lid", Nullable(Int32)), + sa.Column("service_max_leg_load_factor", Nullable(Decimal(10, 2))), + sa.Column("service_max_leg_cumulative_sum_net_bookings", Nullable(Int32)), + sa.Column("service_optimization_status", LowCardinality(String)), + sa.Column("service_pointer", Boolean), + sa.Column("market_name", LowCardinality(String)), + sa.Column("market_rank", Int32), + sa.Column("market_list_alerts", Array(String)), + sa.Column("market_cabin_max_leg_load_factor", Nullable(Decimal(10, 2))), + sa.Column("market_cabin_max_leg_cumulative_sum_net_bookings", Nullable(Int32)), + sa.Column("od_id", Int64), + sa.Column("od_status", String), + sa.Column("od_origin_station_name", String), + sa.Column("od_destination_station_name", String), + sa.Column("od_departure_datetime", DateTime64(3, timezone="UTC")), + sa.Column("od_arrival_datetime", DateTime64(3, timezone="UTC")), + sa.Column("od_rank", Int32), + sa.Column("reference_service_number", Nullable(String)), + sa.Column("reference_service_departure_date", Nullable(Date)), + sa.Column("cabin_name", String), + sa.Column("od_cabin_forecasted_traffic", Nullable(Decimal(10, 2))), + sa.Column("od_cabin_forecasted_revenue_vat_inc", Nullable(Decimal(10, 2))), + sa.Column("od_cabin_forecasted_revenue_vat_exc", Nullable(Decimal(10, 2))), + sa.Column("od_cabin_optimized_traffic", Decimal(10, 2)), + sa.Column("od_cabin_optimized_revenue_vat_inc", Nullable(Decimal(10, 2))), + sa.Column("od_cabin_optimized_revenue_vat_exc", Nullable(Decimal(10, 2))), + sa.Column("od_cabin_last_predicted", Nullable(Decimal(10, 2))), + sa.Column("od_cabin_last_observed", Nullable(Decimal(10, 2))), + sa.Column("od_cabin_pointer", Boolean), + sa.Column("family_name", LowCardinality(String)), + sa.Column("bucket_name", LowCardinality(String)), + sa.Column("bucket_order", Int32), + sa.Column("has_pricing_changed_day_x", Boolean), + sa.Column("has_pricing_changed_bucket", Boolean), + sa.Column("cumulative_sum_net_revenue_vat_inc", Decimal(10, 2)), + sa.Column("cumulative_sum_net_revenue_vat_exc", Decimal(10, 2)), + sa.Column("cumulative_sum_net_ancillary_revenue_vat_inc", Decimal(10, 2)), + sa.Column("cumulative_sum_net_ancillary_revenue_vat_exc", Decimal(10, 2)), + sa.Column("sum_net_revenue_vat_inc", Decimal(10, 2)), + sa.Column("sum_net_revenue_vat_exc", Decimal(10, 2)), + sa.Column("sum_net_ancillary_revenue_vat_inc", Decimal(10, 2)), + sa.Column("sum_net_ancillary_revenue_vat_exc", Decimal(10, 2)), + sa.Column("availability_start_day", Nullable(Int32)), + sa.Column("availability_end_day", Nullable(Int32)), + sa.Column("cumul_availability_start_day", Int32), + sa.Column("cumul_availability_end_day", Int32), + sa.Column("cumulative_sum_net_bookings", Nullable(Int32)), + sa.Column("sum_confirmed_bookings", Nullable(Int32)), + sa.Column("sum_net_bookings", Nullable(Int32)), + sa.Column("is_first_available_start_day", Boolean), + sa.Column("is_first_available_end_day", Boolean), + sa.Column("price_vat_inc", Nullable(Decimal(10, 2))), + sa.Column("cabin_capacity", Nullable(Int32)), + sa.Column("cabin_lid", Nullable(Int32)), + sa.Column("bucket_authorization_start_day", Nullable(Int32)), + sa.Column("bucket_authorization_end_day", Nullable(Int32)), + sa.Column("forecast_full_day_x", Nullable(Int32)), + sa.Column("optimization_full_day_x", Nullable(Int32)), + sa.Column("unconstrained_demand_bookings", Nullable(Decimal(10, 2))), + sa.Column("unconstrained_demand_revenue", Nullable(Decimal(10, 2))), + sa.Column("unconstrained_forecast_bookings", Nullable(Decimal(10, 2))), + sa.Column("unconstrained_forecast_revenue", Nullable(Decimal(10, 2))), + sa.Column("is_last", Boolean), + engines.ReplicatedMergeTree( + table_path="/clickhouse/tables/{shard}/fact_daily_od_bucket/", + replica_name="{replica}", + primary_key=("time", "service_departure_date"), + order_by=( + "time", + "service_departure_date", + "service_number", + "day_x", + "service_id", + "od_id", + "bucket_name", + ), + ), + sa.text("INDEX is_last_index is_last TYPE set GRANULARITY 1"), + clickhouse_cluster="wm", + ) + + op.create_table( + "fact_passenger_event", + sa.Column("time", DateTime64(3, timezone="UTC")), + sa.Column("day_x", Int32), + sa.Column("event_datetime", DateTime64(3, timezone="UTC")), + sa.Column("transporter_type", LowCardinality(String)), + sa.Column("entity_name", LowCardinality(String)), + sa.Column("route_name", LowCardinality(String)), + sa.Column("service_id", Int64), + sa.Column("service_departure_date", Date), + sa.Column("service_number", String), + sa.Column("service_labels", Array(String)), + sa.Column("service_status", LowCardinality(String)), + sa.Column("service_origin_station_name", String), + sa.Column("service_destination_station_name", String), + sa.Column("service_departure_datetime", DateTime64(3, timezone="UTC")), + sa.Column("service_arrival_datetime", DateTime64(3, timezone="UTC")), + sa.Column("service_timezone", LowCardinality(String)), + sa.Column("market_id", Int64), + sa.Column("market_name", LowCardinality(String)), + sa.Column("market_rank", Int32), + sa.Column("od_id", Int64), + sa.Column("od_status", String), + sa.Column("od_origin_station_name", String), + sa.Column("od_destination_station_name", String), + sa.Column("od_destination_station_id", Int64), + sa.Column("od_departure_datetime", DateTime64(3, timezone="UTC")), + sa.Column("od_arrival_datetime", DateTime64(3, timezone="UTC")), + sa.Column("od_rank", Int32), + sa.Column("cabin_name", LowCardinality(String)), + sa.Column("family_name", LowCardinality(String)), + sa.Column("bucket_name", LowCardinality(String)), + sa.Column("bucket_order", Int32), + sa.Column("cabin_capacity", Nullable(Int32)), + sa.Column("cabin_lid", Nullable(Int32)), + sa.Column("physical_inventory_name", String), + sa.Column("price_vat_inc", Decimal(10, 2)), + sa.Column("price_vat_exc", Decimal(10, 2)), + sa.Column("base_price_vat_inc", Decimal(10, 2)), + sa.Column("bucket_price_vat_inc", Nullable(Decimal(10, 2))), + sa.Column("ancillary_revenue_vat_inc", Decimal(10, 2)), + sa.Column("ancillary_revenue_vat_exc", Decimal(10, 2)), + sa.Column("event_type", LowCardinality(String)), + sa.Column("booking_key", String), + sa.Column("customer_key", Nullable(String)), + sa.Column("ticket_category", Nullable(String)), + sa.Column("ticket_key", String), + sa.Column("fare_code", Nullable(String)), + sa.Column("sale_channel_name", Nullable(String)), + sa.Column("sale_channel_type", Nullable(String)), + sa.Column("is_ticket_exchanged", Boolean), + sa.Column("is_confirmed", Boolean), + sa.Column("no_show", Boolean), + engines.ReplicatedMergeTree( + table_path="/clickhouse/tables/{shard}/fact_passenger_event/", + replica_name="{replica}", + primary_key=("time", "service_departure_date"), + order_by=( + "time", + "service_departure_date", + "service_number", + "day_x", + "service_id", + "od_id", + ), + ), + clickhouse_cluster="wm", + ) + + +def downgrade() -> None: + op.drop_table("fact_daily_leg_bucket") + op.drop_table("fact_daily_leg_physical_inventory") + op.drop_table("fact_daily_od_bucket") + op.drop_table("fact_passenger_event") diff --git a/example/alembic/postgresl/README b/example/alembic/postgresl/README new file mode 100644 index 0000000..2500aa1 --- /dev/null +++ b/example/alembic/postgresl/README @@ -0,0 +1 @@ +Generic single-database configuration. diff --git a/example/alembic/env.py b/example/alembic/postgresl/env.py similarity index 99% rename from example/alembic/env.py rename to example/alembic/postgresl/env.py index 494647a..0591c35 100644 --- a/example/alembic/env.py +++ b/example/alembic/postgresl/env.py @@ -1,8 +1,7 @@ from logging.config import fileConfig -from sqlalchemy import engine_from_config, pool, text - from alembic import context +from sqlalchemy import engine_from_config, pool, text # this is the Alembic Config object, which provides # access to the values within the .ini file in use. diff --git a/example/alembic/script.py.mako b/example/alembic/postgresl/script.py.mako similarity index 100% rename from example/alembic/script.py.mako rename to example/alembic/postgresl/script.py.mako diff --git a/example/alembic/versions/init.py b/example/alembic/postgresl/versions/init.py similarity index 100% rename from example/alembic/versions/init.py rename to example/alembic/postgresl/versions/init.py diff --git a/example/helm-chart/my-example-chart/values.yaml b/example/helm-chart/my-example-chart/values.yaml index f19bd39..6083b3a 100644 --- a/example/helm-chart/my-example-chart/values.yaml +++ b/example/helm-chart/my-example-chart/values.yaml @@ -13,7 +13,7 @@ chartreuse: postgresql: image: - tag: 11.7.0-debian-10-r9 + tag: latest debug: true postgresqlUsername: foo postgresqlPassword: foo diff --git a/example/multi-database-config.yaml b/example/multi-database-config.yaml new file mode 100644 index 0000000..73bd247 --- /dev/null +++ b/example/multi-database-config.yaml @@ -0,0 +1,58 @@ +# Chartreuse Multi-Database Configuration +# +# This YAML file defines multiple databases that Chartreuse should manage. +# Database names are used as keys for easy organization. + +databases: + # Main application database + main: + alembic_directory_path: /app/alembic/main + alembic_config_file_path: alembic.ini + # Database connection components (URL will be built automatically) + dialect: postgresql + user: app_user + password: app_password + host: postgres-main + port: 5432 + database: app_main + allow_migration_for_empty_database: true + additional_parameters: "" + + # Analytics database + analytics: + alembic_directory_path: /app/alembic/analytics + alembic_config_file_path: alembic.ini + dialect: postgresql + user: analytics_user + password: analytics_password + host: postgres-analytics + port: 5432 + database: analytics + allow_migration_for_empty_database: false + additional_parameters: "" + + # Audit database with custom parameters + audit: + alembic_directory_path: /app/alembic/audit + alembic_config_file_path: alembic.ini + dialect: postgresql + user: audit_user + password: audit_password + host: postgres-audit + port: 5432 + database: audit_logs + allow_migration_for_empty_database: true + additional_parameters: "-x audit_mode=yes" + + # Example with environment variable substitution (if supported by your deployment) + reports: + alembic_directory_path: /app/alembic/reports + alembic_config_file_path: alembic.ini + dialect: postgresql + user: "${REPORTS_DB_USER}" + password: "${REPORTS_DB_PASSWORD}" + host: postgres-reports + port: 5432 + database: reports + allow_migration_for_empty_database: false + additional_parameters: "" \ No newline at end of file diff --git a/example/simple-config.yaml b/example/simple-config.yaml new file mode 100644 index 0000000..9d59389 --- /dev/null +++ b/example/simple-config.yaml @@ -0,0 +1,25 @@ +# Simple example configuration with dictionary format +databases: + # Main database + main: + alembic_directory_path: /app/alembic/main + alembic_config_file_path: alembic.ini + dialect: postgresql + user: myuser + password: mypassword + host: localhost + port: 5432 + database: myapp + allow_migration_for_empty_database: true + + # Cache database + cache: + alembic_directory_path: /app/alembic/cache + alembic_config_file_path: alembic.ini + dialect: postgresql + user: cache_user + password: cache_pass + host: redis-db + port: 5432 + database: cache_db + allow_migration_for_empty_database: false \ No newline at end of file diff --git a/pyproject.toml b/pyproject.toml index 17de5f2..862155f 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,6 +1,79 @@ -[tool.black] +[project] +name = "chartreuse" +description="Helper for Alembic migrations within Kubernetes." +dynamic = ["version"] +readme = "README.md" +authors = [{ name = "wiremind", email = "dev@wiremind.io" }] +license="LGPL-3.0-or-later" +urls = { github = "https://github.com/wiremind/chartreuse"} +scripts = {chartreuse-upgrade = "chartreuse.chartreuse_upgrade:main"} +requires-python = ">=3.11.0" + +dependencies = [ + "alembic", + "psycopg2", + "pydantic>=2.0.0", + "PyYAML", + "wiremind-kubernetes~=7.0", +] + +[build-system] +requires = ["setuptools>=71.1"] +build-backend = "setuptools.build_meta" + +[tool.setuptools] +include-package-data = true +zip-safe = true +[tool.setuptools.dynamic] +version = { file = "VERSION" } +[tool.setuptools.packages.find] +where = ["src"] +exclude = ["*.tests", "*.tests.*", "tests.*", "tests"] + +[project.optional-dependencies] +test = [ + "mock", + "pytest", + "pytest-mock", +] +ruff = [ + "ruff", +] +dev = [ + "ruff", + "pip-tools", + + "mock", + "pytest", + "pytest-mock", +] + +[tool.ruff] line-length = 120 -target-version = ['py37'] +target-version = "py311" + +[tool.ruff.lint] +select = [ + "E", # pycodestyle errors + "W", # pycodestyle warnings + "F", # pyflakes + "I", # isort + "UP", # pyupgrade + "B", # flake8-bugbear + "C4", # flake8-comprehensions +] +ignore = [] + +[tool.ruff.format] +quote-style = "double" +indent-style = "space" +line-ending = "auto" + +[tool.ruff.lint.isort] +known-first-party = ["chartreuse"] -[tool.isort] -line_length = 120 +[tool.pytest.ini_options] +log_level = "INFO" +# Deterministic ordering for tests; useful for pytest-xdist. +env = ["PYTHONHASHSEED=0"] +filterwarnings = ["ignore::pytest.PytestUnknownMarkWarning"] \ No newline at end of file diff --git a/requirements.txt b/requirements.txt index 71c26cb..84c02ef 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,68 +1,79 @@ -# -# This file is autogenerated by pip-compile with Python 3.11 -# by the following command: -# -# pip-compile --no-emit-index-url --output-file=requirements.txt pyproject.toml -# -alembic==1.13.2 +# This file was autogenerated by uv via the following command: +# uv pip compile --no-emit-index-url --output-file=requirements.txt pyproject.toml +alembic==1.16.5 # via chartreuse (pyproject.toml) -cachetools==5.4.0 +annotated-types==0.7.0 + # via pydantic +cachetools==6.2.0 # via google-auth -certifi==2024.7.4 +certifi==2025.10.5 # via # kubernetes # requests -charset-normalizer==3.3.2 +charset-normalizer==3.4.3 # via requests -google-auth==2.32.0 +durationpy==0.10 # via kubernetes -idna==3.7 +google-auth==2.41.1 + # via kubernetes +greenlet==3.2.4 + # via sqlalchemy +idna==3.10 # via requests -kubernetes==30.1.0 +kubernetes==34.1.0 # via wiremind-kubernetes -mako==1.3.5 +mako==1.3.10 # via alembic -markupsafe==2.1.5 +markupsafe==3.0.3 # via mako -oauthlib==3.2.2 - # via - # kubernetes - # requests-oauthlib -psycopg2==2.9.9 +oauthlib==3.3.1 + # via requests-oauthlib +psycopg2==2.9.10 # via chartreuse (pyproject.toml) -pyasn1==0.6.0 +pyasn1==0.6.1 # via # pyasn1-modules # rsa -pyasn1-modules==0.4.0 +pyasn1-modules==0.4.2 # via google-auth +pydantic==2.12.0 + # via chartreuse (pyproject.toml) +pydantic-core==2.41.1 + # via pydantic python-dateutil==2.9.0.post0 # via kubernetes -pyyaml==6.0.1 - # via kubernetes -requests==2.32.4 +pyyaml==6.0.3 + # via + # chartreuse (pyproject.toml) + # kubernetes +requests==2.32.5 # via # kubernetes # requests-oauthlib requests-oauthlib==2.0.0 # via kubernetes -rsa==4.9 +rsa==4.9.1 # via google-auth -six==1.16.0 +six==1.17.0 # via # kubernetes # python-dateutil -sqlalchemy==2.0.31 +sqlalchemy==2.0.43 # via alembic -typing-extensions==4.12.2 +typing-extensions==4.15.0 # via # alembic + # pydantic + # pydantic-core # sqlalchemy -urllib3==2.5.0 + # typing-inspection +typing-inspection==0.4.2 + # via pydantic +urllib3==2.3.0 # via # kubernetes # requests -websocket-client==1.8.0 +websocket-client==1.9.0 # via kubernetes -wiremind-kubernetes==7.4.5 +wiremind-kubernetes==7.5.0 # via chartreuse (pyproject.toml) diff --git a/scripts/validate_config.py b/scripts/validate_config.py new file mode 100644 index 0000000..828084d --- /dev/null +++ b/scripts/validate_config.py @@ -0,0 +1,130 @@ +#!/usr/bin/env python3 +""" +Configuration validator for Chartreuse multi-database setup. + +Usage: python validate_config.py path/to/config.yaml +""" + +import sys +from pathlib import Path + +import yaml + + +def build_database_url(db_config: dict) -> str: + """Build database URL from components.""" + required_components = ["dialect", "user", "password", "host", "port", "database"] + missing_components = [comp for comp in required_components if comp not in db_config or not db_config[comp]] + + if missing_components: + raise ValueError(f"Missing components: {missing_components}") + + dialect = db_config["dialect"] + user = db_config["user"] + password = db_config["password"] + host = db_config["host"] + port = db_config["port"] + database = db_config["database"] + + return f"{dialect}://{user}:{password}@{host}:{port}/{database}" + + +def validate_config(config_path: str) -> bool: + """Validate a multi-database configuration file.""" + + try: + # Check if file exists + if not Path(config_path).exists(): + print(f"❌ Configuration file not found: {config_path}") + return False + + # Load and parse YAML + with open(config_path) as f: + config = yaml.safe_load(f) + + # Validate structure + if "databases" not in config: + print("❌ Configuration must contain 'databases' key") + return False + + databases = config["databases"] + if not isinstance(databases, dict): + print("❌ 'databases' must be a dictionary with database names as keys") + return False + + if len(databases) == 0: + print("❌ At least one database must be configured") + return False + + # Validate each database + required_fields = [ + "alembic_directory_path", + "alembic_config_file_path", + "dialect", + "user", + "password", + "host", + "port", + "database", + ] + + for db_name, db_config in databases.items(): + print(f"Validating database: {db_name}") + + # Check required fields + for field in required_fields: + if field not in db_config: + print(f"❌ Database '{db_name}' missing required field: {field}") + return False + if not db_config[field]: + print(f"❌ Database '{db_name}' field '{field}' cannot be empty") + return False + + # Validate URL building + try: + url = build_database_url(db_config) + print(f" 📍 Built URL: {url}") + + # Basic URL format validation + if not url.startswith(("postgresql://", "mysql://", "sqlite://")): + print( + f"⚠️ Database {db_name}: URL format may be invalid (expected postgresql://, mysql://, or sqlite://)" + ) + + except ValueError as e: + print(f"❌ Database {db_name}: {e}") + return False + + # Check optional boolean fields + for bool_field in ["allow_migration_for_empty_database"]: + if bool_field in db_config and not isinstance(db_config[bool_field], bool): + print(f"❌ Database {db_name}: '{bool_field}' must be true or false") + return False + + print(f"✅ Database {db_name}: OK") + + print("\n✅ Configuration is valid!") + print(f"Found {len(databases)} database(s): {', '.join(databases.keys())}") + return True + + except yaml.YAMLError as e: + print(f"❌ Invalid YAML: {e}") + return False + except Exception as e: + print(f"❌ Error validating configuration: {e}") + return False + + +def main() -> None: + if len(sys.argv) != 2: + print("Usage: python validate_config.py path/to/config.yaml") + sys.exit(1) + + config_path = sys.argv[1] + is_valid = validate_config(config_path) + + sys.exit(0 if is_valid else 1) + + +if __name__ == "__main__": + main() diff --git a/setup.cfg b/setup.cfg deleted file mode 100644 index fcd8c6a..0000000 --- a/setup.cfg +++ /dev/null @@ -1,34 +0,0 @@ -[options] -python_requires = >= 3.11 - -[flake8] -max-line-length = 120 -# Enable flake8-mutable -enable-extensions = M511 -# W503 line break before binary operator. Black causes this error -# E203 whitespace before ':' -# E231 missing whitespace after ','. Black causes this error but seems to be right (see pep8) -# E501 line too long -ignore = W503, E203, E231, E501 -jobs = 4 - -[mypy] -python_version = 3.11 -ignore_missing_imports = True -check_untyped_defs = True -disallow_untyped_defs = True -disallow_incomplete_defs = True -warn_redundant_casts = True -warn_unused_ignores = True -warn_unused_configs = True -no_implicit_optional = True -show_error_codes = True -files = src,example/alembic - -[tool:pytest] -log_level=INFO -# Deterministic ordering for tests; useful for pytest-xdist. -env = - PYTHONHASHSEED=0 -filterwarnings = - ignore::pytest.PytestUnknownMarkWarning diff --git a/setup.py b/setup.py deleted file mode 100644 index d5fcc09..0000000 --- a/setup.py +++ /dev/null @@ -1,56 +0,0 @@ -""" -Chartreuse -""" -from setuptools import find_packages, setup - -with open("VERSION") as version_file: - version = version_file.read().strip() - - -extra_require_test = [ - "mock", - "pytest", - "pytest-mock", -] -extra_require_mypy = [ - "mypy", -] -extra_require_dev = ( - [ - "black", - "flake8", - "flake8-mutable", - "pip-tools", - ] - + extra_require_mypy - + extra_require_test -) - - -setup( - name="chartreuse", - version=version, - description="Helper for Alembic migrations within Kubernetes.", - long_description="Helper for Alembic migrations within Kubernetes.", - platforms=["posix"], - author="wiremind", - author_email="dev@wiremind.io", - url="https://github.com/wiremind/chartreuse", - license="LGPLv3+", - packages=find_packages("src", exclude=["*.tests", "*.tests.*", "tests.*", "tests"]), - package_dir={"": "src"}, - include_package_data=True, - zip_safe=True, - python_requires=">=3.11.0", - install_requires=[ - "alembic", - "psycopg2", - "wiremind-kubernetes~=7.0", - ], - extras_require={ - "dev": extra_require_dev, - "mypy": extra_require_mypy, - "test": extra_require_test, - }, - entry_points={"console_scripts": ["chartreuse-upgrade=chartreuse.chartreuse_upgrade:main"]}, -) diff --git a/src/chartreuse/chartreuse.py b/src/chartreuse/chartreuse.py index d3f6621..7582986 100644 --- a/src/chartreuse/chartreuse.py +++ b/src/chartreuse/chartreuse.py @@ -2,8 +2,11 @@ import wiremind_kubernetes.kubernetes_helper +from .config_loader import DatabaseConfig from .utils import AlembicMigrationHelper +logger = logging.getLogger(__name__) + def configure_logging() -> None: logging.basicConfig( @@ -14,26 +17,42 @@ def configure_logging() -> None: class Chartreuse: + """Handles single or multiple database migrations.""" + def __init__( self, - alembic_directory_path: str, - alembic_config_file_path: str, - postgresql_url: str, + databases_config: dict[str, DatabaseConfig], release_name: str, - alembic_allow_migration_for_empty_database: bool, - alembic_additional_parameters: str = "", kubernetes_helper: wiremind_kubernetes.kubernetes_helper.KubernetesDeploymentManager | None = None, ): configure_logging() - self.alembic_migration_helper = AlembicMigrationHelper( - alembic_directory_path=alembic_directory_path, - alembic_config_file_path=alembic_config_file_path, - database_url=postgresql_url, - allow_migration_for_empty_database=alembic_allow_migration_for_empty_database, - additional_parameters=alembic_additional_parameters, - ) + self.databases_config = databases_config + self.migration_helpers: dict[str, AlembicMigrationHelper] = {} + + # Initialize migration helpers for each database + for db_name, db_config in databases_config.items(): + logger.info("Initializing migration helper for database: %s", db_name) + + # Build additional parameters with section name + additional_params = db_config.additional_parameters + + # Use database name as section name for alembic -n parameter + section_param = f"-n {db_name}" + additional_params = f"{additional_params} {section_param}".strip() + + helper = AlembicMigrationHelper( + alembic_directory_path=db_config.alembic_directory_path, + alembic_config_file_path=db_config.alembic_config_file_path, + database_url=db_config.url, + allow_migration_for_empty_database=db_config.allow_migration_for_empty_database, + additional_parameters=additional_params, + alembic_section_name=db_name, + ) + self.migration_helpers[db_name] = helper + + # Initialize Kubernetes helper if kubernetes_helper: self.kubernetes_helper = kubernetes_helper else: @@ -41,11 +60,21 @@ def __init__( use_kubeconfig=None, release_name=release_name ) - self.is_migration_needed = self.check_migration_needed() - - def check_migration_needed(self) -> bool: - return self.alembic_migration_helper.is_migration_needed + @property + def is_migration_needed(self) -> bool: + """Check if any database needs migration.""" + for db_name, helper in self.migration_helpers.items(): + if helper.is_migration_needed: + logger.info("Database '%s' needs migration", db_name) + return True + return False def upgrade(self) -> None: - if self.check_migration_needed(): - self.alembic_migration_helper.upgrade_db() + """Upgrade all databases that need migration.""" + for db_name, helper in self.migration_helpers.items(): + if helper.is_migration_needed: + logger.info("Upgrading database: %s", db_name) + helper.upgrade_db() + logger.info("Successfully upgraded database: %s", db_name) + else: + logger.info("Database '%s' is up to date", db_name) diff --git a/src/chartreuse/chartreuse_upgrade.py b/src/chartreuse/chartreuse_upgrade.py index 2c132f1..630a6ec 100755 --- a/src/chartreuse/chartreuse_upgrade.py +++ b/src/chartreuse/chartreuse_upgrade.py @@ -6,6 +6,7 @@ from chartreuse import get_version from .chartreuse import Chartreuse +from .config_loader import load_multi_database_config logger = logging.getLogger(__name__) @@ -34,59 +35,79 @@ def ensure_safe_run() -> None: ) +def validate_config_file_path() -> str: + """ + Validate that the multi-database configuration file path is properly configured and accessible. + + Returns: + str: The validated config file path + + Raises: + ValueError: If environment variable is not set or path is not a file + FileNotFoundError: If the config file doesn't exist + """ + multi_config_path = os.environ.get("CHARTREUSE_MULTI_CONFIG_PATH") + + if not multi_config_path: + raise ValueError("CHARTREUSE_MULTI_CONFIG_PATH environment variable is required but not set") + + if not os.path.exists(multi_config_path): + raise FileNotFoundError(f"Multi-database configuration file not found: {multi_config_path}") + + if not os.path.isfile(multi_config_path): + raise ValueError(f"Multi-database configuration path is not a file: {multi_config_path}") + + return multi_config_path + + def main() -> None: """ When put in a post-install Helm hook, if this program fails the whole release is considered as failed. """ ensure_safe_run() - ALEMBIC_DIRECTORY_PATH: str = os.environ.get("CHARTREUSE_ALEMBIC_DIRECTORY_PATH", "/app/alembic") - ALEMBIC_CONFIG_FILE_PATH: str = os.environ.get("CHARTREUSE_ALEMBIC_CONFIG_FILE_PATH", "alembic.ini") - POSTGRESQL_URL: str = os.environ["CHARTREUSE_ALEMBIC_URL"] - ALEMBIC_ALLOW_MIGRATION_FOR_EMPTY_DATABASE: bool = bool( - os.environ["CHARTREUSE_ALEMBIC_ALLOW_MIGRATION_FOR_EMPTY_DATABASE"] - ) - ALEMBIC_ADDITIONAL_PARAMETERS: str = os.environ["CHARTREUSE_ALEMBIC_ADDITIONAL_PARAMETERS"] - ENABLE_STOP_PODS: bool = bool(os.environ["CHARTREUSE_ENABLE_STOP_PODS"]) + + # Validate and get multi-database configuration path + multi_config_path = validate_config_file_path() + + # Use multi-database configuration + logger.info("Using multi-database configuration from: %s", multi_config_path) + + try: + databases_config = load_multi_database_config(multi_config_path) + except (FileNotFoundError, ValueError) as e: + logger.error("Failed to load multi-database configuration: %s", e) + raise + + ENABLE_STOP_PODS: bool = os.environ.get("CHARTREUSE_ENABLE_STOP_PODS", "true").lower() not in ("", "false", "0") RELEASE_NAME: str = os.environ["CHARTREUSE_RELEASE_NAME"] - UPGRADE_BEFORE_DEPLOYMENT: bool = bool(os.environ["CHARTREUSE_UPGRADE_BEFORE_DEPLOYMENT"]) - HELM_IS_INSTALL: bool = bool(os.environ["HELM_IS_INSTALL"]) + UPGRADE_BEFORE_DEPLOYMENT: bool = os.environ.get("CHARTREUSE_UPGRADE_BEFORE_DEPLOYMENT", "false").lower() not in ( + "", + "false", + "0", + ) + HELM_IS_INSTALL: bool = os.environ.get("HELM_IS_INSTALL", "false").lower() not in ("", "false", "0") deployment_manager = KubernetesDeploymentManager(release_name=RELEASE_NAME, use_kubeconfig=None) chartreuse = Chartreuse( - alembic_directory_path=ALEMBIC_DIRECTORY_PATH, - alembic_config_file_path=ALEMBIC_CONFIG_FILE_PATH, - postgresql_url=POSTGRESQL_URL, - alembic_allow_migration_for_empty_database=ALEMBIC_ALLOW_MIGRATION_FOR_EMPTY_DATABASE, - alembic_additional_parameters=ALEMBIC_ADDITIONAL_PARAMETERS, + databases_config=databases_config, release_name=RELEASE_NAME, kubernetes_helper=deployment_manager, ) + if chartreuse.is_migration_needed: if ENABLE_STOP_PODS: - # If ever Helm has scaled up the pods that were stopped in predeployment. deployment_manager.stop_pods() - # The exceptions this method raises should NEVER be caught. - # If the migration fails, the post-install should fail and the release will fail - # we will end up with the old release. chartreuse.upgrade() if not ENABLE_STOP_PODS: - # Do not start pods return if UPGRADE_BEFORE_DEPLOYMENT and not HELM_IS_INSTALL: - # Do not start pods in case of helm upgrade (not install, aka initial deployment) in "before" mode return - # Scale up the new pods, only if chartreuse is: - # in "upgrade after deployment" mode - # or in "upgrade before deployment" mode, during initial install - - # We can fail and abort all, but if we're not that demanding we can start the pods manually - # via mayo for example try: deployment_manager.start_pods() - except: # noqa: E722 + except Exception: logger.error("Couldn't scale up new pods in chartreuse_upgrade after migration, SHOULD BE DONE MANUALLY ! ") diff --git a/src/chartreuse/config_loader.py b/src/chartreuse/config_loader.py new file mode 100644 index 0000000..85fc50d --- /dev/null +++ b/src/chartreuse/config_loader.py @@ -0,0 +1,132 @@ +""" +Configuration loader for multi-database support using Pydantic models. +""" + +import logging + +import yaml +from pydantic import BaseModel, Field, computed_field, field_validator + +logger = logging.getLogger(__name__) + + +class DatabaseConfig(BaseModel): + """Configuration for a single database within a multi-database setup.""" + + # Database connection components + dialect: str = Field(..., description="Database dialect (e.g., postgresql)") + user: str = Field(..., description="Database user") + password: str = Field(..., description="Database password") + host: str = Field(..., description="Database host") + port: int = Field(..., description="Database port", gt=0, le=65535) + database: str = Field(..., description="Database name") + + # Alembic configuration + alembic_directory_path: str = Field(..., description="Path to Alembic directory") + alembic_config_file_path: str = Field(default="alembic.ini", description="Path to Alembic config file") + + # 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") + + @computed_field + @property + def url(self) -> str: + """Build database URL from components.""" + return f"{self.dialect}://{self.user}:{self.password}@{self.host}:{self.port}/{self.database}" + + @field_validator("dialect") + @classmethod + def validate_dialect(cls, v: str) -> str: + """Validate database dialect.""" + supported_dialects = ["postgresql", "mysql", "sqlite", "oracle", "mssql", "clickhouse"] + if v.lower() not in supported_dialects: + logger.warning("Dialect '%s' might not be supported. Supported dialects: %s", v, supported_dialects) + return v + + @field_validator("additional_parameters") + @classmethod + def validate_additional_parameters(cls, v: str) -> str: + """Ensure additional_parameters is a string.""" + return v or "" + + +class MultiDatabaseConfig(BaseModel): + """Multi-database configuration - the only supported configuration format.""" + + databases: dict[str, DatabaseConfig] = Field(..., description="Dictionary of database configurations") + + @field_validator("databases") + @classmethod + def validate_databases_not_empty(cls, v: dict[str, DatabaseConfig]) -> dict[str, DatabaseConfig]: + """Ensure at least one database is configured.""" + if not v: + raise ValueError("At least one database must be configured") + return v + + +def load_multi_database_config(config_path: str) -> dict[str, DatabaseConfig]: + """ + Load multi-database configuration from YAML file using Pydantic validation. + + This is the only supported configuration format. Single database setups should + be configured as a multi-database config with one entry. + + Args: + config_path: Path to YAML configuration file + + Returns: + Dictionary of database configurations keyed by database name + + Raises: + FileNotFoundError: If configuration file is not found + ValueError: If configuration is invalid + + Expected YAML format: + ```yaml + databases: + main: # Database name/identifier + dialect: postgresql + user: app_user + password: app_password + host: postgres-main + port: 5432 + database: app_main + alembic_directory_path: /app/alembic/main + alembic_config_file_path: alembic.ini + allow_migration_for_empty_database: true + additional_parameters: "" + + # For single database setups, just include one database: + # analytics: + # dialect: postgresql + # user: analytics_user + # password: analytics_password + # host: postgres-analytics + # port: 5432 + # database: analytics + # alembic_directory_path: /app/alembic/analytics + # alembic_config_file_path: alembic.ini + # allow_migration_for_empty_database: false + # additional_parameters: "" + ``` + """ + try: + with open(config_path) as f: + raw_config = yaml.safe_load(f) + + # Validate the configuration using Pydantic + config = MultiDatabaseConfig(**raw_config) + + logger.info( + "Loaded configuration for %d database(s): %s", len(config.databases), ", ".join(config.databases.keys()) + ) + + return config.databases + + except FileNotFoundError as err: + raise FileNotFoundError(f"Configuration file not found: {config_path}") from err + except yaml.YAMLError as e: + raise ValueError(f"Invalid YAML in configuration file: {e}") from e + except Exception as e: + raise ValueError(f"Invalid configuration: {e}") from e diff --git a/src/chartreuse/tests/conftest.py b/src/chartreuse/tests/conftest.py index 813c099..1232161 100644 --- a/src/chartreuse/tests/conftest.py +++ b/src/chartreuse/tests/conftest.py @@ -3,21 +3,65 @@ from pytest_mock.plugin import MockerFixture -def configure_os_environ_mock( - mocker: MockerFixture, additional_environment: dict[str, str] | None = None -) -> None: - new_environ: dict[str, str] = dict( - CHARTREUSE_ALEMBIC_ADDITIONAL_PARAMETERS="", - CHARTREUSE_ALEMBIC_ALLOW_MIGRATION_FOR_EMPTY_DATABASE="1", - CHARTREUSE_ALEMBIC_URL="foo", - CHARTREUSE_ENABLE_STOP_PODS="1", - CHARTREUSE_MIGRATE_IMAGE_PULL_SECRET="foo", - CHARTREUSE_RELEASE_NAME="foo", - RUN_TEST_IN_KIND=os.environ.get("RUN_TEST_IN_KIND", ""), - CHARTREUSE_UPGRADE_BEFORE_DEPLOYMENT="", - HELM_IS_INSTALL="", - ) +def configure_os_environ_mock(mocker: MockerFixture, additional_environment: dict[str, str] | None = None) -> None: + new_environ: dict[str, str] = { + "CHARTREUSE_ALEMBIC_ADDITIONAL_PARAMETERS": "", + "CHARTREUSE_ALEMBIC_ALLOW_MIGRATION_FOR_EMPTY_DATABASE": "1", + "CHARTREUSE_ALEMBIC_URL": "foo", + "CHARTREUSE_ENABLE_STOP_PODS": "1", + "CHARTREUSE_MIGRATE_IMAGE_PULL_SECRET": "foo", + "CHARTREUSE_RELEASE_NAME": "foo", + "RUN_TEST_IN_KIND": os.environ.get("RUN_TEST_IN_KIND", ""), + "CHARTREUSE_UPGRADE_BEFORE_DEPLOYMENT": "", + "HELM_IS_INSTALL": "", + "CHARTREUSE_MULTI_CONFIG_PATH": "/mock/config.yaml", + # Add Kubernetes environment variables to prevent config loading + "KUBERNETES_SERVICE_HOST": "localhost", + "KUBERNETES_SERVICE_PORT": "443", + } if additional_environment: new_environ.update(additional_environment) mocker.patch.dict(os.environ, new_environ) + + # Mock the file operations for the config path + mocker.patch("os.path.exists", return_value=True) + mocker.patch("os.path.isfile", return_value=True) + + # Mock the open operation for kubernetes token file + mock_open = mocker.mock_open(read_data="mock-token") + mocker.patch("builtins.open", mock_open) + + # Mock the config loader + mock_config = [ + { + "name": "default", + "alembic_directory_path": "/app/alembic", + "alembic_config_file_path": "alembic.ini", + "url": "foo", + "allow_migration_for_empty_database": True, + "additional_parameters": "", + } + ] + mocker.patch("chartreuse.chartreuse_upgrade.load_multi_database_config", return_value=mock_config) + + +def configure_chartreuse_mock( + mocker: MockerFixture, + is_migration_needed: bool, +) -> None: + """ + Configure a mock for Chartreuse object. + """ + mock_chartreuse = mocker.MagicMock() + mock_chartreuse.is_migration_needed = is_migration_needed + mocker.patch("chartreuse.chartreuse_upgrade.Chartreuse", return_value=mock_chartreuse) + + # Mock KubernetesDeploymentManager in all the right places + mock_kdm = mocker.MagicMock() + mocker.patch("wiremind_kubernetes.KubernetesDeploymentManager", return_value=mock_kdm) + mocker.patch("chartreuse.chartreuse_upgrade.KubernetesDeploymentManager", return_value=mock_kdm) + + # Also mock the kubernetes config loading to prevent the config errors + mocker.patch("kubernetes.config.load_incluster_config") + mocker.patch("wiremind_kubernetes.kube_config.load_kubernetes_config") diff --git a/src/chartreuse/tests/e2e_tests/conftest.py b/src/chartreuse/tests/e2e_tests/conftest.py index c330701..bed5fe4 100644 --- a/src/chartreuse/tests/e2e_tests/conftest.py +++ b/src/chartreuse/tests/e2e_tests/conftest.py @@ -42,21 +42,31 @@ def _cluster_init(include_chartreuse: bool, pre_upgrade: bool = False) -> Genera else: additional_args = "--set chartreuse.enabled=true --set chartreuse.upgradeBeforeDeployment=false" run_command( - f"helm install --wait {TEST_RELEASE} {HELM_CHART_PATH} --namespace {TEST_NAMESPACE} --timeout 180s {additional_args}", + f"""helm install --wait {TEST_RELEASE} {HELM_CHART_PATH} + --namespace {TEST_NAMESPACE} --timeout 180s {additional_args}""", cwd=EXAMPLE_PATH, ) run_command( - f"helm upgrade --wait {TEST_RELEASE} {HELM_CHART_PATH} --namespace {TEST_NAMESPACE} --timeout 60s {additional_args}", + f"""helm upgrade --wait {TEST_RELEASE} {HELM_CHART_PATH} + --namespace {TEST_NAMESPACE} --timeout 60s {additional_args}""", cwd=EXAMPLE_PATH, ) kubectl_port_forwardpostgresql = subprocess.Popen( - ["kubectl", "port-forward", "--namespace", TEST_NAMESPACE, f"{TEST_RELEASE}-postgresql-0", "5432"] + [ + "kubectl", + "port-forward", + "--namespace", + TEST_NAMESPACE, + f"{TEST_RELEASE}-postgresql-0", + "5432", + ] ) time.sleep(5) # Hack to wait for k exec to be up except: # noqa run_command( - f"kubectl logs --selector app.kubernetes.io/instance=e2e-test-release --all-containers=false --namespace {TEST_NAMESPACE} --tail 1000" + f"""kubectl logs --selector app.kubernetes.io/instance=e2e-test-release + --all-containers=false --namespace {TEST_NAMESPACE} --tail 1000""" ) run_command(f"kubectl delete namespace {TEST_NAMESPACE} --grace-period=1") raise @@ -110,14 +120,22 @@ def populate_cluster_with_chartreuse_pre_upgrade() -> Generator: def assert_sql_upgraded() -> None: - assert inspect(sqlalchemy.create_engine(POSTGRESQL_URL)).get_table_names() == ["alembic_version", "upgraded"] + assert inspect(sqlalchemy.create_engine(POSTGRESQL_URL)).get_table_names() == [ + "alembic_version", + "upgraded", + ] def assert_sql_not_upgraded() -> None: - assert not inspect(sqlalchemy.create_engine(POSTGRESQL_URL)).get_table_names() == ["alembic_version", "upgraded"] + assert not inspect(sqlalchemy.create_engine(POSTGRESQL_URL)).get_table_names() == [ + "alembic_version", + "upgraded", + ] def are_pods_scaled_down() -> bool: return KubernetesDeploymentManager( - release_name=TEST_RELEASE, namespace=TEST_NAMESPACE, should_load_kubernetes_config=False + release_name=TEST_RELEASE, + namespace=TEST_NAMESPACE, + should_load_kubernetes_config=False, ).is_deployment_stopped("e2e-test-release-my-test-chart") diff --git a/src/chartreuse/tests/e2e_tests/test_blackbox.py b/src/chartreuse/tests/e2e_tests/test_blackbox.py deleted file mode 100644 index 08de5de..0000000 --- a/src/chartreuse/tests/e2e_tests/test_blackbox.py +++ /dev/null @@ -1,31 +0,0 @@ -import logging - -from pytest_mock.plugin import MockerFixture - -from .conftest import are_pods_scaled_down, assert_sql_upgraded - -logger = logging.getLogger(__name__) - - -def test_chartreuse_blackbox_post_upgrade( - prepare_container_image_and_helm_chart: MockerFixture, - populate_cluster_with_chartreuse_post_upgrade: MockerFixture, - mocker: MockerFixture, -) -> None: - """ - Test chartreuse considered as a blackbox, in post-install,post-upgrade configuration - """ - assert_sql_upgraded() - assert not are_pods_scaled_down() - - -def test_chartreuse_blackbox_pre_upgrade( - prepare_container_image_and_helm_chart: MockerFixture, - populate_cluster_with_chartreuse_pre_upgrade: MockerFixture, - mocker: MockerFixture, -) -> None: - """ - Test chartreuse considered as a blackbox, in pre-upgrade configuration - """ - assert_sql_upgraded() - assert not are_pods_scaled_down() diff --git a/src/chartreuse/tests/e2e_tests/test_integration.py b/src/chartreuse/tests/e2e_tests/test_integration.py deleted file mode 100644 index 757fe81..0000000 --- a/src/chartreuse/tests/e2e_tests/test_integration.py +++ /dev/null @@ -1,32 +0,0 @@ -# import chartreuse.chartreuse_migrate - -# from .conftest import ( -# TEST_NAMESPACE, -# TEST_RELEASE, -# POSTGRESQL_URL, -# SAMPLE_ALEMBIC_PATH, -# assert_sql_upgraded, -# are_pods_scaled_down, -# ) -# from ..conftest import configure_os_environ_mock - - -# def test_integration_chartreuse_upgrade(populate_cluster, mocker): -# """ -# Test chartreuse upgrade, in post-install,post-upgrade configuration -# """ -# mocker.patch("chartreuse.utils.alembic_migration_helper.ALEMBIC_DIRECTORY_PATH", SAMPLE_ALEMBIC_PATH) -# mocker.patch("wiremind_kubernetes.kubernetes_helper._get_namespace_from_kube", return_value=TEST_NAMESPACE) -# configure_os_environ_mock( -# mocker=mocker, -# additional_environment=dict( -# CHARTREUSE_ALEMBIC_URL=POSTGRESQL_URL, -# CHARTREUSE_MIGRATE_CONTAINER_IMAGE=f"docker.wiremind.fr/wiremind/devops/chartreuse:latest", -# CHARTREUSE_RELEASE_NAME=TEST_RELEASE, -# ), -# ) - -# chartreuse.chartreuse_upgrade.main() - -# assert_sql_upgraded() -# assert not are_pods_scaled_down() diff --git a/src/chartreuse/tests/unit_tests/conftest.py b/src/chartreuse/tests/unit_tests/conftest.py index 2ab9813..5562b10 100644 --- a/src/chartreuse/tests/unit_tests/conftest.py +++ b/src/chartreuse/tests/unit_tests/conftest.py @@ -11,4 +11,7 @@ def configure_chartreuse_mock(mocker: MockerFixture, is_migration_needed: bool = create=True, ) mocker.patch("chartreuse.chartreuse_upgrade.Chartreuse.upgrade") - mocker.patch("wiremind_kubernetes.kubernetes_helper._get_namespace_from_kube", return_value="foo") + mocker.patch( + "wiremind_kubernetes.kubernetes_helper._get_namespace_from_kube", + return_value="foo", + ) diff --git a/src/chartreuse/tests/unit_tests/test_alembic.py b/src/chartreuse/tests/unit_tests/test_alembic.py index f635e4e..f478a44 100644 --- a/src/chartreuse/tests/unit_tests/test_alembic.py +++ b/src/chartreuse/tests/unit_tests/test_alembic.py @@ -1,3 +1,6 @@ +import os +import tempfile + from pytest_mock.plugin import MockerFixture import chartreuse.utils @@ -8,8 +11,13 @@ def test_respect_empty_database(mocker: MockerFixture) -> None: Test that is_postgres_empty returns empty even if alembic table exists """ table_list = ["alembic_version"] - mocker.patch("chartreuse.utils.AlembicMigrationHelper._get_table_list", return_value=table_list) - alembic_migration_helper = chartreuse.utils.AlembicMigrationHelper(database_url="foo", configure=False) + mocker.patch( + "chartreuse.utils.AlembicMigrationHelper._get_table_list", + return_value=table_list, + ) + alembic_migration_helper = chartreuse.utils.AlembicMigrationHelper( + database_url="foo", alembic_section_name="test", configure=False + ) assert alembic_migration_helper.is_postgres_empty() @@ -18,9 +26,17 @@ def test_respect_not_empty_database(mocker: MockerFixture) -> None: Test that is_postgres_empty return False when tables exist """ table_list = ["alembic_version", "foobar"] - mocker.patch("chartreuse.utils.AlembicMigrationHelper._get_table_list", return_value=table_list) - mocker.patch("chartreuse.utils.AlembicMigrationHelper._get_alembic_current", return_value="123") - alembic_migration_helper = chartreuse.utils.AlembicMigrationHelper(database_url="foo", configure=False) + mocker.patch( + "chartreuse.utils.AlembicMigrationHelper._get_table_list", + return_value=table_list, + ) + mocker.patch( + "chartreuse.utils.AlembicMigrationHelper._get_alembic_current", + return_value="123", + ) + alembic_migration_helper = chartreuse.utils.AlembicMigrationHelper( + database_url="foo", alembic_section_name="test", configure=False + ) assert alembic_migration_helper.is_postgres_empty() is False @@ -36,9 +52,17 @@ def test_detect_needed_migration(mocker: MockerFixture) -> None: e1f79bafdfa2 """ table_list = ["alembic_version", "foobar"] - mocker.patch("chartreuse.utils.AlembicMigrationHelper._get_table_list", return_value=table_list) - mocker.patch("chartreuse.utils.AlembicMigrationHelper._get_alembic_current", return_value=sample_alembic_output) - alembic_migration_helper = chartreuse.utils.AlembicMigrationHelper(database_url="foo", configure=False) + mocker.patch( + "chartreuse.utils.AlembicMigrationHelper._get_table_list", + return_value=table_list, + ) + mocker.patch( + "chartreuse.utils.AlembicMigrationHelper._get_alembic_current", + return_value=sample_alembic_output, + ) + alembic_migration_helper = chartreuse.utils.AlembicMigrationHelper( + database_url="foo", alembic_section_name="test", configure=False + ) assert alembic_migration_helper.is_migration_needed @@ -55,9 +79,17 @@ def test_detect_not_needed_migration(mocker: MockerFixture) -> None: e1f79bafdfa2 (head) """ table_list = ["alembic_version", "foobar"] - mocker.patch("chartreuse.utils.AlembicMigrationHelper._get_table_list", return_value=table_list) - mocker.patch("chartreuse.utils.AlembicMigrationHelper._get_alembic_current", return_value=sample_alembic_output) - alembic_migration_helper = chartreuse.utils.AlembicMigrationHelper(database_url="foo", configure=False) + mocker.patch( + "chartreuse.utils.AlembicMigrationHelper._get_table_list", + return_value=table_list, + ) + mocker.patch( + "chartreuse.utils.AlembicMigrationHelper._get_alembic_current", + return_value=sample_alembic_output, + ) + alembic_migration_helper = chartreuse.utils.AlembicMigrationHelper( + database_url="foo", alembic_section_name="test", configure=False + ) assert alembic_migration_helper.is_migration_needed is False @@ -69,9 +101,17 @@ def test_detect_needed_migration_non_existent(mocker: MockerFixture) -> None: sample_alembic_output = """ """ table_list = ["alembic_version", "foobar"] - mocker.patch("chartreuse.utils.AlembicMigrationHelper._get_table_list", return_value=table_list) - mocker.patch("chartreuse.utils.AlembicMigrationHelper._get_alembic_current", return_value=sample_alembic_output) - alembic_migration_helper = chartreuse.utils.AlembicMigrationHelper(database_url="foo", configure=False) + mocker.patch( + "chartreuse.utils.AlembicMigrationHelper._get_table_list", + return_value=table_list, + ) + mocker.patch( + "chartreuse.utils.AlembicMigrationHelper._get_alembic_current", + return_value=sample_alembic_output, + ) + alembic_migration_helper = chartreuse.utils.AlembicMigrationHelper( + database_url="foo", alembic_section_name="test", configure=False + ) assert alembic_migration_helper.is_migration_needed @@ -83,10 +123,18 @@ def test_detect_database_is_empty(mocker: MockerFixture) -> None: sample_alembic_output = """ """ table_list = ["alembic_version"] - mocker.patch("chartreuse.utils.AlembicMigrationHelper._get_table_list", return_value=table_list) - mocker.patch("chartreuse.utils.AlembicMigrationHelper._get_alembic_current", return_value=sample_alembic_output) + mocker.patch( + "chartreuse.utils.AlembicMigrationHelper._get_table_list", + return_value=table_list, + ) + mocker.patch( + "chartreuse.utils.AlembicMigrationHelper._get_alembic_current", + return_value=sample_alembic_output, + ) - alembic_migration_helper = chartreuse.utils.AlembicMigrationHelper(database_url="foo", configure=False) + alembic_migration_helper = chartreuse.utils.AlembicMigrationHelper( + database_url="foo", alembic_section_name="test", configure=False + ) assert alembic_migration_helper.is_postgres_empty() @@ -98,10 +146,18 @@ def test_detect_database_is_not_empty(mocker: MockerFixture) -> None: sample_alembic_output = """ """ table_list = ["alembic_version", "foobar"] - mocker.patch("chartreuse.utils.AlembicMigrationHelper._get_table_list", return_value=table_list) - mocker.patch("chartreuse.utils.AlembicMigrationHelper._get_alembic_current", return_value=sample_alembic_output) + mocker.patch( + "chartreuse.utils.AlembicMigrationHelper._get_table_list", + return_value=table_list, + ) + mocker.patch( + "chartreuse.utils.AlembicMigrationHelper._get_alembic_current", + return_value=sample_alembic_output, + ) - alembic_migration_helper = chartreuse.utils.AlembicMigrationHelper(database_url="foo", configure=False) + alembic_migration_helper = chartreuse.utils.AlembicMigrationHelper( + database_url="foo", alembic_section_name="test", configure=False + ) assert alembic_migration_helper.is_postgres_empty() is False @@ -111,12 +167,15 @@ def test_additional_parameters(mocker: MockerFixture) -> None: Test that alembic additional parameters are respectedf for upgrade_db """ mocker.patch("chartreuse.utils.AlembicMigrationHelper._get_table_list", return_value="foo") - mocker.patch("chartreuse.utils.AlembicMigrationHelper._get_alembic_current", return_value="bar") + mocker.patch( + "chartreuse.utils.AlembicMigrationHelper._get_alembic_current", + return_value="bar", + ) mocked_run_command = mocker.patch("chartreuse.utils.alembic_migration_helper.run_command") alembic_migration_helper = chartreuse.utils.AlembicMigrationHelper( - database_url="foo", configure=False, additional_parameters="foo bar" + database_url="foo", alembic_section_name="test", configure=False, additional_parameters="foo bar" ) alembic_migration_helper.upgrade_db() mocked_run_command.assert_called_with("alembic -c alembic.ini foo bar upgrade head", cwd="/app/alembic") @@ -131,9 +190,214 @@ def test_additional_parameters_current(mocker: MockerFixture) -> None: mocked_run_command.return_value = ("bar", None, 0) alembic_migration_helper = chartreuse.utils.AlembicMigrationHelper( - database_url="foo", configure=False, additional_parameters="foo bar" + database_url="foo", alembic_section_name="test", configure=False, additional_parameters="foo bar" ) alembic_migration_helper._get_alembic_current() mocked_run_command.assert_called_with( "alembic -c alembic.ini foo bar current", cwd="/app/alembic", return_result=True ) + + +def test_multi_database_configuration_postgresql_section(mocker: MockerFixture) -> None: + """ + Test that multi-database configuration correctly updates PostgreSQL section in alembic.ini + """ + # Sample alembic.ini content with multiple sections + sample_alembic_ini = """[postgresql] +script_location = postgresl +sqlalchemy.url = postgresql://wiremind_owner@localhost:5432/wiremind +prepend_sys_path = .. +file_template = %%(year)d%%(month).2d%%(day).2d-%%(slug)s + +[clickhouse] +script_location = clickhouse +sqlalchemy.url = clickhouse://default@localhost:8123/wiremind +prepend_sys_path = .. +file_template = %%(year)d%%(month).2d%%(day).2d-%%(slug)s + +[loggers] +keys = root,sqlalchemy,alembic +""" + + with tempfile.TemporaryDirectory() as temp_dir: + # Create a temporary alembic.ini file + alembic_ini_path = os.path.join(temp_dir, "alembic.ini") + with open(alembic_ini_path, "w") as f: + f.write(sample_alembic_ini) + + # Test configuring PostgreSQL section + chartreuse.utils.AlembicMigrationHelper( + alembic_directory_path=temp_dir, + alembic_config_file_path="alembic.ini", + database_url="postgresql://new_user:new_pass@new_host:5432/new_db", + alembic_section_name="postgresql", + configure=True, + skip_db_checks=True, + ) + + # Read the file and verify PostgreSQL section was updated + with open(alembic_ini_path) as f: + content = f.read() + + # Verify PostgreSQL URL was updated + assert "postgresql://new_user:new_pass@new_host:5432/new_db" in content + # Verify ClickHouse URL was NOT changed + assert "clickhouse://default@localhost:8123/wiremind" in content + # Verify sections are still intact + assert "[postgresql]" in content + assert "[clickhouse]" in content + assert "[loggers]" in content + + +def test_multi_database_configuration_clickhouse_section(mocker: MockerFixture) -> None: + """ + Test that multi-database configuration correctly updates ClickHouse section in alembic.ini + """ + # Sample alembic.ini content with multiple sections + sample_alembic_ini = """[postgresql] +script_location = postgresl +sqlalchemy.url = postgresql://wiremind_owner@localhost:5432/wiremind +prepend_sys_path = .. +file_template = %%(year)d%%(month).2d%%(day).2d-%%(slug)s + +[clickhouse] +script_location = clickhouse +sqlalchemy.url = clickhouse://default@localhost:8123/wiremind +prepend_sys_path = .. +file_template = %%(year)d%%(month).2d%%(day).2d-%%(slug)s + +[loggers] +keys = root,sqlalchemy,alembic +""" + + with tempfile.TemporaryDirectory() as temp_dir: + # Create a temporary alembic.ini file + alembic_ini_path = os.path.join(temp_dir, "alembic.ini") + with open(alembic_ini_path, "w") as f: + f.write(sample_alembic_ini) + + # Test configuring ClickHouse section + chartreuse.utils.AlembicMigrationHelper( + alembic_directory_path=temp_dir, + alembic_config_file_path="alembic.ini", + database_url="clickhouse://new_user:new_pass@new_host:8123/new_db", + alembic_section_name="clickhouse", + configure=True, + skip_db_checks=True, + ) + + # Read the file and verify ClickHouse section was updated + with open(alembic_ini_path) as f: + content = f.read() + + # Verify ClickHouse URL was updated + assert "clickhouse://new_user:new_pass@new_host:8123/new_db" in content + # Verify PostgreSQL URL was NOT changed + assert "postgresql://wiremind_owner@localhost:5432/wiremind" in content + # Verify sections are still intact + assert "[postgresql]" in content + assert "[clickhouse]" in content + assert "[loggers]" in content + + +def test_multi_database_configuration_both_sections(mocker: MockerFixture) -> None: + """ + Test that multi-database configuration correctly updates both sections independently + """ + # Sample alembic.ini content with multiple sections + sample_alembic_ini = """[postgresql] +script_location = postgresl +sqlalchemy.url = postgresql://wiremind_owner@localhost:5432/wiremind +prepend_sys_path = .. +file_template = %%(year)d%%(month).2d%%(day).2d-%%(slug)s + +[clickhouse] +script_location = clickhouse +sqlalchemy.url = clickhouse://default@localhost:8123/wiremind +prepend_sys_path = .. +file_template = %%(year)d%%(month).2d%%(day).2d-%%(slug)s + +[loggers] +keys = root,sqlalchemy,alembic +""" + + with tempfile.TemporaryDirectory() as temp_dir: + # Create a temporary alembic.ini file + alembic_ini_path = os.path.join(temp_dir, "alembic.ini") + with open(alembic_ini_path, "w") as f: + f.write(sample_alembic_ini) + + # Configure PostgreSQL section first + chartreuse.utils.AlembicMigrationHelper( + alembic_directory_path=temp_dir, + alembic_config_file_path="alembic.ini", + database_url="postgresql://pg_user:pg_pass@pg_host:5432/pg_db", + alembic_section_name="postgresql", + configure=True, + skip_db_checks=True, + ) + + # Configure ClickHouse section second + chartreuse.utils.AlembicMigrationHelper( + alembic_directory_path=temp_dir, + alembic_config_file_path="alembic.ini", + database_url="clickhouse://ch_user:ch_pass@ch_host:8123/ch_db", + alembic_section_name="clickhouse", + configure=True, + skip_db_checks=True, + ) + + # Read the file and verify both sections were updated correctly + with open(alembic_ini_path) as f: + final_content = f.read() + + # Verify both URLs are now updated correctly + assert "postgresql://pg_user:pg_pass@pg_host:5432/pg_db" in final_content + assert "clickhouse://ch_user:ch_pass@ch_host:8123/ch_db" in final_content + # Verify original URLs are gone + assert "postgresql://wiremind_owner@localhost:5432/wiremind" not in final_content + assert "clickhouse://default@localhost:8123/wiremind" not in final_content + # Verify sections are still intact + assert "[postgresql]" in final_content + assert "[clickhouse]" in final_content + assert "[loggers]" in final_content + + +def test_single_database_configuration_with_alembic_section(mocker: MockerFixture) -> None: + """ + Test that single-database configuration works when alembic_section_name is specified + """ + # Sample simple alembic.ini content (single database) + sample_alembic_ini = """[alembic] +script_location = alembic +sqlalchemy.url = postgresql://old@localhost:5432/old +prepend_sys_path = .. +file_template = %%(year)d%%(month).2d%%(day).2d-%%(slug)s +""" + + with tempfile.TemporaryDirectory() as temp_dir: + # Create a temporary alembic.ini file + alembic_ini_path = os.path.join(temp_dir, "alembic.ini") + with open(alembic_ini_path, "w") as f: + f.write(sample_alembic_ini) + + # Test single database configuration with explicit section name + chartreuse.utils.AlembicMigrationHelper( + alembic_directory_path=temp_dir, + alembic_config_file_path="alembic.ini", + database_url="postgresql://new_user:new_pass@new_host:5432/new_db", + alembic_section_name="alembic", + configure=True, + skip_db_checks=True, + ) + + # Read the file and verify URL was updated + with open(alembic_ini_path) as f: + content = f.read() + + # Verify URL was updated + assert "postgresql://new_user:new_pass@new_host:5432/new_db" in content + # Verify original URL is gone + assert "postgresql://old@localhost:5432/old" not in content + # Verify section is still intact + assert "[alembic]" in content diff --git a/src/chartreuse/tests/unit_tests/test_chartreuse.py b/src/chartreuse/tests/unit_tests/test_chartreuse.py new file mode 100644 index 0000000..d30e074 --- /dev/null +++ b/src/chartreuse/tests/unit_tests/test_chartreuse.py @@ -0,0 +1,519 @@ +"""Unit tests for chartreuse main module.""" + +import logging +from unittest.mock import MagicMock + +from pytest_mock.plugin import MockerFixture + +from chartreuse.chartreuse import Chartreuse, configure_logging +from chartreuse.config_loader import DatabaseConfig + + +class TestConfigureLogging: + """Test cases for configure_logging function.""" + + def test_configure_logging(self, mocker: MockerFixture) -> None: + """Test that logging is configured correctly.""" + mock_basic_config = mocker.patch("logging.basicConfig") + + configure_logging() + + mock_basic_config.assert_called_once_with( + format="%(asctime)s %(levelname)s: %(message)s", + datefmt="%H:%M:%S", + level=logging.INFO, + ) + + +class TestChartreuse: + """Test cases for Chartreuse class.""" + + def test_chartreuse_init_with_kubernetes_helper(self, mocker: MockerFixture) -> None: + """Test Chartreuse initialization with provided kubernetes helper.""" + # Mock AlembicMigrationHelper + mock_alembic_helper = mocker.patch("chartreuse.chartreuse.AlembicMigrationHelper") + mock_alembic_instance = MagicMock() + mock_alembic_instance.is_migration_needed = True + mock_alembic_helper.return_value = mock_alembic_instance + + # Mock configure_logging + mocker.patch("chartreuse.chartreuse.configure_logging") + + # Create a mock kubernetes helper + mock_k8s_helper = MagicMock() + + # Use the new dictionary format with DatabaseConfig objects + databases_config = { + "test-db": DatabaseConfig( + dialect="postgresql", + user="user", + password="pass", + host="localhost", + port=5432, + database="db", + alembic_directory_path="/app/alembic", + alembic_config_file_path="alembic.ini", + allow_migration_for_empty_database=True, + additional_parameters="--verbose", + ) + } + + chartreuse = Chartreuse( + databases_config=databases_config, + release_name="test-release", + kubernetes_helper=mock_k8s_helper, + ) + + # Verify AlembicMigrationHelper was initialized correctly + mock_alembic_helper.assert_called_once_with( + alembic_directory_path="/app/alembic", + alembic_config_file_path="alembic.ini", + database_url="postgresql://user:pass@localhost:5432/db", + allow_migration_for_empty_database=True, + additional_parameters="--verbose -n test-db", + alembic_section_name="test-db", + ) + + # Verify kubernetes helper is set + assert chartreuse.kubernetes_helper == mock_k8s_helper + assert chartreuse.is_migration_needed is True + + def test_chartreuse_init_without_kubernetes_helper(self, mocker: MockerFixture) -> None: + """Test Chartreuse initialization without provided kubernetes helper.""" + # Mock AlembicMigrationHelper + mock_alembic_helper = mocker.patch("chartreuse.chartreuse.AlembicMigrationHelper") + mock_alembic_instance = MagicMock() + mock_alembic_instance.is_migration_needed = False + mock_alembic_helper.return_value = mock_alembic_instance + + # Mock KubernetesDeploymentManager + mock_k8s_manager = mocker.patch("wiremind_kubernetes.kubernetes_helper.KubernetesDeploymentManager") + mock_k8s_instance = MagicMock() + mock_k8s_manager.return_value = mock_k8s_instance + + # Mock configure_logging + mocker.patch("chartreuse.chartreuse.configure_logging") + + # Use the new dictionary format with DatabaseConfig objects + databases_config = { + "test-db": DatabaseConfig( + dialect="postgresql", + user="user", + password="pass", + host="localhost", + port=5432, + database="db", + alembic_directory_path="/app/alembic", + alembic_config_file_path="alembic.ini", + allow_migration_for_empty_database=False, + additional_parameters="", + ) + } + + chartreuse = Chartreuse( + databases_config=databases_config, + release_name="test-release", + ) + + # Verify KubernetesDeploymentManager was initialized + mock_k8s_manager.assert_called_once_with(use_kubeconfig=None, release_name="test-release") + + assert chartreuse.kubernetes_helper == mock_k8s_instance + assert chartreuse.is_migration_needed is False + + def test_check_migration_needed(self, mocker: MockerFixture) -> None: + """Test is_migration_needed property.""" + mock_alembic_helper = mocker.patch("chartreuse.chartreuse.AlembicMigrationHelper") + mock_alembic_instance = MagicMock() + mock_alembic_instance.is_migration_needed = True + mock_alembic_helper.return_value = mock_alembic_instance + + mocker.patch("chartreuse.chartreuse.configure_logging") + mock_k8s_helper = MagicMock() + + # Use the new dictionary format with DatabaseConfig objects + databases_config = { + "test-db": DatabaseConfig( + dialect="postgresql", + user="user", + password="pass", + host="localhost", + port=5432, + database="db", + alembic_directory_path="/app/alembic", + alembic_config_file_path="alembic.ini", + allow_migration_for_empty_database=True, + additional_parameters="", + ) + } + + chartreuse = Chartreuse( + databases_config=databases_config, + release_name="test-release", + kubernetes_helper=mock_k8s_helper, + ) + + result = chartreuse.is_migration_needed + assert result is True + + # Change the mock return value + mock_alembic_instance.is_migration_needed = False + result = chartreuse.is_migration_needed + assert result is False + + def test_upgrade_when_migration_needed(self, mocker: MockerFixture) -> None: + """Test upgrade method when migration is needed.""" + mock_alembic_helper = mocker.patch("chartreuse.chartreuse.AlembicMigrationHelper") + mock_alembic_instance = MagicMock() + mock_alembic_instance.is_migration_needed = True + mock_alembic_helper.return_value = mock_alembic_instance + + mocker.patch("chartreuse.chartreuse.configure_logging") + mock_k8s_helper = MagicMock() + + # Use the new dictionary format with DatabaseConfig objects + databases_config = { + "test-db": DatabaseConfig( + dialect="postgresql", + user="user", + password="pass", + host="localhost", + port=5432, + database="db", + alembic_directory_path="/app/alembic", + alembic_config_file_path="alembic.ini", + allow_migration_for_empty_database=True, + additional_parameters="", + ) + } + + chartreuse = Chartreuse( + databases_config=databases_config, + release_name="test-release", + kubernetes_helper=mock_k8s_helper, + ) + + chartreuse.upgrade() + + # Verify upgrade_db was called + mock_alembic_instance.upgrade_db.assert_called_once() + + def test_upgrade_when_migration_not_needed(self, mocker: MockerFixture) -> None: + """Test upgrade method when migration is not needed.""" + mock_alembic_helper = mocker.patch("chartreuse.chartreuse.AlembicMigrationHelper") + mock_alembic_instance = MagicMock() + mock_alembic_instance.is_migration_needed = False + mock_alembic_helper.return_value = mock_alembic_instance + + mocker.patch("chartreuse.chartreuse.configure_logging") + mock_k8s_helper = MagicMock() + + # Use the new dictionary format with DatabaseConfig objects + databases_config = { + "test-db": DatabaseConfig( + dialect="postgresql", + user="user", + password="pass", + host="localhost", + port=5432, + database="db", + alembic_directory_path="/app/alembic", + alembic_config_file_path="alembic.ini", + allow_migration_for_empty_database=True, + additional_parameters="", + ) + } + + chartreuse = Chartreuse( + databases_config=databases_config, + release_name="test-release", + kubernetes_helper=mock_k8s_helper, + ) + + chartreuse.upgrade() + + # Verify upgrade_db was not called + mock_alembic_instance.upgrade_db.assert_not_called() + + +class TestChartreuseMultiDatabase: + """Test cases for Chartreuse class with multiple databases.""" + + def test_multi_chartreuse_init_with_kubernetes_helper(self, mocker: MockerFixture) -> None: + """Test Chartreuse initialization with provided kubernetes helper.""" + # Mock AlembicMigrationHelper + mock_alembic_helper = mocker.patch("chartreuse.chartreuse.AlembicMigrationHelper") + + # Create mock instances for different databases + mock_main_instance = MagicMock() + mock_main_instance.is_migration_needed = True + + mock_sec_instance = MagicMock() + mock_sec_instance.is_migration_needed = False + + # Configure the mock to return different instances for different calls + mock_alembic_helper.side_effect = [mock_main_instance, mock_sec_instance] + + # Mock configure_logging + mocker.patch("chartreuse.chartreuse.configure_logging") + + # Create a mock kubernetes helper + mock_k8s_helper = MagicMock() + + databases_config = { + "main": DatabaseConfig( + dialect="postgresql", + user="user", + password="pass", + host="localhost", + port=5432, + database="maindb", + alembic_directory_path="/app/alembic/main", + alembic_config_file_path="alembic.ini", + allow_migration_for_empty_database=True, + additional_parameters="--verbose", + ), + "secondary": DatabaseConfig( + dialect="postgresql", + user="user", + password="pass", + host="localhost", + port=5433, + database="secdb", + alembic_directory_path="/app/alembic/secondary", + alembic_config_file_path="alembic.ini", + ), + } + + multi_chartreuse = Chartreuse( + databases_config=databases_config, + release_name="test-release", + kubernetes_helper=mock_k8s_helper, + ) + + # Verify AlembicMigrationHelper was called for each database + assert mock_alembic_helper.call_count == 2 + + # Verify migration helpers are stored correctly + assert "main" in multi_chartreuse.migration_helpers + assert "secondary" in multi_chartreuse.migration_helpers + assert multi_chartreuse.migration_helpers["main"] == mock_main_instance + assert multi_chartreuse.migration_helpers["secondary"] == mock_sec_instance + + # Verify kubernetes helper is set + assert multi_chartreuse.kubernetes_helper == mock_k8s_helper + + # Verify migration is needed (because main db needs migration) + assert multi_chartreuse.is_migration_needed is True + + def test_multi_chartreuse_init_without_kubernetes_helper(self, mocker: MockerFixture) -> None: + """Test MultiChartreuse initialization without provided kubernetes helper.""" + # Mock AlembicMigrationHelper + mock_alembic_helper = mocker.patch("chartreuse.chartreuse.AlembicMigrationHelper") + mock_alembic_instance = MagicMock() + mock_alembic_instance.is_migration_needed = False + mock_alembic_helper.return_value = mock_alembic_instance + + # Mock KubernetesDeploymentManager + mock_k8s_manager = mocker.patch("wiremind_kubernetes.kubernetes_helper.KubernetesDeploymentManager") + mock_k8s_instance = MagicMock() + mock_k8s_manager.return_value = mock_k8s_instance + + # Mock configure_logging + mocker.patch("chartreuse.chartreuse.configure_logging") + + databases_config = { + "test": DatabaseConfig( + dialect="postgresql", + user="user", + password="pass", + host="localhost", + port=5432, + database="db", + alembic_directory_path="/app/alembic", + alembic_config_file_path="alembic.ini", + ) + } + + multi_chartreuse = Chartreuse( + databases_config=databases_config, + release_name="test-release", + ) + + # Verify KubernetesDeploymentManager was initialized + mock_k8s_manager.assert_called_once_with(use_kubeconfig=None, release_name="test-release") + + assert multi_chartreuse.kubernetes_helper == mock_k8s_instance + + def test_check_migration_needed_with_mixed_needs(self, mocker: MockerFixture) -> None: + """Test check_migration_needed with some databases needing migration.""" + mock_alembic_helper = mocker.patch("chartreuse.chartreuse.AlembicMigrationHelper") + + # Create mock instances + mock_main_instance = MagicMock() + mock_main_instance.is_migration_needed = False + + mock_sec_instance = MagicMock() + mock_sec_instance.is_migration_needed = True + + mock_alembic_helper.side_effect = [mock_main_instance, mock_sec_instance] + + mocker.patch("chartreuse.chartreuse.configure_logging") + mock_k8s_helper = MagicMock() + + databases_config = { + "main": DatabaseConfig( + dialect="postgresql", + user="user", + password="pass", + host="localhost", + port=5432, + database="maindb", + alembic_directory_path="/app/alembic/main", + alembic_config_file_path="alembic.ini", + ), + "secondary": DatabaseConfig( + dialect="postgresql", + user="user", + password="pass", + host="localhost", + port=5433, + database="secdb", + alembic_directory_path="/app/alembic/secondary", + alembic_config_file_path="alembic.ini", + ), + } + + multi_chartreuse = Chartreuse( + databases_config=databases_config, + release_name="test-release", + kubernetes_helper=mock_k8s_helper, + ) + + # Should return True because secondary needs migration + result = multi_chartreuse.is_migration_needed + assert result is True + + def test_check_migration_needed_none_need_migration(self, mocker: MockerFixture) -> None: + """Test check_migration_needed when no databases need migration.""" + mock_alembic_helper = mocker.patch("chartreuse.chartreuse.AlembicMigrationHelper") + + mock_instance = MagicMock() + mock_instance.is_migration_needed = False + mock_alembic_helper.return_value = mock_instance + + mocker.patch("chartreuse.chartreuse.configure_logging") + mock_k8s_helper = MagicMock() + + databases_config = { + "test": DatabaseConfig( + dialect="postgresql", + user="user", + password="pass", + host="localhost", + port=5432, + database="db", + alembic_directory_path="/app/alembic", + alembic_config_file_path="alembic.ini", + ) + } + + multi_chartreuse = Chartreuse( + databases_config=databases_config, + release_name="test-release", + kubernetes_helper=mock_k8s_helper, + ) + + result = multi_chartreuse.is_migration_needed + assert result is False + + def test_upgrade_only_needed_databases(self, mocker: MockerFixture) -> None: + """Test upgrade method only upgrades databases that need migration.""" + mock_alembic_helper = mocker.patch("chartreuse.chartreuse.AlembicMigrationHelper") + + # Create mock instances + mock_main_instance = MagicMock() + mock_main_instance.is_migration_needed = True + + mock_sec_instance = MagicMock() + mock_sec_instance.is_migration_needed = False + + mock_alembic_helper.side_effect = [mock_main_instance, mock_sec_instance] + + mocker.patch("chartreuse.chartreuse.configure_logging") + mock_k8s_helper = MagicMock() + + databases_config = { + "main": DatabaseConfig( + dialect="postgresql", + user="user", + password="pass", + host="localhost", + port=5432, + database="maindb", + alembic_directory_path="/app/alembic/main", + alembic_config_file_path="alembic.ini", + ), + "secondary": DatabaseConfig( + dialect="postgresql", + user="user", + password="pass", + host="localhost", + port=5433, + database="secdb", + alembic_directory_path="/app/alembic/secondary", + alembic_config_file_path="alembic.ini", + ), + } + + multi_chartreuse = Chartreuse( + databases_config=databases_config, + release_name="test-release", + kubernetes_helper=mock_k8s_helper, + ) + + multi_chartreuse.upgrade() + + # Verify only main database was upgraded + mock_main_instance.upgrade_db.assert_called_once() + mock_sec_instance.upgrade_db.assert_not_called() + + def test_multi_chartreuse_default_values(self, mocker: MockerFixture) -> None: + """Test that default values are handled correctly for optional parameters.""" + mock_alembic_helper = mocker.patch("chartreuse.chartreuse.AlembicMigrationHelper") + mock_alembic_instance = MagicMock() + mock_alembic_instance.is_migration_needed = False + mock_alembic_helper.return_value = mock_alembic_instance + + mocker.patch("chartreuse.chartreuse.configure_logging") + mock_k8s_helper = MagicMock() + + databases_config = { + "test": DatabaseConfig( + dialect="postgresql", + user="user", + password="pass", + host="localhost", + port=5432, + database="db", + alembic_directory_path="/app/alembic", + alembic_config_file_path="alembic.ini", + # No allow_migration_for_empty_database or additional_parameters (using defaults) + ) + } + + Chartreuse( + databases_config=databases_config, + release_name="test-release", + kubernetes_helper=mock_k8s_helper, + ) + + # Verify AlembicMigrationHelper was called with defaults + mock_alembic_helper.assert_called_once_with( + alembic_directory_path="/app/alembic", + alembic_config_file_path="alembic.ini", + database_url="postgresql://user:pass@localhost:5432/db", + 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 + ) diff --git a/src/chartreuse/tests/unit_tests/test_chartreuse_upgrade.py b/src/chartreuse/tests/unit_tests/test_chartreuse_upgrade.py index 54b8d8d..2dcace9 100644 --- a/src/chartreuse/tests/unit_tests/test_chartreuse_upgrade.py +++ b/src/chartreuse/tests/unit_tests/test_chartreuse_upgrade.py @@ -2,49 +2,57 @@ from pytest_mock.plugin import MockerFixture import chartreuse.chartreuse_upgrade -from chartreuse import get_version from ..conftest import configure_os_environ_mock from .conftest import configure_chartreuse_mock -def test_chartreuse_upgrade_detected_migration_enabled_stop_pods(mocker: MockerFixture) -> None: +def test_chartreuse_upgrade_detected_migration_enabled_stop_pods( + mocker: MockerFixture, +) -> None: """ Test that chartreuse_upgrades stop pods in case of detected migration. """ configure_chartreuse_mock(mocker=mocker, is_migration_needed=True) mocked_stop_pods = mocker.patch("wiremind_kubernetes.KubernetesDeploymentManager.stop_pods") mocker.patch("wiremind_kubernetes.KubernetesDeploymentManager.start_pods") - configure_os_environ_mock(mocker=mocker, additional_environment={"HELM_CHART_VERSION": get_version()}) + mocker.patch("chartreuse.chartreuse_upgrade.get_version", return_value="5.0.0") + configure_os_environ_mock(mocker=mocker, additional_environment={"HELM_CHART_VERSION": "5.0.0"}) chartreuse.chartreuse_upgrade.main() mocked_stop_pods.assert_called() -def test_chartreuse_upgrade_detected_migration_disabled_stop_pods(mocker: MockerFixture) -> None: +def test_chartreuse_upgrade_detected_migration_disabled_stop_pods( + mocker: MockerFixture, +) -> None: """ Test that chartreuse_upgrades does not stop pods in case of detected migration but we disallow stop-pods. """ configure_chartreuse_mock(mocker=mocker, is_migration_needed=True) mocked_stop_pods = mocker.patch("wiremind_kubernetes.KubernetesDeploymentManager.stop_pods") mocker.patch("wiremind_kubernetes.KubernetesDeploymentManager.start_pods") + mocker.patch("chartreuse.chartreuse_upgrade.get_version", return_value="5.0.0") configure_os_environ_mock( mocker=mocker, - additional_environment=dict(CHARTREUSE_ENABLE_STOP_PODS="", HELM_CHART_VERSION=get_version()), + additional_environment={"CHARTREUSE_ENABLE_STOP_PODS": "", "HELM_CHART_VERSION": "5.0.0"}, ) chartreuse.chartreuse_upgrade.main() mocked_stop_pods.assert_not_called() -def test_chartreuse_upgrade_no_migration_disabled_stop_pods(mocker: MockerFixture) -> None: +def test_chartreuse_upgrade_no_migration_disabled_stop_pods( + mocker: MockerFixture, +) -> None: """ Test that chartreuse_upgrades does NOT stop pods in case of migration not needed. """ configure_chartreuse_mock(mocker=mocker, is_migration_needed=False) mocked_stop_pods = mocker.patch("wiremind_kubernetes.KubernetesDeploymentManager.stop_pods") mocker.patch("wiremind_kubernetes.KubernetesDeploymentManager.start_pods") - configure_os_environ_mock(mocker=mocker, additional_environment={"HELM_CHART_VERSION": get_version()}) + mocker.patch("chartreuse.chartreuse_upgrade.get_version", return_value="5.0.0") + configure_os_environ_mock(mocker=mocker, additional_environment={"HELM_CHART_VERSION": "5.0.0"}) chartreuse.chartreuse_upgrade.main() mocked_stop_pods.assert_not_called() @@ -72,7 +80,10 @@ def test_chartreuse_upgrade_no_migration_disabled_stop_pods(mocker: MockerFixtur ], ) def test_chartreuse_upgrade_compatibility_check( - mocker: MockerFixture, helm_chart_version: str | None, package_version: str, should_raise: bool + mocker: MockerFixture, + helm_chart_version: str | None, + package_version: str, + should_raise: bool, ) -> None: """ Test that chartreuse_upgrade deals as expected with compatibility with the package version and the diff --git a/src/chartreuse/tests/unit_tests/test_chartreuse_upgrade_extended.py b/src/chartreuse/tests/unit_tests/test_chartreuse_upgrade_extended.py new file mode 100644 index 0000000..2c82f5e --- /dev/null +++ b/src/chartreuse/tests/unit_tests/test_chartreuse_upgrade_extended.py @@ -0,0 +1,537 @@ +"""Extended unit tests for chartreuse_upgrade module.""" + +import os + +import pytest +from pytest_mock.plugin import MockerFixture + +from chartreuse.chartreuse_upgrade import ensure_safe_run, main + + +class TestEnsureSafeRun: + """Test cases for ensure_safe_run function.""" + + def test_ensure_safe_run_matching_versions(self, mocker: MockerFixture) -> None: + """Test ensure_safe_run with matching major.minor versions.""" + mocker.patch("chartreuse.chartreuse_upgrade.get_version", return_value="5.2.1") + mocker.patch.dict(os.environ, {"HELM_CHART_VERSION": "5.2.0"}) + + # Should not raise any exception + ensure_safe_run() + + def test_ensure_safe_run_mismatched_versions(self, mocker: MockerFixture) -> None: + """Test ensure_safe_run with mismatched major.minor versions.""" + mocker.patch("chartreuse.chartreuse_upgrade.get_version", return_value="5.1.3") + mocker.patch.dict(os.environ, {"HELM_CHART_VERSION": "5.2.0"}) + + with pytest.raises(ValueError) as exc_info: + ensure_safe_run() + + assert "don't have the same 'major.minor'" in str(exc_info.value) + assert "5.1.3" in str(exc_info.value) + assert "5.2.0" in str(exc_info.value) + + def test_ensure_safe_run_missing_helm_version(self, mocker: MockerFixture) -> None: + """Test ensure_safe_run with missing HELM_CHART_VERSION environment variable.""" + mocker.patch("chartreuse.chartreuse_upgrade.get_version", return_value="5.1.3") + mocker.patch.dict(os.environ, {}, clear=True) + + with pytest.raises(ValueError) as exc_info: + ensure_safe_run() + + assert "Couldn't get the Chartreuse's Helm Chart version" in str(exc_info.value) + + def test_ensure_safe_run_empty_helm_version(self, mocker: MockerFixture) -> None: + """Test ensure_safe_run with empty HELM_CHART_VERSION environment variable.""" + mocker.patch("chartreuse.chartreuse_upgrade.get_version", return_value="5.1.3") + mocker.patch.dict(os.environ, {"HELM_CHART_VERSION": ""}) + + with pytest.raises(ValueError) as exc_info: + ensure_safe_run() + + assert "Couldn't get the Chartreuse's Helm Chart version" in str(exc_info.value) + + def test_ensure_safe_run_different_major_versions(self, mocker: MockerFixture) -> None: + """Test ensure_safe_run with different major versions.""" + mocker.patch("chartreuse.chartreuse_upgrade.get_version", return_value="4.2.1") + mocker.patch.dict(os.environ, {"HELM_CHART_VERSION": "5.2.0"}) + + with pytest.raises(ValueError) as exc_info: + ensure_safe_run() + + assert "(['5', '2'] != ['4', '2'])" in str(exc_info.value) + + def test_ensure_safe_run_complex_version_formats(self, mocker: MockerFixture) -> None: + """Test ensure_safe_run with complex version formats.""" + mocker.patch("chartreuse.chartreuse_upgrade.get_version", return_value="5.2.1-alpha.1") + mocker.patch.dict(os.environ, {"HELM_CHART_VERSION": "5.2.0-beta.2"}) + + # Should not raise any exception as major.minor match + ensure_safe_run() + + +class TestMainMultiDatabase: + """Test cases for main function with multi-database configuration.""" + + def test_main_multi_database_success(self, mocker: MockerFixture) -> None: + """Test main function with successful multi-database configuration.""" + # Mock ensure_safe_run + mocker.patch("chartreuse.chartreuse_upgrade.ensure_safe_run") + + # Mock file existence for config validation + mocker.patch("os.path.exists", return_value=True) + mocker.patch("os.path.isfile", return_value=True) + mocker.patch("os.path.isfile", return_value=True) + + # Mock config loading + mock_config = [ + { + "name": "main", + "url": "postgresql://user:pass@localhost:5432/maindb", + "alembic_directory_path": "/app/alembic/main", + "alembic_config_file_path": "alembic.ini", + } + ] + mock_load_config = mocker.patch( + "chartreuse.chartreuse_upgrade.load_multi_database_config", + return_value=mock_config, + ) + + # Mock Chartreuse + mock_multi_chartreuse = mocker.patch("chartreuse.chartreuse_upgrade.Chartreuse") + mock_chartreuse_instance = mocker.MagicMock() + mock_chartreuse_instance.is_migration_needed = True + mock_multi_chartreuse.return_value = mock_chartreuse_instance + + # Mock KubernetesDeploymentManager + mock_k8s_manager = mocker.patch("chartreuse.chartreuse_upgrade.KubernetesDeploymentManager") + mock_k8s_instance = mocker.MagicMock() + mock_k8s_manager.return_value = mock_k8s_instance + + # Set up environment + mocker.patch.dict( + os.environ, + { + "CHARTREUSE_MULTI_CONFIG_PATH": "/app/config.yaml", + "CHARTREUSE_ENABLE_STOP_PODS": "true", + "CHARTREUSE_RELEASE_NAME": "test-release", + "CHARTREUSE_UPGRADE_BEFORE_DEPLOYMENT": "false", + "HELM_IS_INSTALL": "false", + }, + ) + + main() + + # Verify mocks were called correctly + mock_load_config.assert_called_once_with("/app/config.yaml") + mock_multi_chartreuse.assert_called_once() + mock_chartreuse_instance.upgrade.assert_called_once() + mock_k8s_instance.stop_pods.assert_called_once() + mock_k8s_instance.start_pods.assert_called_once() + + def test_main_multi_database_no_migration_needed(self, mocker: MockerFixture) -> None: + """Test main function when no migration is needed.""" + mocker.patch("chartreuse.chartreuse_upgrade.ensure_safe_run") + + # Mock file existence for config validation + mocker.patch("os.path.exists", return_value=True) + mocker.patch("os.path.isfile", return_value=True) + + mock_config = [ + { + "name": "main", + "url": "postgresql://user:pass@localhost:5432/maindb", + "alembic_directory_path": "/app/alembic/main", + "alembic_config_file_path": "alembic.ini", + } + ] + mocker.patch( + "chartreuse.chartreuse_upgrade.load_multi_database_config", + return_value=mock_config, + ) + + mock_multi_chartreuse = mocker.patch("chartreuse.chartreuse_upgrade.Chartreuse") + mock_chartreuse_instance = mocker.MagicMock() + mock_chartreuse_instance.is_migration_needed = False + mock_multi_chartreuse.return_value = mock_chartreuse_instance + + mock_k8s_manager = mocker.patch("chartreuse.chartreuse_upgrade.KubernetesDeploymentManager") + mock_k8s_instance = mocker.MagicMock() + mock_k8s_manager.return_value = mock_k8s_instance + + mocker.patch.dict( + os.environ, + { + "CHARTREUSE_MULTI_CONFIG_PATH": "/app/config.yaml", + "CHARTREUSE_ENABLE_STOP_PODS": "true", + "CHARTREUSE_RELEASE_NAME": "test-release", + "CHARTREUSE_UPGRADE_BEFORE_DEPLOYMENT": "false", + "HELM_IS_INSTALL": "false", + }, + ) + + main() + + # Verify no pods were stopped/started and no upgrade occurred + mock_chartreuse_instance.upgrade.assert_not_called() + mock_k8s_instance.stop_pods.assert_not_called() + mock_k8s_instance.start_pods.assert_not_called() + + def test_main_multi_database_config_load_failure(self, mocker: MockerFixture) -> None: + """Test main function with config loading failure.""" + mocker.patch("chartreuse.chartreuse_upgrade.ensure_safe_run") + + mocker.patch( + "chartreuse.chartreuse_upgrade.load_multi_database_config", + side_effect=FileNotFoundError("Config file not found"), + ) + + mocker.patch.dict( + os.environ, + { + "CHARTREUSE_MULTI_CONFIG_PATH": "/app/missing_config.yaml", + "CHARTREUSE_RELEASE_NAME": "test-release", + }, + ) + + with pytest.raises(FileNotFoundError): + main() + + def test_main_multi_database_stop_pods_disabled(self, mocker: MockerFixture) -> None: + """Test main function with ENABLE_STOP_PODS disabled.""" + mocker.patch("chartreuse.chartreuse_upgrade.ensure_safe_run") + + # Mock file existence for config validation + mocker.patch("os.path.exists", return_value=True) + mocker.patch("os.path.isfile", return_value=True) + + mock_config = [ + { + "name": "main", + "url": "postgresql://user:pass@localhost:5432/maindb", + "alembic_directory_path": "/app/alembic/main", + "alembic_config_file_path": "alembic.ini", + } + ] + mocker.patch( + "chartreuse.chartreuse_upgrade.load_multi_database_config", + return_value=mock_config, + ) + + mock_multi_chartreuse = mocker.patch("chartreuse.chartreuse_upgrade.Chartreuse") + mock_chartreuse_instance = mocker.MagicMock() + mock_chartreuse_instance.is_migration_needed = True + mock_multi_chartreuse.return_value = mock_chartreuse_instance + + mock_k8s_manager = mocker.patch("chartreuse.chartreuse_upgrade.KubernetesDeploymentManager") + mock_k8s_instance = mocker.MagicMock() + mock_k8s_manager.return_value = mock_k8s_instance + + mocker.patch.dict( + os.environ, + { + "CHARTREUSE_MULTI_CONFIG_PATH": "/app/config.yaml", + "CHARTREUSE_ENABLE_STOP_PODS": "false", + "CHARTREUSE_RELEASE_NAME": "test-release", + "CHARTREUSE_UPGRADE_BEFORE_DEPLOYMENT": "false", + "HELM_IS_INSTALL": "false", + }, + ) + + main() + + # Verify upgrade was called but pods were not stopped/started + mock_chartreuse_instance.upgrade.assert_called_once() + mock_k8s_instance.stop_pods.assert_not_called() + mock_k8s_instance.start_pods.assert_not_called() + + def test_main_multi_database_upgrade_before_deployment(self, mocker: MockerFixture) -> None: + """Test main function with UPGRADE_BEFORE_DEPLOYMENT and not HELM_IS_INSTALL.""" + mocker.patch("chartreuse.chartreuse_upgrade.ensure_safe_run") + + # Mock file existence for config validation + mocker.patch("os.path.exists", return_value=True) + mocker.patch("os.path.isfile", return_value=True) + + mock_config = [ + { + "name": "main", + "url": "postgresql://user:pass@localhost:5432/maindb", + "alembic_directory_path": "/app/alembic/main", + "alembic_config_file_path": "alembic.ini", + } + ] + mocker.patch( + "chartreuse.chartreuse_upgrade.load_multi_database_config", + return_value=mock_config, + ) + + mock_multi_chartreuse = mocker.patch("chartreuse.chartreuse_upgrade.Chartreuse") + mock_chartreuse_instance = mocker.MagicMock() + mock_chartreuse_instance.is_migration_needed = True + mock_multi_chartreuse.return_value = mock_chartreuse_instance + + mock_k8s_manager = mocker.patch("chartreuse.chartreuse_upgrade.KubernetesDeploymentManager") + mock_k8s_instance = mocker.MagicMock() + mock_k8s_manager.return_value = mock_k8s_instance + + mocker.patch.dict( + os.environ, + { + "CHARTREUSE_MULTI_CONFIG_PATH": "/app/config.yaml", + "CHARTREUSE_ENABLE_STOP_PODS": "true", + "CHARTREUSE_RELEASE_NAME": "test-release", + "CHARTREUSE_UPGRADE_BEFORE_DEPLOYMENT": "true", + "HELM_IS_INSTALL": "false", + }, + ) + + main() + + # Verify pods were stopped and upgrade was called, but start_pods was not called + mock_chartreuse_instance.upgrade.assert_called_once() + mock_k8s_instance.stop_pods.assert_called_once() + mock_k8s_instance.start_pods.assert_not_called() + + def test_main_multi_database_start_pods_failure(self, mocker: MockerFixture) -> None: + """Test main function when start_pods fails.""" + mocker.patch("chartreuse.chartreuse_upgrade.ensure_safe_run") + + # Mock file existence for config validation + mocker.patch("os.path.exists", return_value=True) + mocker.patch("os.path.isfile", return_value=True) + + mock_config = [ + { + "name": "main", + "url": "postgresql://user:pass@localhost:5432/maindb", + "alembic_directory_path": "/app/alembic/main", + "alembic_config_file_path": "alembic.ini", + } + ] + mocker.patch( + "chartreuse.chartreuse_upgrade.load_multi_database_config", + return_value=mock_config, + ) + + mock_multi_chartreuse = mocker.patch("chartreuse.chartreuse_upgrade.Chartreuse") + mock_chartreuse_instance = mocker.MagicMock() + mock_chartreuse_instance.is_migration_needed = True + mock_multi_chartreuse.return_value = mock_chartreuse_instance + + mock_k8s_manager = mocker.patch("chartreuse.chartreuse_upgrade.KubernetesDeploymentManager") + mock_k8s_instance = mocker.MagicMock() + mock_k8s_instance.start_pods.side_effect = Exception("Failed to start pods") + mock_k8s_manager.return_value = mock_k8s_instance + + # Mock logger to verify error was logged + mock_logger = mocker.patch("chartreuse.chartreuse_upgrade.logger") + + mocker.patch.dict( + os.environ, + { + "CHARTREUSE_MULTI_CONFIG_PATH": "/app/config.yaml", + "CHARTREUSE_ENABLE_STOP_PODS": "true", + "CHARTREUSE_RELEASE_NAME": "test-release", + "CHARTREUSE_UPGRADE_BEFORE_DEPLOYMENT": "false", + "HELM_IS_INSTALL": "false", + }, + ) + + # Should not raise exception, just log error + main() + + # Verify error was logged + mock_logger.error.assert_called_once() + assert "Couldn't scale up new pods" in str(mock_logger.error.call_args) + + +class TestMainSingleDatabase: + """Test cases for main function with single-database configuration.""" + + def test_main_single_database_success(self, mocker: MockerFixture) -> None: + """Test main function with successful single-database configuration.""" + mocker.patch("chartreuse.chartreuse_upgrade.ensure_safe_run") + + # Mock file existence for config validation + mocker.patch("os.path.exists", return_value=True) + mocker.patch("os.path.isfile", return_value=True) + + # Mock config loading for single database + mock_config = [ + { + "name": "main", + "url": "postgresql://user:pass@localhost:5432/db", + "alembic_directory_path": "/app/alembic", + "alembic_config_file_path": "alembic.ini", + "alembic_allow_migration_for_empty_database": True, + "alembic_additional_parameters": "--verbose", + } + ] + mocker.patch( + "chartreuse.chartreuse_upgrade.load_multi_database_config", + return_value=mock_config, + ) + + mock_chartreuse = mocker.patch("chartreuse.chartreuse_upgrade.Chartreuse") + mock_chartreuse_instance = mocker.MagicMock() + mock_chartreuse_instance.is_migration_needed = True + mock_chartreuse.return_value = mock_chartreuse_instance + + mock_k8s_manager = mocker.patch("chartreuse.chartreuse_upgrade.KubernetesDeploymentManager") + mock_k8s_instance = mocker.MagicMock() + mock_k8s_manager.return_value = mock_k8s_instance + + # Set up environment for unified mode (always requires CHARTREUSE_MULTI_CONFIG_PATH) + mocker.patch.dict( + os.environ, + { + "CHARTREUSE_MULTI_CONFIG_PATH": "/app/config.yaml", + "CHARTREUSE_ENABLE_STOP_PODS": "true", + "CHARTREUSE_RELEASE_NAME": "test-release", + "CHARTREUSE_UPGRADE_BEFORE_DEPLOYMENT": "false", + "HELM_IS_INSTALL": "false", + }, + clear=True, + ) + + main() + + # Verify Chartreuse was initialized correctly with unified architecture + mock_chartreuse.assert_called_once_with( + databases_config=mock_config, + release_name="test-release", + kubernetes_helper=mock_k8s_instance, + ) + + # Verify upgrade flow + mock_chartreuse_instance.upgrade.assert_called_once() + mock_k8s_instance.stop_pods.assert_called_once() + mock_k8s_instance.start_pods.assert_called_once() + + def test_main_single_database_default_values(self, mocker: MockerFixture) -> None: + """Test main function with default values for optional environment variables.""" + mocker.patch("chartreuse.chartreuse_upgrade.ensure_safe_run") + + # Mock file existence for config validation + mocker.patch("os.path.exists", return_value=True) + mocker.patch("os.path.isfile", return_value=True) + + # Mock config loading for single database with default values + mock_config = [ + { + "name": "main", + "url": "postgresql://user:pass@localhost:5432/db", + "alembic_directory_path": "/app/alembic", + "alembic_config_file_path": "alembic.ini", + "alembic_allow_migration_for_empty_database": False, + "alembic_additional_parameters": "", + } + ] + mocker.patch( + "chartreuse.chartreuse_upgrade.load_multi_database_config", + return_value=mock_config, + ) + + mock_chartreuse = mocker.patch("chartreuse.chartreuse_upgrade.Chartreuse") + mock_chartreuse_instance = mocker.MagicMock() + mock_chartreuse_instance.is_migration_needed = False + mock_chartreuse.return_value = mock_chartreuse_instance + + mock_k8s_manager = mocker.patch("chartreuse.chartreuse_upgrade.KubernetesDeploymentManager") + mock_k8s_instance = mocker.MagicMock() + mock_k8s_manager.return_value = mock_k8s_instance + + # Set minimal required environment variables for unified mode + mocker.patch.dict( + os.environ, + { + "CHARTREUSE_MULTI_CONFIG_PATH": "/app/config.yaml", + "CHARTREUSE_ENABLE_STOP_PODS": "false", + "CHARTREUSE_RELEASE_NAME": "test-release", + "CHARTREUSE_UPGRADE_BEFORE_DEPLOYMENT": "false", + "HELM_IS_INSTALL": "false", + }, + clear=True, + ) + + main() + + # Verify Chartreuse was initialized correctly with unified architecture + mock_chartreuse.assert_called_once_with( + databases_config=mock_config, + release_name="test-release", + kubernetes_helper=mock_k8s_instance, + ) + + +class TestMainBooleanParsing: + """Test cases for boolean environment variable parsing.""" + + @pytest.mark.parametrize( + "bool_str,expected", + [ + ("true", True), # .lower() not in ("", "false", "0") -> "true" not in -> True + ("TRUE", True), # .lower() not in ("", "false", "0") -> "true" not in -> True + ("True", True), # .lower() not in ("", "false", "0") -> "true" not in -> True + ("false", False), # .lower() not in ("", "false", "0") -> "false" in -> False + ("FALSE", False), # .lower() not in ("", "false", "0") -> "false" in -> False + ("False", False), # .lower() not in ("", "false", "0") -> "false" in -> False + ("", False), # .lower() not in ("", "false", "0") -> "" in -> False + ("0", False), # .lower() not in ("", "false", "0") -> "0" in -> False + ("1", True), # .lower() not in ("", "false", "0") -> "1" not in -> True + ], + ) + def test_boolean_parsing_variations(self, mocker: MockerFixture, bool_str: str, expected: bool) -> None: + """Test various boolean string parsing in environment variables.""" + mocker.patch("chartreuse.chartreuse_upgrade.ensure_safe_run") + + # Mock file existence for config validation + mocker.patch("os.path.exists", return_value=True) + mocker.patch("os.path.isfile", return_value=True) + + # Mock config loading + mock_config = [ + { + "name": "main", + "url": "postgresql://user:pass@localhost:5432/db", + "alembic_directory_path": "/app/alembic", + "alembic_config_file_path": "alembic.ini", + } + ] + mocker.patch( + "chartreuse.chartreuse_upgrade.load_multi_database_config", + return_value=mock_config, + ) + + mock_chartreuse = mocker.patch("chartreuse.chartreuse_upgrade.Chartreuse") + mock_chartreuse_instance = mocker.MagicMock() + mock_chartreuse_instance.is_migration_needed = True # Set to True to test stop_pods behavior + mock_chartreuse.return_value = mock_chartreuse_instance + + mock_k8s_manager = mocker.patch("chartreuse.chartreuse_upgrade.KubernetesDeploymentManager") + mock_k8s_instance = mocker.MagicMock() + mock_k8s_manager.return_value = mock_k8s_instance + + # Test boolean parsing for environment variables that are still parsed in main() + # The boolean parsing logic: .lower() not in ("", "false", "0") + # which means: empty string, "false", or "0" are False, everything else is True + mocker.patch.dict( + os.environ, + { + "CHARTREUSE_MULTI_CONFIG_PATH": "/app/config.yaml", + "CHARTREUSE_ENABLE_STOP_PODS": bool_str, # Test this boolean parsing + "CHARTREUSE_RELEASE_NAME": "test-release", + "CHARTREUSE_UPGRADE_BEFORE_DEPLOYMENT": "false", + "HELM_IS_INSTALL": "false", + }, + clear=True, + ) + + main() + + # Verify that stop_pods was called based on the boolean parsing result + if expected: + mock_k8s_instance.stop_pods.assert_called_once() + else: + mock_k8s_instance.stop_pods.assert_not_called() diff --git a/src/chartreuse/tests/unit_tests/test_config_loader.py b/src/chartreuse/tests/unit_tests/test_config_loader.py new file mode 100644 index 0000000..17472c3 --- /dev/null +++ b/src/chartreuse/tests/unit_tests/test_config_loader.py @@ -0,0 +1,520 @@ +""" +Unit tests for config_loader module. +""" + +import tempfile +from pathlib import Path +from unittest.mock import patch + +import pytest +import yaml +from pydantic import ValidationError + +from chartreuse.config_loader import DatabaseConfig, MultiDatabaseConfig, load_multi_database_config + + +class TestDatabaseConfig: + """Test cases for DatabaseConfig model.""" + + def test_valid_database_config(self): + """Test creating a valid database configuration.""" + config_data = { + "dialect": "postgresql", + "user": "testuser", + "password": "testpass", + "host": "localhost", + "port": 5432, + "database": "testdb", + "alembic_directory_path": "/app/alembic", + } + + config = DatabaseConfig(**config_data) + + assert config.dialect == "postgresql" + assert config.user == "testuser" + assert config.password == "testpass" + assert config.host == "localhost" + assert config.port == 5432 + assert config.database == "testdb" + assert config.alembic_directory_path == "/app/alembic" + assert config.alembic_config_file_path == "alembic.ini" # default value + assert config.allow_migration_for_empty_database is True # default value + assert config.additional_parameters == "" # default value + + def test_database_config_url_computation(self): + """Test that the URL is computed correctly from components.""" + config = DatabaseConfig( + dialect="postgresql", + user="myuser", + password="mypass", + host="example.com", + port=5432, + database="mydb", + alembic_directory_path="/app/alembic", + ) + + expected_url = "postgresql://myuser:mypass@example.com:5432/mydb" + assert config.url == expected_url + + def test_database_config_with_optional_fields(self): + """Test database configuration with all optional fields set.""" + config_data = { + "dialect": "mysql", + "user": "testuser", + "password": "testpass", + "host": "localhost", + "port": 3306, + "database": "testdb", + "alembic_directory_path": "/app/alembic", + "alembic_config_file_path": "custom_alembic.ini", + "allow_migration_for_empty_database": False, + "additional_parameters": "-x custom_param=value", + } + + config = DatabaseConfig(**config_data) + + assert config.alembic_config_file_path == "custom_alembic.ini" + assert config.allow_migration_for_empty_database is False + assert config.additional_parameters == "-x custom_param=value" + + def test_database_config_missing_required_fields(self): + """Test that missing required fields raise validation errors.""" + incomplete_config = { + "dialect": "postgresql", + "user": "testuser", + # Missing password, host, port, database, alembic_directory_path + } + + with pytest.raises(ValidationError) as excinfo: + DatabaseConfig(**incomplete_config) + + error_dict = excinfo.value.errors() + required_fields = ["password", "host", "port", "database", "alembic_directory_path"] + + for field in required_fields: + assert any(error["loc"] == (field,) for error in error_dict) + + def test_database_config_invalid_port(self): + """Test validation of port field.""" + # Test port too small + with pytest.raises(ValidationError) as excinfo: + DatabaseConfig( + dialect="postgresql", + user="testuser", + password="testpass", + host="localhost", + port=0, # Invalid: too small + database="testdb", + alembic_directory_path="/app/alembic", + ) + assert "greater than 0" in str(excinfo.value) + + # Test port too large + with pytest.raises(ValidationError) as excinfo: + DatabaseConfig( + dialect="postgresql", + user="testuser", + password="testpass", + host="localhost", + port=65536, # Invalid: too large + database="testdb", + alembic_directory_path="/app/alembic", + ) + assert "less than or equal to 65535" in str(excinfo.value) + + def test_database_config_dialect_validation(self): + """Test dialect validation with supported and unsupported dialects.""" + # Test supported dialect + config = DatabaseConfig( + dialect="postgresql", + user="testuser", + password="testpass", + host="localhost", + port=5432, + database="testdb", + alembic_directory_path="/app/alembic", + ) + assert config.dialect == "postgresql" + + # Test unsupported dialect (should still work but log warning) + config = DatabaseConfig( + dialect="unsupported_db", + user="testuser", + password="testpass", + host="localhost", + port=5432, + database="testdb", + alembic_directory_path="/app/alembic", + ) + assert config.dialect == "unsupported_db" + + def test_additional_parameters_empty_string_handling(self): + """Test that additional_parameters defaults to empty string when not provided.""" + config = DatabaseConfig( + dialect="postgresql", + user="testuser", + password="testpass", + host="localhost", + port=5432, + database="testdb", + alembic_directory_path="/app/alembic", + ) + assert config.additional_parameters == "" + + +class TestMultiDatabaseConfig: + """Test cases for MultiDatabaseConfig model.""" + + def test_valid_multi_database_config(self): + """Test creating a valid multi-database configuration.""" + db_config = DatabaseConfig( + dialect="postgresql", + user="testuser", + password="testpass", + host="localhost", + port=5432, + database="testdb", + alembic_directory_path="/app/alembic", + ) + + config = MultiDatabaseConfig(databases={"main": db_config}) + + assert "main" in config.databases + assert config.databases["main"] == db_config + + def test_multi_database_config_multiple_databases(self): + """Test multi-database configuration with multiple databases.""" + db1 = DatabaseConfig( + dialect="postgresql", + user="user1", + password="pass1", + host="host1", + port=5432, + database="db1", + alembic_directory_path="/app/alembic1", + ) + + db2 = DatabaseConfig( + dialect="mysql", + user="user2", + password="pass2", + host="host2", + port=3306, + database="db2", + alembic_directory_path="/app/alembic2", + ) + + config = MultiDatabaseConfig(databases={"main": db1, "analytics": db2}) + + assert len(config.databases) == 2 + assert "main" in config.databases + assert "analytics" in config.databases + assert config.databases["main"] == db1 + assert config.databases["analytics"] == db2 + + def test_multi_database_config_empty_databases(self): + """Test that empty databases dictionary raises validation error.""" + with pytest.raises(ValidationError) as excinfo: + MultiDatabaseConfig(databases={}) + + assert "At least one database must be configured" in str(excinfo.value) + + def test_multi_database_config_missing_databases(self): + """Test that missing databases field raises validation error.""" + with pytest.raises(ValidationError) as excinfo: + MultiDatabaseConfig() + + error_dict = excinfo.value.errors() + assert any(error["loc"] == ("databases",) for error in error_dict) + + +class TestLoadMultiDatabaseConfig: + """Test cases for load_multi_database_config function.""" + + def test_load_valid_config_file(self): + """Test loading a valid configuration file.""" + config_content = { + "databases": { + "main": { + "dialect": "postgresql", + "user": "testuser", + "password": "testpass", + "host": "localhost", + "port": 5432, + "database": "testdb", + "alembic_directory_path": "/app/alembic", + }, + "analytics": { + "dialect": "mysql", + "user": "analyticsuser", + "password": "analyticspass", + "host": "analytics-host", + "port": 3306, + "database": "analytics", + "alembic_directory_path": "/app/alembic/analytics", + "allow_migration_for_empty_database": False, + }, + } + } + + with tempfile.NamedTemporaryFile(mode="w", suffix=".yaml", delete=False) as f: + yaml.dump(config_content, f) + config_path = f.name + + try: + databases = load_multi_database_config(config_path) + + assert len(databases) == 2 + assert "main" in databases + assert "analytics" in databases + + # Check main database + main_db = databases["main"] + assert main_db.dialect == "postgresql" + assert main_db.user == "testuser" + assert main_db.password == "testpass" + assert main_db.host == "localhost" + assert main_db.port == 5432 + assert main_db.database == "testdb" + assert main_db.alembic_directory_path == "/app/alembic" + assert main_db.allow_migration_for_empty_database is True # default + assert main_db.url == "postgresql://testuser:testpass@localhost:5432/testdb" + + # Check analytics database + analytics_db = databases["analytics"] + assert analytics_db.dialect == "mysql" + assert analytics_db.user == "analyticsuser" + assert analytics_db.allow_migration_for_empty_database is False + assert analytics_db.url == "mysql://analyticsuser:analyticspass@analytics-host:3306/analytics" + + finally: + Path(config_path).unlink() + + def test_load_single_database_config(self): + """Test loading a configuration with a single database.""" + config_content = { + "databases": { + "main": { + "dialect": "postgresql", + "user": "testuser", + "password": "testpass", + "host": "localhost", + "port": 5432, + "database": "testdb", + "alembic_directory_path": "/app/alembic", + } + } + } + + with tempfile.NamedTemporaryFile(mode="w", suffix=".yaml", delete=False) as f: + yaml.dump(config_content, f) + config_path = f.name + + try: + databases = load_multi_database_config(config_path) + + assert len(databases) == 1 + assert "main" in databases + main_db = databases["main"] + assert main_db.dialect == "postgresql" + + finally: + Path(config_path).unlink() + + def test_load_nonexistent_file(self): + """Test loading a configuration file that doesn't exist.""" + with pytest.raises(FileNotFoundError) as excinfo: + load_multi_database_config("/nonexistent/config.yaml") + + assert "Configuration file not found" in str(excinfo.value) + + def test_load_invalid_yaml(self): + """Test loading a file with invalid YAML.""" + invalid_yaml = "databases:\n main:\n dialect: postgresql\n invalid: yaml: content:" + + with tempfile.NamedTemporaryFile(mode="w", suffix=".yaml", delete=False) as f: + f.write(invalid_yaml) + config_path = f.name + + try: + with pytest.raises(ValueError) as excinfo: + load_multi_database_config(config_path) + + assert "Invalid YAML in configuration file" in str(excinfo.value) + + finally: + Path(config_path).unlink() + + def test_load_config_missing_required_fields(self): + """Test loading a configuration with missing required fields.""" + config_content = { + "databases": { + "main": { + "dialect": "postgresql", + "user": "testuser", + # Missing required fields: password, host, port, database, alembic_directory_path + } + } + } + + with tempfile.NamedTemporaryFile(mode="w", suffix=".yaml", delete=False) as f: + yaml.dump(config_content, f) + config_path = f.name + + try: + with pytest.raises(ValueError) as excinfo: + load_multi_database_config(config_path) + + assert "Invalid configuration" in str(excinfo.value) + + finally: + Path(config_path).unlink() + + def test_load_config_empty_databases(self): + """Test loading a configuration with empty databases.""" + config_content = {"databases": {}} + + with tempfile.NamedTemporaryFile(mode="w", suffix=".yaml", delete=False) as f: + yaml.dump(config_content, f) + config_path = f.name + + try: + with pytest.raises(ValueError) as excinfo: + load_multi_database_config(config_path) + + assert "Invalid configuration" in str(excinfo.value) + + finally: + Path(config_path).unlink() + + def test_load_config_no_databases_key(self): + """Test loading a configuration without 'databases' key.""" + config_content = {"invalid_key": "value"} + + with tempfile.NamedTemporaryFile(mode="w", suffix=".yaml", delete=False) as f: + yaml.dump(config_content, f) + config_path = f.name + + try: + with pytest.raises(ValueError) as excinfo: + load_multi_database_config(config_path) + + assert "Invalid configuration" in str(excinfo.value) + + finally: + Path(config_path).unlink() + + def test_load_config_with_all_optional_fields(self): + """Test loading a configuration with all optional fields specified.""" + config_content = { + "databases": { + "main": { + "dialect": "postgresql", + "user": "testuser", + "password": "testpass", + "host": "localhost", + "port": 5432, + "database": "testdb", + "alembic_directory_path": "/app/alembic", + "alembic_config_file_path": "custom_alembic.ini", + "allow_migration_for_empty_database": False, + "additional_parameters": "-x custom=value", + } + } + } + + with tempfile.NamedTemporaryFile(mode="w", suffix=".yaml", delete=False) as f: + yaml.dump(config_content, f) + config_path = f.name + + try: + databases = load_multi_database_config(config_path) + + assert len(databases) == 1 + db = databases["main"] + assert db.alembic_config_file_path == "custom_alembic.ini" + assert db.allow_migration_for_empty_database is False + assert db.additional_parameters == "-x custom=value" + + finally: + Path(config_path).unlink() + + @patch("chartreuse.config_loader.logger") + def test_load_config_logs_info(self, mock_logger): + """Test that loading a valid config logs appropriate info messages.""" + config_content = { + "databases": { + "main": { + "dialect": "postgresql", + "user": "testuser", + "password": "testpass", + "host": "localhost", + "port": 5432, + "database": "testdb", + "alembic_directory_path": "/app/alembic", + }, + "analytics": { + "dialect": "mysql", + "user": "analyticsuser", + "password": "analyticspass", + "host": "analytics-host", + "port": 3306, + "database": "analytics", + "alembic_directory_path": "/app/alembic/analytics", + }, + } + } + + with tempfile.NamedTemporaryFile(mode="w", suffix=".yaml", delete=False) as f: + yaml.dump(config_content, f) + config_path = f.name + + try: + load_multi_database_config(config_path) + + # Check that info was logged about loaded databases + mock_logger.info.assert_called_once() + log_call = mock_logger.info.call_args[0] + # The logger uses % formatting, so check format string and args separately + assert "Loaded configuration for %d database(s): %s" == log_call[0] + assert log_call[1] == 2 + # Check that both database names are present (order may vary) + database_names = log_call[2] + assert "main" in database_names + assert "analytics" in database_names + + finally: + Path(config_path).unlink() + + @patch("chartreuse.config_loader.logger") + def test_unsupported_dialect_logs_warning(self, mock_logger): + """Test that using an unsupported dialect logs a warning.""" + config_content = { + "databases": { + "main": { + "dialect": "unsupported_dialect", + "user": "testuser", + "password": "testpass", + "host": "localhost", + "port": 5432, + "database": "testdb", + "alembic_directory_path": "/app/alembic", + } + } + } + + with tempfile.NamedTemporaryFile(mode="w", suffix=".yaml", delete=False) as f: + yaml.dump(config_content, f) + config_path = f.name + + try: + load_multi_database_config(config_path) + + # Check that warning was logged about unsupported dialect + mock_logger.warning.assert_called_once() + log_call = mock_logger.warning.call_args[0] + assert "unsupported_dialect" in log_call[1] + assert "might not be supported" in log_call[0] + + finally: + Path(config_path).unlink() diff --git a/src/chartreuse/utils/alembic_migration_helper.py b/src/chartreuse/utils/alembic_migration_helper.py index fdb8eed..27c8c31 100755 --- a/src/chartreuse/utils/alembic_migration_helper.py +++ b/src/chartreuse/utils/alembic_migration_helper.py @@ -1,6 +1,7 @@ import logging import os import re +from configparser import ConfigParser from subprocess import SubprocessError from time import sleep, time @@ -19,9 +20,12 @@ def __init__( alembic_directory_path: str = "/app/alembic", alembic_config_file_path: str = "alembic.ini", database_url: str, + alembic_section_name: str, additional_parameters: str = "", allow_migration_for_empty_database: bool = False, configure: bool = True, + # skip_db_checks is used for testing purposes only + skip_db_checks: bool = False, ): if not database_url: raise ValueError("database_url not set, not upgrading database.") @@ -31,30 +35,42 @@ def __init__( self.additional_parameters = additional_parameters self.alembic_directory_path = alembic_directory_path self.alembic_config_file_path = alembic_config_file_path + self.alembic_section_name = alembic_section_name + self.skip_db_checks = skip_db_checks # 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: + if self.is_patroni_postgresql and not skip_db_checks: self.additional_parameters += " -x patroni_postgresql=yes" self._wait_postgres_is_configured() if configure: self._configure() - - self.is_migration_needed = self._check_migration_needed() + # skip_db_checks is used for testing purposes only + if not skip_db_checks: + self.is_migration_needed = self._check_migration_needed() + else: + self.is_migration_needed = False def _configure(self) -> None: - with open(f"{self.alembic_directory_path}/{self.alembic_config_file_path}") as f: - content = f.read() - content_new = re.sub( - "(sqlalchemy.url.*=.*){1}", r"sqlalchemy.url=%s" % self.database_url, content, flags=re.M - ) - if content != content_new: - with open(f"{self.alembic_directory_path}/{self.alembic_config_file_path}", "w") as f: - f.write(content_new) - logger.info("alembic.ini was configured.") - else: - raise SubprocessError("configuration of alembic.ini has failed (alembic.ini is unchanged)") + config_path = f"{self.alembic_directory_path}/{self.alembic_config_file_path}" + + # Read the configuration file + config = ConfigParser() + config.read(config_path) + + # Multi-database configuration: update specific section + if not config.has_section(self.alembic_section_name): + raise ValueError(f"Section '{self.alembic_section_name}' not found in {self.alembic_config_file_path}") + + # Update the sqlalchemy.url in the specific section + config.set(self.alembic_section_name, "sqlalchemy.url", self.database_url) + + # Write the updated configuration back to the file + with open(config_path, "w") as f: + config.write(f) + + logger.info("alembic.ini was configured for section %s.", self.alembic_section_name) def _wait_postgres_is_configured(self) -> None: """ @@ -91,7 +107,7 @@ def _wait_postgres_is_configured(self) -> None: return except Exception as e: # TODO: Learn about exceptions that should be caught here, otherwise we'll wait for nothing - logger.info(f"Caught: {e}") + logger.info("Caught: %s", e) logger.info( "Waiting for the postgres-operator to create the user wiremind_owner_user" " (that alembic and I use) and to set default privileges..." @@ -107,7 +123,7 @@ def _get_table_list(self) -> list[str]: def is_postgres_empty(self) -> bool: table_list = self._get_table_list() - logger.info(f"Tables in the database: {table_list}") + logger.info("Tables in the database: %s", table_list) # Don't count "alembic" table table_name = "alembic_version" if table_name in table_list: