Update Azure's schema changes during 2025.2.1 upgrade (#2965)

Co-authored-by: Chuck McAndrew <[email protected]>

Author: viniciusdc

src/_nebari/stages/terraform_state/__init__.py

@@ -219,7 +219,6 @@ def deploy(
         # terraform show command, inside check_immutable_fields
         with super().deploy(stage_outputs, disable_prompt, tofu_init=False):
             env_mapping = {}
-
             with modified_environ(**env_mapping):
                 yield
 
@@ -262,6 +261,7 @@ def check_immutable_fields(self):
 
     def get_nebari_config_state(self) -> dict:
         directory = str(self.output_directory / self.stage_prefix)
+
         tf_state = opentofu.show(directory)
         nebari_config_state = None
 

src/_nebari/upgrade.py

@@ -34,6 +34,7 @@
     get_k8s_version_prefix,
     get_provider_config_block_name,
     load_yaml,
+    update_tfstate_file,
     yaml,
 )
 from _nebari.version import __version__, rounded_ver_parse
@@ -1649,8 +1650,61 @@ def _version_specific_upgrade(
             https://www.nebari.dev/docs/how-tos/kubernetes-version-upgrade
             """
         )
-
         rich.print(text)
+
+        # If the Nebari provider is Azure, we must handle a major version upgrade
+        # of the Azure Terraform provider (from 3.x to 4.x). This involves schema changes
+        # that can cause validation issues. The following steps will attempt to migrate
+        # your state file automatically. For details, see:
+        # https://github.com/nebari-dev/nebari/issues/2964
+
+        if config.get("provider", "") == "azure":
+            rich.print("\n ⚠️ Azure Provider Upgrade Notice ⚠️")
+            rich.print(
+                textwrap.dedent(
+                    """
+                    In this Nebari release, the Azure Terraform provider has been upgraded
+                    from version 3.97.1 to 4.7.0. This major update includes internal schema
+                    changes for certain resources, most notably the `azurerm_storage_account`.
+
+                    Nebari will attempt to update your Terraform state automatically to
+                    accommodate these changes. However, if you skip this automatic migration,
+                    you may encounter validation errors during redeployment.
+
+                    For detailed information on the Azure provider 4.x changes, please visit:
+                    https://registry.terraform.io/providers/hashicorp/azurerm/latest/docs/guides/4.0-upgrade-guide
+                    """
+                )
+            )
+
+            # Prompt user for confirmation
+            continue_ = kwargs.get("attempt_fixes", False) or Confirm.ask(
+                "Nebari can automatically apply the necessary state migrations. Continue?",
+                default=False,
+            )
+
+            if not continue_:
+                rich.print(
+                    "You have chosen to skip the automatic state migration. This may lead "
+                    "to validation errors during deployment.\n\nFor instructions on manually "
+                    "updating your Terraform state, please refer to:\n"
+                    "https://github.com/nebari-dev/nebari/issues/2964"
+                )
+                exit
+            else:
+                # In this case the full path in the tfstate file is
+                # resources.instances.attributes.enable_https_traffic_only
+                MIGRATION_STATE = {
+                    "enable_https_traffic_only": "https_traffic_only_enabled"
+                }
+                state_filepath = (
+                    config_filename.parent
+                    / "stages/01-terraform-state/azure/terraform.tfstate"
+                )
+
+                # Perform the state file update
+                update_tfstate_file(state_filepath, MIGRATION_STATE)
+
         rich.print("Ready to upgrade to Nebari version [green]2025.2.1[/green].")
 
         return config

src/_nebari/utils.py

@@ -15,6 +15,7 @@
 from pathlib import Path
 from typing import Any, Dict, List, Set
 
+import rich
 from ruamel.yaml import YAML
 
 from _nebari import constants
@@ -472,3 +473,49 @@ def modified(self):
 
     def __repr__(self):
         return f"{self.__class__.__name__}(diff={json.dumps(self.diff)})"
+
+
+def update_tfstate_file(state_filepath: Path, migration_map: dict) -> None:
+    """
+    Updates a Terraform state file by replacing deprecated attributes with their new
+    counterparts.
+
+    Originally introduced in the Nebari `2025.2.1` release to accommodate major schema
+    changes from Terraform cloud providers, this function can be extended for future
+    migrations or patches. By centralizing the replacement logic, it ensures a clean,
+    modular design that keeps upgrade steps concise.
+
+    :param state_filepath: A Path object pointing to the Terraform state file.
+    :param migration_map: A dictionary where keys are old attribute paths and values are new attribute paths.
+    """
+    if not state_filepath.exists():
+        rich.print(
+            f"[red]No Terraform state file found at {state_filepath}. Skipping migration.[/red]"
+        )
+        return
+
+    try:
+        with open(state_filepath, "r") as f:
+            state = json.load(f)
+    except json.JSONDecodeError:
+        rich.print(
+            f"[red]Invalid JSON structure in {state_filepath}. Skipping migration.[/red]"
+        )
+        return
+
+    # Traverse the resources → instances → attributes hierarchy
+    # and apply the specified attribute replacements.
+    for resource in state.get("resources", []):
+        for instance in resource.get("instances", []):
+            attributes = instance.get("attributes", {})
+            for old_attr, new_attr in migration_map.items():
+                if old_attr in attributes:
+                    attributes[new_attr] = attributes.pop(old_attr)
+
+    # Save the modified state back to disk
+    with open(state_filepath, "w") as f:
+        json.dump(state, f, indent=2)
+
+    rich.print(
+        f" ✅ [green]Successfully updated the Terraform state file: {state_filepath}[/green]"
+    )