|
| 1 | +# Copyright Materialize, Inc. and contributors. All rights reserved. |
| 2 | +# |
| 3 | +# Use of this software is governed by the Business Source License |
| 4 | +# included in the LICENSE file at the root of this repository. |
| 5 | +# |
| 6 | +# As of the Change Date specified in that file, in accordance with |
| 7 | +# the Business Source License, use of this software will be governed |
| 8 | +# by the Apache License, Version 2.0. |
| 9 | + |
| 10 | +from collections.abc import Callable |
| 11 | +from dataclasses import dataclass |
| 12 | +from random import Random |
| 13 | +from typing import TypedDict |
| 14 | + |
| 15 | +from materialize.mz_version import MzVersion |
| 16 | +from materialize.mzcompose import get_default_system_parameters |
| 17 | +from materialize.mzcompose.composition import Composition |
| 18 | +from materialize.mzcompose.services.materialized import DeploymentStatus, Materialized |
| 19 | +from materialize.mzcompose.services.postgres import METADATA_STORE |
| 20 | + |
| 21 | + |
| 22 | +class MaterializedUpgradeArgs(TypedDict): |
| 23 | + """Arguments for the Materialized service constructor required for 0dt upgrades.""" |
| 24 | + |
| 25 | + name: str |
| 26 | + image: str | None |
| 27 | + deploy_generation: int |
| 28 | + system_parameter_defaults: dict[str, str] |
| 29 | + external_metadata_store: bool |
| 30 | + restart: str |
| 31 | + |
| 32 | + |
| 33 | +@dataclass |
| 34 | +class UpgradeStep: |
| 35 | + """Represents a single upgrade step with its service name and action.""" |
| 36 | + |
| 37 | + new_service: Materialized |
| 38 | + previous_service: Materialized |
| 39 | + upgrade: Callable[[], None] |
| 40 | + |
| 41 | + |
| 42 | +def generate_materialized_upgrade_args( |
| 43 | + versions: list[MzVersion | None], |
| 44 | +) -> list[MaterializedUpgradeArgs]: |
| 45 | + """ |
| 46 | + Constructs a list of required Materialized arguments for 0dt upgrades. |
| 47 | + Requires there to be an mz_1 and mz_2 service already in the composition. |
| 48 | + """ |
| 49 | + # We use the first version to get the system parameters since the defaults for |
| 50 | + # newer versions include cutting edge features than can break backwards compatibility. |
| 51 | + # TODO (multiversion1): Get minimal system parameters by default to avoid cutting edge features. |
| 52 | + system_parameter_defaults = get_default_system_parameters(versions[0]) |
| 53 | + |
| 54 | + return [ |
| 55 | + MaterializedUpgradeArgs( |
| 56 | + image=f"materialize/materialized:{version}" if version else None, |
| 57 | + # Cycle through mz_1 and mz_2 for upgrades since spinning up services have a cost. |
| 58 | + name=f"mz_{(i % 2) + 1}", |
| 59 | + # Generation number for the service. Required to start services in read only mode. |
| 60 | + deploy_generation=i, |
| 61 | + system_parameter_defaults=system_parameter_defaults, |
| 62 | + # To share the same metadata store between services |
| 63 | + external_metadata_store=True, |
| 64 | + # To restart when container exits due to promotion |
| 65 | + restart="on-failure", |
| 66 | + ) |
| 67 | + for i, version in enumerate(versions) |
| 68 | + ] |
| 69 | + |
| 70 | + |
| 71 | +def generate_random_upgrade_path( |
| 72 | + versions: list[MzVersion], |
| 73 | + rng: Random | None = None, |
| 74 | +) -> list[MzVersion]: |
| 75 | + """ |
| 76 | + Generates a random upgrade path between the given versions. |
| 77 | + """ |
| 78 | + selected_versions = [] |
| 79 | + |
| 80 | + rng = rng or Random() |
| 81 | + # For each version in the input list, randomly select it with a 50% chance. |
| 82 | + for v in versions: |
| 83 | + if rng.random() < 0.5: |
| 84 | + selected_versions.append(v) |
| 85 | + |
| 86 | + # Always include at least one version to avoid empty paths. |
| 87 | + if len(selected_versions) == 0: |
| 88 | + selected_versions.append(rng.choice(versions)) |
| 89 | + |
| 90 | + return selected_versions |
| 91 | + |
| 92 | + |
| 93 | +class Materialized0dtUpgrader: |
| 94 | + """ |
| 95 | + Manages a sequence of Materialized service upgrades using zero-downtime deployments. |
| 96 | +
|
| 97 | + Args: |
| 98 | + materialized_services: List of Materialized instances representing each upgrade step |
| 99 | + """ |
| 100 | + |
| 101 | + def __init__(self, c: Composition, materialized_services: list[Materialized]): |
| 102 | + self.materialized_services = materialized_services |
| 103 | + self.c = c |
| 104 | + |
| 105 | + def create_upgrade_steps_list(self) -> list[UpgradeStep]: |
| 106 | + """ |
| 107 | + Returns a list of upgrade step actions from the second service onward. |
| 108 | +
|
| 109 | + Each step is a closure that, when called, will perform |
| 110 | + the upgrade step to the corresponding service. |
| 111 | + """ |
| 112 | + |
| 113 | + def create_upgrade_action( |
| 114 | + current_service: Materialized, |
| 115 | + previous_service: Materialized, |
| 116 | + ): |
| 117 | + def upgrade() -> None: |
| 118 | + with self.c.override(current_service): |
| 119 | + current_service_image = ( |
| 120 | + current_service.config.get("image") or "current" |
| 121 | + ) |
| 122 | + previous_service_image = previous_service.config.get("image") |
| 123 | + |
| 124 | + print(f"Bringing up {current_service_image}") |
| 125 | + self.c.up(current_service.name) |
| 126 | + print(f"Awaiting promotion of {current_service_image}") |
| 127 | + self.c.await_mz_deployment_status( |
| 128 | + DeploymentStatus.READY_TO_PROMOTE, current_service.name |
| 129 | + ) |
| 130 | + self.c.promote_mz(current_service.name) |
| 131 | + print(f"Awaiting leader status of {current_service_image}") |
| 132 | + self.c.await_mz_deployment_status( |
| 133 | + DeploymentStatus.IS_LEADER, current_service.name |
| 134 | + ) |
| 135 | + |
| 136 | + print(f"Killing {previous_service_image}") |
| 137 | + self.c.kill(previous_service.name, wait=True) |
| 138 | + |
| 139 | + return upgrade |
| 140 | + |
| 141 | + services = self.materialized_services |
| 142 | + steps = [] |
| 143 | + for idx in range(1, len(services)): |
| 144 | + current_service = services[idx] |
| 145 | + previous_service = services[idx - 1] |
| 146 | + |
| 147 | + steps.append( |
| 148 | + UpgradeStep( |
| 149 | + new_service=current_service, |
| 150 | + previous_service=previous_service, |
| 151 | + upgrade=create_upgrade_action(current_service, previous_service), |
| 152 | + ) |
| 153 | + ) |
| 154 | + return steps |
| 155 | + |
| 156 | + def initialize(self) -> tuple[Materialized, list[UpgradeStep]]: |
| 157 | + """ |
| 158 | + Initialize the with the first service. Returns a list where |
| 159 | + each step is a closure that, when called, will perform the upgrade step to the corresponding service. |
| 160 | + """ |
| 161 | + first_service = self.materialized_services[0] |
| 162 | + with self.c.override(first_service): |
| 163 | + print(f"Bringing up {first_service.name}") |
| 164 | + self.c.up(first_service.name) |
| 165 | + |
| 166 | + return first_service, self.create_upgrade_steps_list() |
| 167 | + |
| 168 | + def print_upgrade_path(self) -> None: |
| 169 | + """ |
| 170 | + Print the upgrade steps. |
| 171 | + """ |
| 172 | + |
| 173 | + def image_to_string(image: str | None) -> str: |
| 174 | + return "current" if image is None else image.split(":")[-1] |
| 175 | + |
| 176 | + print( |
| 177 | + f"Upgrade path: {str.join(' -> ', [image_to_string(service.config.get('image')) for service in self.materialized_services])}" |
| 178 | + ) |
| 179 | + |
| 180 | + def cleanup(self) -> None: |
| 181 | + """ |
| 182 | + Cleanup after upgrade. |
| 183 | + """ |
| 184 | + print("Cleaning up upgrade path") |
| 185 | + # Ensure all services are killed and removed |
| 186 | + self.c.kill( |
| 187 | + *[service.name for service in self.materialized_services], wait=True |
| 188 | + ) |
| 189 | + self.c.rm( |
| 190 | + *[service.name for service in self.materialized_services], |
| 191 | + destroy_volumes=True, |
| 192 | + ) |
| 193 | + self.c.kill(METADATA_STORE, wait=True) |
| 194 | + self.c.rm( |
| 195 | + METADATA_STORE, |
| 196 | + destroy_volumes=True, |
| 197 | + ) |
0 commit comments