qontract-reconcile 0.10.2.dev394__py3-none-any.whl → 0.10.2.dev427__py3-none-any.whl

reconcile/acs_rbac.py

@@ -65,7 +65,7 @@ class AcsRole(BaseModel):
     assignments: list[AssignmentPair]
     permission_set_name: str
     access_scope: AcsAccessScope
-    system_default: bool | None
+    system_default: bool | None = None
 
     @classmethod
     def build(cls, permission: Permission, usernames: list[str]) -> Self:
@@ -151,7 +151,7 @@ class AcsRbacIntegration(QontractReconcileIntegration[NoParams]):
                 for permission in role.oidc_permissions or []:
                     if isinstance(permission, OidcPermissionAcsV1):
                         permission_usernames[
-                            Permission(**permission.dict(by_alias=True))
+                            Permission(**permission.model_dump(by_alias=True))
                         ].append(user.org_username)
         return list(starmap(AcsRole.build, permission_usernames.items()))
 

reconcile/aus/advanced_upgrade_service.py

@@ -1,13 +1,13 @@
 import logging
 from collections import defaultdict
 from datetime import timedelta
-from typing import Optional
+from typing import Annotated, Optional
 
 from pydantic import (
     BaseModel,
     Field,
     ValidationError,
-    validator,
+    field_validator,
 )
 from pydantic.dataclasses import dataclass
 
@@ -75,8 +75,10 @@ from reconcile.utils.ocm_base_client import (
     OCMBaseClient,
     init_ocm_base_client,
 )
+from reconcile.utils.semver_helper import make_semver
 
 QONTRACT_INTEGRATION = "advanced-upgrade-scheduler"
+QONTRACT_INTEGRATION_VERSION = make_semver(0, 1, 0)
 
 
 class AdvancedUpgradeServiceIntegration(OCMClusterUpgradeSchedulerOrgIntegration):
@@ -289,13 +291,16 @@ class OrganizationLabelSet(BaseModel):
 
     blocked_versions: CSV | None = Field(alias=aus_label_key("blocked-versions"))
 
-    sector_max_parallel_upgrades: dict[str, str] = labelset_groupfield(
-        group_prefix=aus_label_key("sector-max-parallel-upgrades.")
-    )
+    sector_max_parallel_upgrades: Annotated[
+        dict[str, str],
+        labelset_groupfield(
+            group_prefix=aus_label_key("sector-max-parallel-upgrades.")
+        ),
+    ]
 
-    sector_deps: dict[str, CSV] = labelset_groupfield(
-        group_prefix=aus_label_key("sector-deps.")
-    )
+    sector_deps: Annotated[
+        dict[str, CSV], labelset_groupfield(group_prefix=aus_label_key("sector-deps."))
+    ]
     """
     Each sector with dependencies is represented as a `sector-deps.<sector-name>` label
     with a CSV formatted list of dependant sectors. The custom `labelset_groupfield``
@@ -357,7 +362,7 @@ def _build_org_upgrade_spec(
         ]
 
     org_labelset = build_labelset(org_labels, OrganizationLabelSet)
-    final_org = org.copy(deep=True)
+    final_org = org.model_copy(deep=True)
     final_org.blocked_versions = org_labelset.blocked_versions  # type: ignore
     final_org.sectors = org_labelset.sector_dependencies()
     final_org.inherit_version_data = inherit_version_data
@@ -375,6 +380,7 @@ def _build_org_upgrade_spec(
                     org=org_upgrade_spec.org,
                     upgradePolicy=upgrade_policy,
                     cluster=c.ocm_cluster,
+                    cluster_labels=c.labels,
                     health=cluster_health,
                     nodePools=node_pool_specs_by_cluster_id.get(c.ocm_cluster.id) or [],
                 )
@@ -411,7 +417,7 @@ class ClusterUpgradePolicyLabelSet(BaseModel):
     """
 
     soak_days: int = Field(alias=aus_label_key("soak-days"), ge=0)
-    workloads: CSV = Field(alias=aus_label_key("workloads"), csv_min_items=1)
+    workloads: CSV = Field(alias=aus_label_key("workloads"), min_length=1)
     schedule: str = Field(alias=aus_label_key("schedule"))
     mutexes: CSV | None = Field(alias=aus_label_key("mutexes"))
     sector: str | None = Field(alias=aus_label_key("sector"))
@@ -419,14 +425,14 @@ class ClusterUpgradePolicyLabelSet(BaseModel):
     version_gate_approvals: CSV | None = Field(
         alias=aus_label_key("version-gate-approvals")
     )
-    _schedule_validator = validator("schedule", allow_reuse=True)(cron_validator)
+    _schedule_validator = field_validator("schedule")(cron_validator)
 
     def build_labels_dict(self) -> dict[str, str]:
         """
         Build a dictionary of all labels in this labelset.
         """
         labels = {}
-        for k, v in self.dict(by_alias=True).items():
+        for k, v in self.model_dump(by_alias=True).items():
             if v is None:
                 continue
             if isinstance(v, list):

reconcile/aus/base.py

@@ -16,7 +16,7 @@ from typing import (
 )
 
 from croniter import croniter
-from pydantic import BaseModel, Extra
+from pydantic import BaseModel
 from requests.exceptions import HTTPError
 from semver import VersionInfo
 
@@ -46,7 +46,7 @@ from reconcile.aus.models import (
     OrganizationUpgradeSpec,
     Sector,
 )
-from reconcile.aus.version_gates import HANDLERS
+from reconcile.aus.version_gates import HANDLERS, sts_version_gate_handler
 from reconcile.gql_definitions.advanced_upgrade_service.aus_organization import (
     query as aus_organizations_query,
 )
@@ -79,7 +79,9 @@ from reconcile.utils.datetime_util import (
 from reconcile.utils.defer import defer
 from reconcile.utils.disabled_integrations import integration_is_enabled
 from reconcile.utils.filtering import remove_none_values_from_dict
+from reconcile.utils.jobcontroller.controller import build_job_controller
 from reconcile.utils.ocm.addons import AddonService, AddonServiceV1, AddonServiceV2
+from reconcile.utils.ocm.base import LabelContainer
 from reconcile.utils.ocm.clusters import (
     OCMCluster,
 )
@@ -102,6 +104,7 @@ from reconcile.utils.runtime.integration import (
     PydanticRunParams,
     QontractReconcileIntegration,
 )
+from reconcile.utils.secret_reader import SecretReaderBase
 from reconcile.utils.semver_helper import (
     get_version_prefix,
     parse_semver,
@@ -110,6 +113,18 @@ from reconcile.utils.semver_helper import (
 from reconcile.utils.state import init_state
 
 MIN_DELTA_MINUTES = 6
+STS_GATE_LABEL = "api.openshift.com/gate-sts"
+AUS_VERSION_GATE_APPROVALS_LABEL = "sre-capabilities.aus.version-gate-approvals"
+
+
+class RosaRoleUpgradeHandlerParams(PydanticRunParams):
+    job_controller_cluster: str
+    job_controller_namespace: str
+    rosa_job_service_account: str
+    rosa_role: str
+    rosa_job_image: str | None = None
+    integration_name: str
+    integration_version: str
 
 
 class AdvancedUpgradeSchedulerBaseIntegrationParams(PydanticRunParams):
@@ -117,6 +132,7 @@ class AdvancedUpgradeSchedulerBaseIntegrationParams(PydanticRunParams):
     ocm_organization_ids: set[str] | None = None
     excluded_ocm_organization_ids: set[str] | None = None
     ignore_sts_clusters: bool = False
+    rosa_role_upgrade_handler_params: RosaRoleUpgradeHandlerParams | None = None
 
 
 class ReconcileError(Exception):
@@ -404,15 +420,20 @@ class AbstractUpgradePolicy(ABC, BaseModel):
 
     cluster: OCMCluster
 
-    id: str | None
-    next_run: str | None
-    schedule: str | None
+    id: str | None = None
+    next_run: str | None = None
+    schedule: str | None = None
     schedule_type: str
     version: str
-    state: str | None
+    state: str | None = None
 
     @abstractmethod
-    def create(self, ocm_api: OCMBaseClient) -> None:
+    def create(
+        self,
+        ocm_api: OCMBaseClient,
+        rosa_role_upgrade_handler_params: RosaRoleUpgradeHandlerParams | None = None,
+        secret_reader: SecretReaderBase | None = None,
+    ) -> None:
         pass
 
     @abstractmethod
@@ -430,7 +451,7 @@ def addon_upgrade_policy_soonest_next_run() -> str:
     return to_utc_seconds_iso_format(next_run)
 
 
-class AddonUpgradePolicy(AbstractUpgradePolicy):
+class AddonUpgradePolicy(AbstractUpgradePolicy, arbitrary_types_allowed=True):
     """Class to create and delete Addon upgrade policies in OCM"""
 
     addon_id: str
@@ -439,7 +460,12 @@ class AddonUpgradePolicy(AbstractUpgradePolicy):
     class Config:
         arbitrary_types_allowed = True
 
-    def create(self, ocm_api: OCMBaseClient) -> None:
+    def create(
+        self,
+        ocm_api: OCMBaseClient,
+        rosa_role_upgrade_handler_params: RosaRoleUpgradeHandlerParams | None = None,
+        secret_reader: SecretReaderBase | None = None,
+    ) -> None:
         self.addon_service.create_addon_upgrade_policy(
             ocm_api=ocm_api,
             cluster_id=self.cluster.id,
@@ -472,14 +498,63 @@ class AddonUpgradePolicy(AbstractUpgradePolicy):
 class ClusterUpgradePolicy(AbstractUpgradePolicy):
     """Class to create ClusterUpgradePolicies in OCM"""
 
-    def create(self, ocm_api: OCMBaseClient) -> None:
+    organization_id: str
+    cluster_labels: LabelContainer
+
+    def create(
+        self,
+        ocm_api: OCMBaseClient,
+        rosa_role_upgrade_handler_params: RosaRoleUpgradeHandlerParams | None = None,
+        secret_reader: SecretReaderBase | None = None,
+    ) -> None:
         policy = {
             "version": self.version,
             "schedule_type": "manual",
             "next_run": self.next_run,
         }
+        if (
+            rosa_role_upgrade_handler_params
+            and secret_reader
+            and self.should_upgrade_roles()
+        ):
+            logging.info(f"Updating account and operator roles for {self.cluster.name}")
+            sts_gate_handler = sts_version_gate_handler.STSGateHandler(
+                job_controller=build_job_controller(
+                    integration=rosa_role_upgrade_handler_params.integration_name,
+                    integration_version=rosa_role_upgrade_handler_params.integration_version,
+                    cluster=rosa_role_upgrade_handler_params.job_controller_cluster,
+                    namespace=rosa_role_upgrade_handler_params.job_controller_namespace,
+                    secret_reader=secret_reader,
+                    dry_run=False,
+                ),
+                aws_iam_role=rosa_role_upgrade_handler_params.rosa_role,
+                rosa_job_service_account=rosa_role_upgrade_handler_params.rosa_job_service_account,
+                rosa_job_image=rosa_role_upgrade_handler_params.rosa_job_image,
+            )
+            if not sts_gate_handler.upgrade_rosa_roles_v2(
+                ocm_api=ocm_api,
+                cluster=self.cluster,
+                dry_run=False,
+                upgrade_version=self.version,
+                ocm_org_id=self.organization_id,
+            ):
+                logging.error(
+                    f"Failed to update account and operator roles for {self.cluster.name}"
+                )
         create_upgrade_policy(ocm_api, self.cluster.id, policy)
 
+    def should_upgrade_roles(self) -> bool:
+        handler_csv = self.cluster_labels.get_label_value(
+            AUS_VERSION_GATE_APPROVALS_LABEL
+        )
+        if not handler_csv:
+            return False
+        return (
+            self.cluster.is_sts()
+            and self.cluster.is_rosa_classic()
+            and STS_GATE_LABEL in set(handler_csv.split(","))
+        )
+
     def delete(self, ocm_api: OCMBaseClient) -> None:
         raise NotImplementedError("ClusterUpgradePolicy.delete() not implemented")
 
@@ -497,7 +572,12 @@ class ClusterUpgradePolicy(AbstractUpgradePolicy):
 class ControlPlaneUpgradePolicy(AbstractUpgradePolicy):
     """Class to create and delete ControlPlanUpgradePolicies in OCM"""
 
-    def create(self, ocm_api: OCMBaseClient) -> None:
+    def create(
+        self,
+        ocm_api: OCMBaseClient,
+        rosa_role_upgrade_handler_params: RosaRoleUpgradeHandlerParams | None = None,
+        secret_reader: SecretReaderBase | None = None,
+    ) -> None:
         policy = {
             "version": self.version,
             "schedule_type": "manual",
@@ -521,10 +601,16 @@ class ControlPlaneUpgradePolicy(AbstractUpgradePolicy):
 
 
 class NodePoolUpgradePolicy(AbstractUpgradePolicy):
-    node_pool: str
     """Class to create NodePoolUpgradePolicies in OCM"""
 
-    def create(self, ocm_api: OCMBaseClient) -> None:
+    node_pool: str
+
+    def create(
+        self,
+        ocm_api: OCMBaseClient,
+        rosa_role_upgrade_handler_params: RosaRoleUpgradeHandlerParams | None = None,
+        secret_reader: SecretReaderBase | None = None,
+    ) -> None:
         policy = {
             "version": self.version,
             "schedule_type": "manual",
@@ -550,13 +636,19 @@ class NodePoolUpgradePolicy(AbstractUpgradePolicy):
         return f"node pool upgrade policy - {remove_none_values_from_dict(details)}"
 
 
-class UpgradePolicyHandler(BaseModel, extra=Extra.forbid):
+class UpgradePolicyHandler(BaseModel, extra="forbid"):
     """Class to handle upgrade policy actions"""
 
     action: str
     policy: AbstractUpgradePolicy
 
-    def act(self, dry_run: bool, ocm_api: OCMBaseClient) -> None:
+    def act(
+        self,
+        dry_run: bool,
+        ocm_api: OCMBaseClient,
+        rosa_role_upgrade_handler_params: RosaRoleUpgradeHandlerParams | None = None,
+        secret_reader: SecretReaderBase | None = None,
+    ) -> None:
         logging.info(f"{self.action} {self.policy.summarize()}")
         if dry_run:
             return
@@ -566,7 +658,7 @@ class UpgradePolicyHandler(BaseModel, extra=Extra.forbid):
         elif self.action == "delete":
             self.policy.delete(ocm_api)
         elif self.action == "create":
-            self.policy.create(ocm_api)
+            self.policy.create(ocm_api, rosa_role_upgrade_handler_params, secret_reader)
 
 
 def fetch_current_state(
@@ -584,6 +676,7 @@ def fetch_current_state(
             )
             current_state.extend(
                 AddonUpgradePolicy(
+                    organization_id=spec.org.org_id,
                     id=addon_upgrade_policy.id,
                     addon_id=addon_spec.addon.addon.id,
                     cluster=spec.cluster,
@@ -620,6 +713,8 @@ def fetch_current_state(
             for upgrade_policy in upgrade_policies:
                 policy = upgrade_policy | {
                     "cluster": spec.cluster,
+                    "organization_id": spec.org.org_id,
+                    "cluster_labels": spec.cluster_labels,
                 }
                 current_state.append(ClusterUpgradePolicy(**policy))
 
@@ -1018,6 +1113,8 @@ def _create_upgrade_policy(
         )
     return ClusterUpgradePolicy(
         cluster=spec.cluster,
+        organization_id=spec.org.org_id,
+        cluster_labels=spec.cluster_labels,
         version=version,
         schedule_type="manual",
         next_run=next_schedule,
@@ -1125,11 +1222,11 @@ def calculate_diff(
                         action="create",
                         policy=AddonUpgradePolicy(
                             action="create",
+                            organization_id=spec.org.org_id,
                             cluster=spec.cluster,
                             version=version,
                             schedule_type="manual",
                             addon_id=addon_id,
-                            upgrade_type="ADDON",
                             addon_service=addon_service,
                         ),
                     )
@@ -1190,6 +1287,8 @@ def act(
     dry_run: bool,
     diffs: list[UpgradePolicyHandler],
     ocm_api: OCMBaseClient,
+    rosa_role_upgrade_handler_params: RosaRoleUpgradeHandlerParams | None = None,
+    secret_reader: SecretReaderBase | None = None,
     addon_id: str | None = None,
 ) -> None:
     diffs.sort(key=sort_diffs)
@@ -1202,7 +1301,7 @@ def act(
         ):
             continue
         try:
-            diff.act(dry_run, ocm_api)
+            diff.act(dry_run, ocm_api, rosa_role_upgrade_handler_params, secret_reader)
         except HTTPError as e:
             logging.error(f"{policy.cluster.name}: {e}: {e.response.text}")
 

reconcile/aus/cluster_version_data.py

@@ -1,4 +1,3 @@
-import json
 from collections.abc import Iterable
 from datetime import datetime
 from typing import (
@@ -20,6 +19,17 @@ class WorkloadHistory(BaseModel):
     soak_days: float = 0.0
     reporting: list[str] = Field(default_factory=list)
 
+    def __eq__(self, value: Any) -> bool:
+        if isinstance(value, WorkloadHistory):
+            return super().__eq__(value)
+
+        if isinstance(value, dict):
+            return self.soak_days == value.get("soak_days", 0.0) and sorted(
+                self.reporting
+            ) == sorted(value.get("reporting", []))
+
+        return False
+
 
 class VersionHistory(BaseModel):
     workloads: dict[str, WorkloadHistory] = Field(default_factory=dict)
@@ -38,7 +48,7 @@ class Stats(BaseModel):
 
     min_version: str
     min_version_per_workload: dict[str, str] = Field(default_factory=dict)
-    inherited: Optional["Stats"]
+    inherited: Optional["Stats"] = None
 
     def inherit(self, added: "Stats") -> None:
         """adds the provided stats to our inherited data
@@ -93,12 +103,12 @@ class VersionData(BaseModel):
     upgrade policies.
     """
 
-    check_in: datetime | None
+    check_in: datetime | None = None
     versions: dict[str, VersionHistory] = Field(default_factory=dict)
-    stats: Stats | None
+    stats: Stats | None = None
 
     def jsondict(self) -> dict[str, Any]:
-        return json.loads(self.json(exclude_none=True))
+        return self.model_dump(mode="json", exclude_none=True)
 
     def save(self, state: State, ocm_name: str) -> None:
         state.add(ocm_name, self.jsondict(), force=True)

reconcile/aus/models.py

@@ -14,6 +14,7 @@ from reconcile.aus.healthchecks import AUSClusterHealth
 from reconcile.gql_definitions.fragments.aus_organization import AUSOCMOrganization
 from reconcile.gql_definitions.fragments.upgrade_policy import ClusterUpgradePolicyV1
 from reconcile.utils.ocm.addons import OCMAddonInstallation
+from reconcile.utils.ocm.base import LabelContainer
 from reconcile.utils.ocm.clusters import OCMCluster
 from reconcile.utils.semver_helper import parse_semver
 
@@ -33,6 +34,7 @@ class ClusterUpgradeSpec(BaseModel):
 
     org: AUSOCMOrganization
     cluster: OCMCluster
+    cluster_labels: LabelContainer | None = None
     upgrade_policy: ClusterUpgradePolicyV1 = Field(..., alias="upgradePolicy")
     health: AUSClusterHealth
     node_pools: list[NodePoolSpec] = Field(default_factory=list, alias="nodePools")
@@ -232,7 +234,7 @@ class SectorConfigError(Exception):
 
 class Sector(BaseModel):
     name: str
-    max_parallel_upgrades: str | None
+    max_parallel_upgrades: str | None = None
     dependencies: list[Sector] = Field(default_factory=list)
     _specs: dict[str, ClusterUpgradeSpec] = PrivateAttr(default_factory=dict)
 

reconcile/aus/ocm_addons_upgrade_scheduler_org.py

@@ -260,6 +260,7 @@ def calculate_diff(
                 aus.UpgradePolicyHandler(
                     action="delete",
                     policy=aus.AddonUpgradePolicy(
+                        organization_id=org_upgrade_spec.org.org_id,
                         cluster=current.cluster,
                         version=current.schedule_type,
                         id=current.id,

reconcile/aus/ocm_upgrade_scheduler.py

@@ -82,7 +82,14 @@ class OCMClusterUpgradeSchedulerIntegration(
                 version_data,
                 integration=self.name,
             )
-            aus.act(dry_run, diffs, ocm_api)
+
+            aus.act(
+                dry_run,
+                diffs,
+                ocm_api,
+                self.params.rosa_role_upgrade_handler_params,
+                self.secret_reader,
+            )
 
     def expose_version_data_metrics(
         self,

reconcile/aus/ocm_upgrade_scheduler_org.py

@@ -1,4 +1,5 @@
 from collections import defaultdict
+from dataclasses import dataclass
 
 from reconcile.aus.healthchecks import (
     AUSClusterHealthCheckProvider,
@@ -13,6 +14,7 @@ from reconcile.aus.node_pool_spec import get_node_pool_specs_by_org_cluster
 from reconcile.aus.ocm_upgrade_scheduler import OCMClusterUpgradeSchedulerIntegration
 from reconcile.gql_definitions.fragments.aus_organization import AUSOCMOrganization
 from reconcile.gql_definitions.fragments.ocm_environment import OCMEnvironment
+from reconcile.utils.ocm.base import LabelContainer
 from reconcile.utils.ocm.clusters import (
     OCMCluster,
     discover_clusters_for_organizations,
@@ -22,6 +24,12 @@ from reconcile.utils.ocm_base_client import init_ocm_base_client
 QONTRACT_INTEGRATION = "ocm-upgrade-scheduler-org"
 
 
+@dataclass
+class ClusterUpgradeSpecWithLabels:
+    cluster: OCMCluster
+    cluster_labels: LabelContainer
+
+
 class OCMClusterUpgradeSchedulerOrgIntegration(OCMClusterUpgradeSchedulerIntegration):
     @property
     def name(self) -> str:
@@ -60,7 +68,10 @@ class OCMClusterUpgradeSchedulerOrgIntegration(OCMClusterUpgradeSchedulerIntegra
                     specs=self._build_cluster_upgrade_specs(
                         org=org,
                         clusters_by_name={
-                            c.ocm_cluster.name: c.ocm_cluster
+                            c.ocm_cluster.name: ClusterUpgradeSpecWithLabels(
+                                cluster=c.ocm_cluster,
+                                cluster_labels=c.labels,
+                            )
                             for c in clusters_by_org[org.org_id]
                         },
                         cluster_health_provider=build_cluster_health_providers_for_organization(
@@ -79,7 +90,7 @@ class OCMClusterUpgradeSchedulerOrgIntegration(OCMClusterUpgradeSchedulerIntegra
     def _build_cluster_upgrade_specs(
         self,
         org: AUSOCMOrganization,
-        clusters_by_name: dict[str, OCMCluster],
+        clusters_by_name: dict[str, ClusterUpgradeSpecWithLabels],
         cluster_health_provider: AUSClusterHealthCheckProvider,
         node_pool_specs_by_cluster_id: dict[str, list[NodePoolSpec]],
     ) -> list[ClusterUpgradeSpec]:
@@ -87,12 +98,16 @@ class OCMClusterUpgradeSchedulerOrgIntegration(OCMClusterUpgradeSchedulerIntegra
             ClusterUpgradeSpec(
                 org=org,
                 upgradePolicy=cluster.upgrade_policy,
-                cluster=clusters_by_name[cluster.name],
+                cluster=clusters_by_name[cluster.name].cluster,
+                cluster_labels=clusters_by_name[cluster.name].cluster_labels,
                 health=cluster_health_provider.cluster_health(
-                    cluster_external_id=clusters_by_name[cluster.name].external_id,
+                    cluster_external_id=clusters_by_name[
+                        cluster.name
+                    ].cluster.external_id,
                     org_id=org.org_id,
                 ),
-                nodePools=node_pool_specs_by_cluster_id.get(ocm_cluster.id) or [],
+                nodePools=node_pool_specs_by_cluster_id.get(ocm_cluster.cluster.id)
+                or [],
             )
             for cluster in org.upgrade_policy_clusters or []
             # clusters that are not in the UUID dict will be ignored because

reconcile/aus/version_gates/sts_version_gate_handler.py

@@ -6,6 +6,7 @@ from reconcile.utils.ocm.base import OCMCluster, OCMVersionGate
 from reconcile.utils.ocm_base_client import OCMBaseClient
 from reconcile.utils.rosa.rosa_cli import RosaCliError
 from reconcile.utils.rosa.session import RosaSession
+from reconcile.utils.semver_helper import get_version_prefix
 
 GATE_LABEL = "api.openshift.com/gate-sts"
 
@@ -63,6 +64,24 @@ class STSGateHandler(GateHandler):
             )
             return False
 
+        return self.upgrade_rosa_roles(
+            cluster=cluster,
+            version_raw_id_prefix=gate.version_raw_id_prefix,
+            dry_run=dry_run,
+            ocm_api=ocm_api,
+            ocm_org_id=ocm_org_id,
+        )
+
+    def upgrade_rosa_roles(
+        self,
+        cluster: OCMCluster,
+        version_raw_id_prefix: str,
+        dry_run: bool,
+        ocm_api: OCMBaseClient,
+        ocm_org_id: str,
+    ) -> bool:
+        if not cluster.aws:
+            return False
         rosa = RosaSession(
             aws_account_id=cluster.aws.aws_account_id,
             aws_region=cluster.region.id,
@@ -83,7 +102,7 @@ class STSGateHandler(GateHandler):
                 )
             rosa.upgrade_account_roles(
                 role_prefix=account_role_prefix,
-                minor_version=gate.version_raw_id_prefix,
+                minor_version=version_raw_id_prefix,
                 channel_group=cluster.version.channel_group,
                 dry_run=dry_run,
             )
@@ -98,3 +117,37 @@ class STSGateHandler(GateHandler):
             e.write_logs_to_logger(logging.error)
             return False
         return True
+
+    def upgrade_rosa_roles_v2(
+        self,
+        cluster: OCMCluster,
+        upgrade_version: str,
+        dry_run: bool,
+        ocm_api: OCMBaseClient,
+        ocm_org_id: str,
+    ) -> bool:
+        if not cluster.aws:
+            return False
+        rosa = RosaSession(
+            aws_account_id=cluster.aws.aws_account_id,
+            aws_region=cluster.region.id,
+            aws_iam_role=self.aws_iam_role,
+            ocm_org_id=ocm_org_id,
+            ocm_api=ocm_api,
+            job_controller=self.job_controller,
+            image=self.rosa_job_image,
+            service_account=self.rosa_job_service_account,
+        )
+        policy_version = get_version_prefix(upgrade_version)
+        try:
+            rosa.upgrade_rosa_roles(
+                cluster_name=cluster.name,
+                upgrade_version=upgrade_version,
+                policy_version=policy_version,
+                dry_run=dry_run,
+            )
+        except RosaCliError as e:
+            logging.error(f"Failed to upgrade roles for cluster {cluster.name}: {e}")
+            e.write_logs_to_logger(logging.error)
+            return False
+        return True