Disconnect pylibjuju model connectors

pylibjuju model connections throw JujuRedirectExceptions
when the default controller is changed. The errors did not
stop further bootstrap procedure but they will be printed
on the console. To avoid these printouts, gracefully
close the connections. This should be done before
bootstrapping a new controller which changes the default
controller.

pylibjuju provides 2 types of connections. One for the
controller and second one at the model level.
controller.get_model() creates a new connection. Access to
any model entities like Application, Unit will be handled
by the connection at the model level.
Maintain a list of model connectors when new connections
are opened and a disconnect method to disconnect the
populated list of model connections. Additionally
disconnect the controller connection in disconnect
method.

This is not an issue if the connections are opened with
same default controller through out the process lifetime.

sunbeam-python/sunbeam/core/juju.py

@@ -251,6 +251,14 @@ class JujuHelper:
 
     def __init__(self, controller: Controller):
         self.controller = controller
+        self.model_connectors: list[Model] = []
+
+    async def disconnect(self):
+        """Disconnect all connections to juju controller."""
+        LOG.debug("Disconnecting juju controller and model connections")
+        for model in self.model_connectors:
+            await model.disconnect()
+        await self.controller.disconnect()
 
     async def get_clouds(self) -> dict:
         """Return clouds available on controller."""
@@ -268,7 +276,9 @@ async def get_model(self, model: str) -> Model:
         :model: Name of the model
         """
         try:
-            return await self.controller.get_model(model)
+            model_impl = await self.controller.get_model(model)
+            self.model_connectors.append(model_impl)
+            return model_impl
         except Exception as e:
             if "HTTP 400" in str(e) or "HTTP 404" in str(e):
                 raise ModelNotFoundException(f"Model {model!r} not found")
@@ -292,9 +302,11 @@ async def add_model(
         old_home = os.environ["HOME"]
         os.environ["HOME"] = os.environ["SNAP_REAL_HOME"]
         try:
-            return await self.controller.add_model(
+            model_impl = await self.controller.add_model(
                 model, cloud_name=cloud, credential_name=credential, config=config
             )
+            self.model_connectors.append(model_impl)
+            return model_impl
         finally:
             os.environ["HOME"] = old_home
 

sunbeam-python/sunbeam/provider/local/commands.py

@@ -216,11 +216,10 @@ def deployment_type(self) -> Tuple[str, Type[Deployment]]:
 
 
 def get_sunbeam_machine_plans(
-    deployment: Deployment, manifest: Manifest
+    deployment: Deployment, jhelper: JujuHelper, manifest: Manifest
 ) -> list[BaseStep]:
     plans: list[BaseStep] = []
     client = deployment.get_client()
-    jhelper = JujuHelper(deployment.get_connected_controller())
     proxy_settings = deployment.get_proxy_settings()
     fqdn = utils.get_fqdn()
 
@@ -248,11 +247,13 @@ def get_sunbeam_machine_plans(
 
 
 def get_k8s_plans(
-    deployment: Deployment, manifest: Manifest, accept_defaults: bool
+    deployment: Deployment,
+    jhelper: JujuHelper,
+    manifest: Manifest,
+    accept_defaults: bool,
 ) -> list[BaseStep]:
     plans: list[BaseStep] = []
     client = deployment.get_client()
-    jhelper = JujuHelper(deployment.get_connected_controller())
     fqdn = utils.get_fqdn()
 
     k8s_tfhelper = deployment.get_tfhelper("k8s-plan")
@@ -416,7 +417,6 @@ def get_juju_user_plans(
 
 def get_juju_bootstrap_plans(
     deployment: Deployment,
-    jhelper: JujuHelper,
     bootstrap_args: list,
 ):
     """Get Juju Bootstrap related plans."""
@@ -425,7 +425,7 @@ def get_juju_bootstrap_plans(
     bootstrap_args.extend(["--config", "controller-service-type=loadbalancer"])
 
     return [
-        AddK8SCloudInClientStep(deployment, jhelper),
+        AddK8SCloudInClientStep(deployment),
         BootstrapJujuStep(
             client,
             f"{deployment.name}{K8S_CLOUD_SUFFIX}",
@@ -453,6 +453,7 @@ def deploy_and_migrate_juju_controller(
     plan4: list[BaseStep] = []
     plan5: list[BaseStep] = []
     plan6: list[BaseStep] = []
+    plan7: list[BaseStep] = []
 
     client = deployment.get_client()
     fqdn = utils.get_fqdn()
@@ -481,29 +482,35 @@ def deploy_and_migrate_juju_controller(
     run_plan(plan2, console, show_hints)
 
     plan3 = get_juju_spaces_plans(deployment, jhelper, management_cidr)
-    plan3.extend(get_sunbeam_machine_plans(deployment, manifest))
-    plan3.extend(get_k8s_plans(deployment, manifest, accept_defaults))
-    plan3.extend(get_juju_bootstrap_plans(deployment, jhelper, juju_bootstrap_args))
+    plan3.extend(get_sunbeam_machine_plans(deployment, jhelper, manifest))
+    plan3.extend(get_k8s_plans(deployment, jhelper, manifest, accept_defaults))
     run_plan(plan3, console, show_hints)
+    # Disconnect all pylibjuju connections before bootstrapping new controller
+    run_sync(jhelper.disconnect())
+    del jhelper
 
-    plan4 = get_juju_migrate_plans(
+    plan4 = get_juju_bootstrap_plans(deployment, juju_bootstrap_args)
+    run_plan(plan4, console, show_hints)
+
+    plan5 = get_juju_migrate_plans(
         deployment, lxd_controller, deployment.controller, data_location
     )
-    run_plan(plan4, console, show_hints)
+    run_plan(plan5, console, show_hints)
     client.cluster.set_juju_controller_migrated()
 
     # Reload deployment with sunbeam-controller admin user credentials
     deployment.reload_credentials()
     jhelper = JujuHelper(deployment.get_connected_controller())
 
-    plan5 = [
+    plan6 = [
         CreateJujuUserStep(fqdn),
     ]
-    plan5_results = run_plan(plan5, console, show_hints)
-    token = get_step_message(plan5_results, CreateJujuUserStep)
+    plan6_results = run_plan(plan6, console, show_hints)
+    token = get_step_message(plan6_results, CreateJujuUserStep)
 
-    plan6 = get_juju_user_plans(deployment, jhelper, data_location, token)
-    run_plan(plan6, console, show_hints)
+    plan7 = get_juju_user_plans(deployment, jhelper, data_location, token)
+    run_plan(plan7, console, show_hints)
+    run_sync(jhelper.disconnect())
 
 
 @click.command()
@@ -684,8 +691,8 @@ def bootstrap(
         run_plan(plan12, console, show_hints)
 
         plan13 = get_juju_spaces_plans(deployment, jhelper, management_cidr)
-        plan13.extend(get_sunbeam_machine_plans(deployment, manifest))
-        plan13.extend(get_k8s_plans(deployment, manifest, accept_defaults))
+        plan13.extend(get_sunbeam_machine_plans(deployment, jhelper, manifest))
+        plan13.extend(get_k8s_plans(deployment, jhelper, manifest, accept_defaults))
         plan13.append(AddK8SCloudStep(deployment, jhelper))
         run_plan(plan13, console, show_hints)
     else:

sunbeam-python/sunbeam/steps/k8s.py

@@ -320,10 +320,9 @@ def run(self, status: Status | None = None) -> Result:
 class AddK8SCloudInClientStep(BaseStep, JujuStepHelper):
     _KUBECONFIG = K8S_KUBECONFIG_KEY
 
-    def __init__(self, deployment: Deployment, jhelper: JujuHelper):
+    def __init__(self, deployment: Deployment):
         super().__init__("Add K8S cloud in client", "Adding K8S cloud to Juju client")
         self.client = deployment.get_client()
-        self.jhelper = jhelper
         self.cloud_name = f"{deployment.name}{K8S_CLOUD_SUFFIX}"
         self.credential_name = f"{self.cloud_name}{CREDENTIAL_SUFFIX}"
 

sunbeam-python/tests/unit/sunbeam/core/test_juju.py

@@ -182,6 +182,7 @@ def jhelper_base(tmp_path: Path) -> juju.JujuHelper:
     jhelper = juju.JujuHelper.__new__(juju.JujuHelper)
     jhelper.data_location = tmp_path
     jhelper.controller = AsyncMock()  # type: ignore
+    jhelper.model_connectors = []
     return jhelper