modules/libjuju/juju/client/connector.py

import asyncio
import logging
import copy

import macaroonbakery.httpbakery as httpbakery
from juju.client.connection import Connection
from juju.client.gocookies import go_to_py_cookie, GoCookieJar
from juju.client.jujudata import FileJujuData
from juju.errors import JujuConnectionError, JujuError

log = logging.getLogger('connector')


class NoConnectionException(Exception):
    '''Raised by Connector when the connection method is called
    and there is no current connection.'''
    pass


class Connector:
    '''This class abstracts out a reconnectable client that can connect
    to controllers and models found in the Juju data files.
    '''
    def __init__(
        self,
        loop=None,
        max_frame_size=None,
        bakery_client=None,
        jujudata=None,
    ):
        '''Initialize a connector that will use the given parameters
        by default when making a new connection'''
        self.max_frame_size = max_frame_size
        self.loop = loop or asyncio.get_event_loop()
        self.bakery_client = bakery_client
        self._connection = None
        self.controller_name = None
        self.model_name = None
        self.jujudata = jujudata or FileJujuData()

    def is_connected(self):
        '''Report whether there is a currently connected controller or not'''
        return self._connection is not None

    def connection(self):
        '''Return the current connection; raises an exception if there
        is no current connection.'''
        if not self.is_connected():
            raise NoConnectionException('not connected')
        return self._connection

    async def connect(self, **kwargs):
        """Connect to an arbitrary Juju model.

        kwargs are passed through to Connection.connect()
        """
        kwargs.setdefault('loop', self.loop)
        kwargs.setdefault('max_frame_size', self.max_frame_size)
        kwargs.setdefault('bakery_client', self.bakery_client)
        if 'macaroons' in kwargs:
            if not kwargs['bakery_client']:
                kwargs['bakery_client'] = httpbakery.Client()
            if not kwargs['bakery_client'].cookies:
                kwargs['bakery_client'].cookies = GoCookieJar()
            jar = kwargs['bakery_client'].cookies
            for macaroon in kwargs.pop('macaroons'):
                jar.set_cookie(go_to_py_cookie(macaroon))
        self._connection = await Connection.connect(**kwargs)

    async def disconnect(self):
        """Shut down the watcher task and close websockets.
        """
        if self._connection:
            log.debug('Closing model connection')
            await self._connection.close()
            self._connection = None

    async def connect_controller(self, controller_name=None):
        """Connect to a controller by name. If the name is empty, it
        connect to the current controller.
        """
        if not controller_name:
            controller_name = self.jujudata.current_controller()
        if not controller_name:
            raise JujuConnectionError('No current controller')

        controller = self.jujudata.controllers()[controller_name]
        # TODO change Connection so we can pass all the endpoints
        # instead of just the first.
        endpoint = controller['api-endpoints'][0]
        accounts = self.jujudata.accounts().get(controller_name, {})

        await self.connect(
            endpoint=endpoint,
            uuid=None,
            username=accounts.get('user'),
            password=accounts.get('password'),
            cacert=controller.get('ca-cert'),
            bakery_client=self.bakery_client_for_controller(controller_name),
        )
        self.controller_name = controller_name

    async def connect_model(self, model_name=None):
        """Connect to a model by name. If either controller or model
        parts of the name are empty, the current controller and/or model
        will be used.

        :param str model: <controller>:<model>
        """

        try:
            controller_name, model_name = self.jujudata.parse_model(model_name)
            controller = self.jujudata.controllers().get(controller_name)
        except JujuError as e:
            raise JujuConnectionError(e.message) from e
        if controller is None:
            raise JujuConnectionError('Controller {} not found'.format(
                controller_name))
        # TODO change Connection so we can pass all the endpoints
        # instead of just the first one.
        endpoint = controller['api-endpoints'][0]
        account = self.jujudata.accounts().get(controller_name, {})
        models = self.jujudata.models().get(controller_name, {}).get('models',
                                                                     {})
        if model_name not in models:
            raise JujuConnectionError('Model not found: {}'.format(model_name))

        # TODO if there's no record for the required model name, connect
        # to the controller to find out the model's uuid, then connect
        # to that. This will let connect_model work with models that
        # haven't necessarily synced with the local juju data,
        # and also remove the need for base.CleanModel to
        # subclass JujuData.
        await self.connect(
            endpoint=endpoint,
            uuid=models[model_name]['uuid'],
            username=account.get('user'),
            password=account.get('password'),
            cacert=controller.get('ca-cert'),
            bakery_client=self.bakery_client_for_controller(controller_name),
        )
        self.controller_name = controller_name
        self.model_name = controller_name + ':' + model_name

    def bakery_client_for_controller(self, controller_name):
        '''Make a copy of the bakery client with a the appropriate controller's
        cookiejar in it.
        '''
        bakery_client = self.bakery_client
        if bakery_client:
            bakery_client = copy.copy(bakery_client)
        else:
            bakery_client = httpbakery.Client()
        bakery_client.cookies = self.jujudata.cookies_for_controller(
            controller_name)
        return bakery_client