Add etcd commands for admin (#84)

* Add a set of commands under "backend.ai admin etcd", just like the manager CLI.

Author: adrysn

src/ai/backend/client/cli/admin/__init__.py

@@ -10,7 +10,7 @@ def admin():
 
 def _attach_command():
     from . import (  # noqa
-        agents, domains, groups, images, keypairs, resources, resource_policies,
+        agents, domains, etcd, groups, images, keypairs, resources, resource_policies,
         scaling_groups, sessions, users, vfolders
     )
 

src/ai/backend/client/cli/admin/etcd.py

@@ -0,0 +1,88 @@
+import json
+import sys
+
+import click
+
+from . import admin
+from ..pretty import print_pretty, print_error, print_fail
+from ...session import Session
+
+
+@admin.group(invoke_without_command=False)
+@click.pass_context
+def etcd(ctx):
+    '''
+    List and manage ETCD configurations.
+    (admin privilege required)
+    '''
+    if ctx.invoked_subcommand is not None:
+        return
+
+
+@etcd.command()
+@click.argument('key', type=str, metavar='KEY')
+@click.option('-p', '--prefix', is_flag=True, default=False,
+              help='Get all keys prefixed with the given key.')
+def get(key, prefix):
+    '''
+    Get a ETCD value(s).
+
+    KEY: Name of ETCD key.
+    '''
+    with Session() as session:
+        try:
+            data = session.EtcdConfig.get(key, prefix)
+        except Exception as e:
+            print_error(e)
+            sys.exit(1)
+        data = json.dumps(data, indent=2) if data else 'null'
+        print_pretty(data)
+
+
+@etcd.command()
+@click.argument('key', type=str, metavar='KEY')
+@click.argument('value', type=str, metavar='VALUE')
+def set(key, value):
+    '''
+    Set new key and value on ETCD.
+
+    KEY: Name of ETCD key.
+    VALUE: Value to set.
+    '''
+    with Session() as session:
+        try:
+            value = json.loads(value)
+            print_pretty('Value converted to a dictionary.')
+        except json.JSONDecodeError:
+            pass
+        try:
+            data = session.EtcdConfig.set(key, value)
+        except Exception as e:
+            print_error(e)
+            sys.exit(1)
+        if data.get('result', False) != 'ok':
+            print_fail('Unable to set key/value.')
+        else:
+            print_pretty('Successfully set key/value.')
+
+
+@etcd.command()
+@click.argument('key', type=str, metavar='KEY')
+@click.option('-p', '--prefix', is_flag=True, default=False,
+              help='Delete all keys prefixed with the given key.')
+def delete(key, prefix):
+    '''
+    Delete key(s) from ETCD.
+
+    KEY: Name of ETCD key.
+    '''
+    with Session() as session:
+        try:
+            data = session.EtcdConfig.delete(key, prefix)
+        except Exception as e:
+            print_error(e)
+            sys.exit(1)
+        if data.get('result', False) != 'ok':
+            print_fail('Unable to delete key/value.')
+        else:
+            print_pretty('Successfully deleted key/value.')

src/ai/backend/client/etcd.py

@@ -0,0 +1,74 @@
+from .base import api_function
+from .request import Request
+
+__all__ = (
+    'EtcdConfig',
+)
+
+
+class EtcdConfig:
+    '''
+    Provides a way to get or set ETCD configurations.
+
+    .. note::
+
+      All methods in this function class require your API access key to
+      have the *superadmin* privilege.
+    '''
+
+    session = None
+    '''The client session instance that this function class is bound to.'''
+
+    @api_function
+    @classmethod
+    async def get(cls, key: str, prefix: bool = False) -> dict:
+        '''
+        Get configuration from ETCD with given key.
+
+        :param key: Name of the key to fetch.
+        :param prefix: get all keys prefixed with the give key.
+        '''
+        rqst = Request(cls.session, 'POST', '/config/get')
+        rqst.set_json({
+            'key': key,
+            'prefix': prefix,
+        })
+        async with rqst.fetch() as resp:
+            data = await resp.json()
+            return data.get('result', None)
+
+    @api_function
+    @classmethod
+    async def set(cls, key: str, value: str) -> dict:
+        '''
+        Set configuration into ETCD with given key and value.
+
+        :param key: Name of the key to set.
+        :param value: Value to set.
+        '''
+        rqst = Request(cls.session, 'POST', '/config/set')
+        rqst.set_json({
+            'key': key,
+            'value': value,
+        })
+        async with rqst.fetch() as resp:
+            data = await resp.json()
+            return data
+
+    @api_function
+    @classmethod
+    async def delete(cls, key: str, prefix: bool = False) -> dict:
+        '''
+        Delete configuration from ETCD with given key.
+
+        :param key: Name of the key to delete.
+        :param prefix: delete all keys prefixed with the give key.
+        '''
+        rqst = Request(cls.session, 'POST', '/config/delete')
+        rqst.set_json({
+            'key': key,
+            'prefix': prefix,
+        })
+        async with rqst.fetch() as resp:
+            data = await resp.json()
+            return data

src/ai/backend/client/session.py

@@ -72,10 +72,13 @@ class BaseSession(metaclass=abc.ABCMeta):
 
     __slots__ = (
         '_config', '_closed', 'aiohttp_session',
-        'Admin', 'Agent', 'AgentWathcer', 'Auth', 'Domain', 'Group', 'ScalingGroup',
-        'Admin', 'Agent', 'AgentWatcher', 'Domain', 'Group', 'ScalingGroup',
-        'Image', 'Kernel', 'KeyPair', 'Manager', 'Resource',
-        'KeypairResourcePolicy', 'User', 'VFolder',
+        'Manager', 'Admin',
+        'Agent', 'AgentWatcher', 'ScalingGroup',
+        'Image', 'Kernel',
+        'Domain', 'Group', 'Auth', 'User', 'KeyPair',
+        'EtcdConfig',
+        'Resource', 'KeypairResourcePolicy',
+        'VFolder',
     )
 
     def __init__(self, *, config: APIConfig = None):
@@ -134,6 +137,7 @@ async def _create_aiohttp_session():
         from .admin import Admin
         from .agent import Agent, AgentWatcher
         from .auth import Auth
+        from .etcd import EtcdConfig
         from .domain import Domain
         from .group import Group
         from .image import Image
@@ -182,6 +186,15 @@ async def _create_aiohttp_session():
         bound to this session.
         '''
 
+        self.EtcdConfig = type('EtcdConfig', (BaseFunction, ), {
+            **EtcdConfig.__dict__,
+            'session': self,
+        })
+        '''
+        The :class:`~ai.backend.client.EtcdConfig` function proxy
+        bound to this session.
+        '''
+
         self.Domain = type('Domain', (BaseFunction, ), {
             **Domain.__dict__,
             'session': self,
@@ -335,6 +348,7 @@ def __init__(self, *, config: APIConfig = None):
         from .admin import Admin
         from .agent import Agent, AgentWatcher
         from .auth import Auth
+        from .etcd import EtcdConfig
         from .group import Group
         from .image import Image
         from .kernel import Kernel
@@ -382,6 +396,15 @@ def __init__(self, *, config: APIConfig = None):
         bound to this session.
         '''
 
+        self.EtcdConfig = type('EtcdConfig', (BaseFunction, ), {
+            **EtcdConfig.__dict__,
+            'session': self,
+        })
+        '''
+        The :class:`~ai.backend.client.EtcdConfig` function proxy
+        bound to this session.
+        '''
+
         self.Group = type('Group', (BaseFunction, ), {
             **Group.__dict__,
             'session': self,