From 8d38fd852e8d367966cdace1c5d444b80d2fcaf1 Mon Sep 17 00:00:00 2001
From: Caner Derici <caner.derici@canonical.com>
Date: Fri, 13 Oct 2023 14:42:21 -0600
Subject: [PATCH 1/2] Improve jujudata to find controller name by endpoint

Fixes #771
---
 juju/client/connector.py | 13 +++++++++----
 juju/client/jujudata.py  | 20 +++++++++++++++++++-
 juju/errors.py           |  2 ++
 3 files changed, 30 insertions(+), 5 deletions(-)

diff --git a/juju/client/connector.py b/juju/client/connector.py
index a3a54b1d..e51f610c 100644
--- a/juju/client/connector.py
+++ b/juju/client/connector.py
@@ -7,9 +7,9 @@
 import macaroonbakery.httpbakery as httpbakery
 from juju.client.connection import Connection
 from juju.client.gocookies import GoCookieJar, go_to_py_cookie
-from juju.client.jujudata import FileJujuData
+from juju.client.jujudata import FileJujuData, API_ENDPOINTS_KEY
 from juju.client.proxy.factory import proxy_from_config
-from juju.errors import JujuConnectionError, JujuError
+from juju.errors import JujuConnectionError, JujuError, PylibjujuProgrammingError
 from juju.client import client
 
 log = logging.getLogger('connector')
@@ -82,6 +82,11 @@ async def connect(self, **kwargs):
                 await self._connection.close()
             self._connection = await Connection.connect(**kwargs)
 
+        if not self.controller_name:
+            if 'endpoint' not in kwargs:
+                raise PylibjujuProgrammingError("Please report this error to the maintainers.")
+            self.controller_name = self.jujudata.controller_name_by_endpoint(kwargs['endpoint'])
+
         # Check if we support the target controller
         if not self._connection.info['server-version'].startswith(SUPPORTED_JUJU_API_PREFIX):
             raise JujuConnectionError("juju server-version %s not supported" % self._connection.info["server-version"])
@@ -108,7 +113,7 @@ async def connect_controller(self, controller_name=None, specified_facades=None)
             raise JujuConnectionError('No current controller')
 
         controller = self.jujudata.controllers()[controller_name]
-        endpoints = controller['api-endpoints']
+        endpoints = controller[API_ENDPOINTS_KEY]
         accounts = self.jujudata.accounts().get(controller_name, {})
 
         proxy = proxy_from_config(controller.get('proxy-config', None))
@@ -142,7 +147,7 @@ async def connect_model(self, model_name=None, **kwargs):
         if controller is None:
             raise JujuConnectionError('Controller {} not found'.format(
                 controller_name))
-        endpoints = controller['api-endpoints']
+        endpoints = controller[API_ENDPOINTS_KEY]
         account = self.jujudata.accounts().get(controller_name, {})
         models = self.jujudata.models().get(controller_name, {}).get('models',
                                                                      {})
diff --git a/juju/client/jujudata.py b/juju/client/jujudata.py
index a4c1d5ff..00e78340 100644
--- a/juju/client/jujudata.py
+++ b/juju/client/jujudata.py
@@ -10,9 +10,10 @@
 import yaml
 from juju import tag
 from juju.client.gocookies import GoCookieJar
-from juju.errors import JujuError
+from juju.errors import JujuError, PylibjujuProgrammingError
 from juju.utils import juju_config_dir
 
+API_ENDPOINTS_KEY = 'api-endpoints'
 
 class NoModelException(Exception):
     pass
@@ -126,6 +127,23 @@ def load_credential(self, cloud, name=None):
         except (KeyError, FileNotFoundError):
             return None, None
 
+    def controller_name_by_endpoint(self, endpoint):
+        """Finds the controller that has the given endpoints, returns the name.
+
+        :param str endpoint: The endpoint of the controller we're looking for
+        """
+        for controller_name, controller in self.controllers().items():
+            if isinstance(endpoint, str):
+                if endpoint in controller[API_ENDPOINTS_KEY]:
+                    return controller_name
+            elif isinstance(endpoint, list):
+                for e in endpoint:
+                    if e in controller[API_ENDPOINTS_KEY]:
+                        return controller_name
+            else:
+                raise PylibjujuProgrammingError()
+        raise JujuError(f'Unable to find controller with endpoint {endpoint}')
+
     def controllers(self):
         return self._load_yaml('controllers.yaml', 'controllers')
 
diff --git a/juju/errors.py b/juju/errors.py
index 6ca9711e..b43e8482 100644
--- a/juju/errors.py
+++ b/juju/errors.py
@@ -85,6 +85,8 @@ class JujuUnitError(JujuError):
 class JujuBackupError(JujuError):
     pass
 
+class PylibjujuProgrammingError(Exception):
+    pass
 
 class JujuConfigError(JujuError):
     """Exception raised during processing a configuration key-value pair

From 6e3d7b9bab10ff52b2131aa4480007c42c4ec61c Mon Sep 17 00:00:00 2001
From: Caner Derici <caner.derici@canonical.com>
Date: Fri, 13 Oct 2023 14:49:30 -0600
Subject: [PATCH 2/2] Fix linter

---
 juju/client/jujudata.py | 1 +
 juju/errors.py          | 2 ++
 2 files changed, 3 insertions(+)

diff --git a/juju/client/jujudata.py b/juju/client/jujudata.py
index 00e78340..b34c64a8 100644
--- a/juju/client/jujudata.py
+++ b/juju/client/jujudata.py
@@ -15,6 +15,7 @@
 
 API_ENDPOINTS_KEY = 'api-endpoints'
 
+
 class NoModelException(Exception):
     pass
 
diff --git a/juju/errors.py b/juju/errors.py
index b43e8482..61026dce 100644
--- a/juju/errors.py
+++ b/juju/errors.py
@@ -85,9 +85,11 @@ class JujuUnitError(JujuError):
 class JujuBackupError(JujuError):
     pass
 
+
 class PylibjujuProgrammingError(Exception):
     pass
 
+
 class JujuConfigError(JujuError):
     """Exception raised during processing a configuration key-value pair
     in a config set for an application.