[cli]: pass db connector as click context (sonic-net#1029)

it is better not let every command to have its
own connector and connect to db.

it is also not good to have a global db variable.

Here, the idea is to have a single db connector
and pass the connector as a click context

Signed-off-by: Guohan Lu <[email protected]>

Author: lguohan

config/main.py

@@ -19,6 +19,7 @@
 from config_mgmt import ConfigMgmtDPB
 from utilities_common.intf_filter import parse_interface_in_filter
 from utilities_common.util_base import UtilHelper
+from utilities_common.db import Db
 from portconfig import get_child_ports, get_port_config_file_name
 
 import aaa
@@ -50,7 +51,6 @@
 CFG_LOOPBACK_NO="<0-999>"
 
 asic_type = None
-config_db = None
 
 # ========================== Syslog wrappers ==========================
 
@@ -691,7 +691,7 @@ def _abort_if_false(ctx, param, value):
         ctx.abort()
 
 
-def _get_disabled_services_list():
+def _get_disabled_services_list(config_db):
     disabled_services_list = []
 
     feature_table = config_db.get_table('FEATURE')
@@ -713,7 +713,7 @@ def _get_disabled_services_list():
 
     return disabled_services_list
 
-def _stop_services():
+def _stop_services(config_db):
     # This list is order-dependent. Please add services in the order they should be stopped
     # on Mellanox platform pmon is stopped by syncd
     services_to_stop = [
@@ -730,7 +730,7 @@ def _stop_services():
     if asic_type == 'mellanox' and 'pmon' in services_to_stop:
         services_to_stop.remove('pmon')
 
-    disabled_services = _get_disabled_services_list()
+    disabled_services = _get_disabled_services_list(config_db)
 
     for service in disabled_services:
         if service in services_to_stop:
@@ -739,7 +739,7 @@ def _stop_services():
     execute_systemctl(services_to_stop, SYSTEMCTL_ACTION_STOP)
 
 
-def _reset_failed_services():
+def _reset_failed_services(config_db):
     # This list is order-independent. Please keep list in alphabetical order
     services_to_reset = [
         'bgp',
@@ -762,7 +762,7 @@ def _reset_failed_services():
         'telemetry'
     ]
 
-    disabled_services = _get_disabled_services_list()
+    disabled_services = _get_disabled_services_list(config_db)
 
     for service in disabled_services:
         if service in services_to_reset:
@@ -771,7 +771,7 @@ def _reset_failed_services():
     execute_systemctl(services_to_reset, SYSTEMCTL_ACTION_RESET_FAILED)
 
 
-def _restart_services():
+def _restart_services(config_db):
     # This list is order-dependent. Please add services in the order they should be started
     # on Mellanox platform pmon is started by syncd
     services_to_restart = [
@@ -790,7 +790,7 @@ def _restart_services():
         'telemetry'
     ]
 
-    disabled_services = _get_disabled_services_list()
+    disabled_services = _get_disabled_services_list(config_db)
 
     for service in disabled_services:
         if service in services_to_restart:
@@ -908,7 +908,8 @@ def validate_mirror_session_config(config_db, session_name, dst_port, src_port,
 
 # This is our main entrypoint - the main 'config' command
 @click.group(cls=AbbreviationGroup, context_settings=CONTEXT_SETTINGS)
-def config():
+@click.pass_context
+def config(ctx):
     """SONiC command line - 'config' command"""
     #
     # Load asic_type for further use
@@ -932,10 +933,9 @@ def config():
 
     SonicDBConfig.load_sonic_global_db_config()
 
-    global config_db
+    ctx.obj = Db()
 
-    config_db = ConfigDBConnector()
-    config_db.connect()
+pass_db = click.make_pass_decorator(Db, ensure=True)
 
 config.add_command(aaa.aaa)
 config.add_command(aaa.tacacs)
@@ -1060,7 +1060,8 @@ def load(filename, yes):
 @click.option('-l', '--load-sysinfo', is_flag=True, help='load system default information (mac, portmap etc) first.')
 @click.option('-n', '--no_service_restart', default=False, is_flag=True, help='Do not restart docker services')
 @click.argument('filename', required=False)
-def reload(filename, yes, load_sysinfo, no_service_restart):
+@pass_db
+def reload(db, filename, yes, load_sysinfo, no_service_restart):
     """Clear current configuration and import a previous saved config DB dump file.
        <filename> : Names of configuration file(s) to load, separated by comma with no spaces in between
     """
@@ -1102,7 +1103,7 @@ def reload(filename, yes, load_sysinfo, no_service_restart):
     #Stop services before config push
     if not no_service_restart:
         log_info("'reload' stopping services...")
-        _stop_services()
+        _stop_services(db.cfgdb)
 
     """ In Single AISC platforms we have single DB service. In multi-ASIC platforms we have a global DB
         service running in the host + DB services running in each ASIC namespace created per ASIC.
@@ -1175,9 +1176,9 @@ def reload(filename, yes, load_sysinfo, no_service_restart):
     # We first run "systemctl reset-failed" to remove the "failed"
     # status from all services before we attempt to restart them
     if not no_service_restart:
-        _reset_failed_services()
+        _reset_failed_services(db.cfgdb)
         log_info("'reload' restarting services...")
-        _restart_services()
+        _restart_services(db.cfgdb)
 
 @config.command("load_mgmt_config")
 @click.option('-y', '--yes', is_flag=True, callback=_abort_if_false,
@@ -1208,14 +1209,15 @@ def load_mgmt_config(filename):
 @click.option('-y', '--yes', is_flag=True, callback=_abort_if_false,
                 expose_value=False, prompt='Reload config from minigraph?')
 @click.option('-n', '--no_service_restart', default=False, is_flag=True, help='Do not restart docker services')
-def load_minigraph(no_service_restart):
+@pass_db
+def load_minigraph(db, no_service_restart):
     """Reconfigure based on minigraph."""
     log_info("'load_minigraph' executing...")
 
     #Stop services before config push
     if not no_service_restart:
         log_info("'load_minigraph' stopping services...")
-        _stop_services()
+        _stop_services(db.cfgdb)
 
     # For Single Asic platform the namespace list has the empty string
     # for mulit Asic platform the empty string to generate the config
@@ -1269,10 +1271,10 @@ def load_minigraph(no_service_restart):
     # We first run "systemctl reset-failed" to remove the "failed"
     # status from all services before we attempt to restart them
     if not no_service_restart:
-        _reset_failed_services()
+        _reset_failed_services(db.cfgdb)
         #FIXME: After config DB daemon is implemented, we'll no longer need to restart every service.
         log_info("'load_minigraph' restarting services...")
-        _restart_services()
+        _restart_services(db.cfgdb)
     click.echo("Please note setting loaded from minigraph will be lost after system reboot. To preserve setting, run `config save`.")
 
 
@@ -3810,25 +3812,27 @@ def feature():
 @feature.command('state', short_help="Enable/disable a feature")
 @click.argument('name', metavar='<feature-name>', required=True)
 @click.argument('state', metavar='<state>', required=True, type=click.Choice(["enabled", "disabled"]))
-def feature_state(name, state):
+@pass_db
+def feature_state(db, name, state):
     """Enable/disable a feature"""
-    state_data = config_db.get_entry('FEATURE', name)
+    state_data = db.cfgdb.get_entry('FEATURE', name)
 
     if not state_data:
         click.echo("Feature '{}' doesn't exist".format(name))
         sys.exit(1)
 
-    config_db.mod_entry('FEATURE', name, {'state': state})
+    db.cfgdb.mod_entry('FEATURE', name, {'state': state})
 
 #
 # 'autorestart' command ('config feature autorestart ...')
 #
 @feature.command(name='autorestart', short_help="Enable/disable autosrestart of a feature")
 @click.argument('name', metavar='<feature-name>', required=True)
 @click.argument('autorestart', metavar='<autorestart>', required=True, type=click.Choice(["enabled", "disabled"]))
-def feature_autorestart(name, autorestart):
+@pass_db
+def feature_autorestart(db, name, autorestart):
     """Enable/disable autorestart of a feature"""
-    feature_table = config_db.get_table('FEATURE')
+    feature_table = db.cfgdb.get_table('FEATURE')
     if not feature_table:
         click.echo("Unable to retrieve feature table from Config DB.")
         sys.exit(1)
@@ -3837,7 +3841,7 @@ def feature_autorestart(name, autorestart):
         click.echo("Unable to retrieve feature '{}'".format(name))
         sys.exit(1)
 
-    config_db.mod_entry('FEATURE', name, {'auto_restart': autorestart})
+    db.cfgdb.mod_entry('FEATURE', name, {'auto_restart': autorestart})
 
 if __name__ == '__main__':
     config()

show/main.py

@@ -19,6 +19,7 @@
 from swsssdk import ConfigDBConnector
 from swsssdk import SonicV2Connector
 from portconfig import get_child_ports
+from utilities_common.db import Db
 
 import mlnx
 
@@ -31,8 +32,6 @@
 
 VLAN_SUB_INTERFACE_SEPARATOR = '.'
 
-config_db = None
-
 try:
     # noinspection PyPep8Naming
     import ConfigParser as configparser
@@ -557,12 +556,13 @@ def get_bgp_neighbor_ip_to_name(ip, static_neighbors, dynamic_neighbors):
 # This is our entrypoint - the main "show" command
 # TODO: Consider changing function name to 'show' for better understandability
 @click.group(cls=AliasedGroup, context_settings=CONTEXT_SETTINGS)
-def cli():
+@click.pass_context
+def cli(ctx):
     """SONiC command line - 'show' command"""
-    global config_db
 
-    config_db = ConfigDBConnector()
-    config_db.connect()
+    ctx.obj = Db()
+
+pass_db = click.make_pass_decorator(Db, ensure=True)
 
 #
 # 'vrf' command ("show vrf")
@@ -3064,14 +3064,15 @@ def feature():
     pass
 
 #
-# 'state' subcommand (show feature status)
+# 'status' subcommand (show feature status)
 #
 @feature.command('status', short_help="Show feature state")
 @click.argument('feature_name', required=False)
-def autorestart(feature_name):
+@pass_db
+def feature_status(db, feature_name):
     header = ['Feature', 'State', 'AutoRestart']
     body = []
-    feature_table = config_db.get_table('FEATURE')
+    feature_table = db.cfgdb.get_table('FEATURE')
     if feature_name:
         if feature_table and feature_table.has_key(feature_name):
             body.append([feature_name, feature_table[feature_name]['state'], \
@@ -3089,10 +3090,11 @@ def autorestart(feature_name):
 #
 @feature.command('autorestart', short_help="Show auto-restart state for a feature")
 @click.argument('feature_name', required=False)
-def autorestart(feature_name):
+@pass_db
+def feature_autorestart(db, feature_name):
     header = ['Feature', 'AutoRestart']
     body = []
-    feature_table = config_db.get_table('FEATURE')
+    feature_table = db.cfgdb.get_table('FEATURE')
     if feature_name:
         if feature_table and feature_table.has_key(feature_name):
             body.append([feature_name, feature_table[feature_name]['auto_restart']])

tests/config_test.py

@@ -1,5 +1,9 @@
+import traceback
+
 from click.testing import CliRunner
 
+from utilities_common.db import Db
+
 load_minigraph_command_output="""\
 Executing stop of service telemetry...
 Executing stop of service swss...
@@ -52,16 +56,20 @@ def test_load_minigraph(self, get_cmd_module, setup_single_broacom_asic):
         result = runner.invoke(config.config.commands["load_minigraph"], ["-y"])
         print result.exit_code
         print result.output
+        traceback.print_tb(result.exc_info[2])
         assert result.exit_code == 0
         assert "\n".join([ l.rstrip() for l in result.output.split('\n')]) == load_minigraph_command_output
 
     def test_load_minigraph_with_disabled_telemetry(self, get_cmd_module, setup_single_broacom_asic):
         (config, show) = get_cmd_module
+        db = Db()
         runner = CliRunner()
-        runner.invoke(config.config.commands["feature"].commands["state"], ["telemetry", "disabled"])
-        result = runner.invoke(show.cli.commands["feature"].commands["status"], ["telemetry"])
+        result = runner.invoke(config.config.commands["feature"].commands["state"], ["telemetry", "disabled"], obj=db)
+        assert result.exit_code == 0
+        result = runner.invoke(show.cli.commands["feature"].commands["status"], ["telemetry"], obj=db)
         print result.output
-        result = runner.invoke(config.config.commands["load_minigraph"], ["-y"])
+        assert result.exit_code == 0
+        result = runner.invoke(config.config.commands["load_minigraph"], ["-y"], obj=db)
         print result.exit_code
         print result.output
         assert result.exit_code == 0

tests/conftest.py

@@ -58,12 +58,6 @@ def get_cmd_module():
     import config.main as config
     import show.main as show
 
-    config_db = ConfigDBConnector()
-    config_db.connect()
-
-    config.config_db = config_db
-    show.config_db = config_db
-
     config.run_command = _dummy_run_command
 
     return (config, show)

tests/feature_test.py

@@ -1,5 +1,7 @@
 from click.testing import CliRunner
 
+from utilities_common.db import Db
+
 show_feature_status_output="""\
 Feature     State     AutoRestart
 ----------  --------  -------------
@@ -122,22 +124,26 @@ def test_show_unknown_autorestart_status(self, get_cmd_module):
 
     def test_config_bgp_feature_state(self, get_cmd_module):
         (config, show) = get_cmd_module
+        db = Db()
         runner = CliRunner()
-        result = runner.invoke(config.config.commands["feature"].commands["state"], ["bgp", "disabled"])
+        result = runner.invoke(config.config.commands["feature"].commands["state"], ["bgp", "disabled"], obj=db)
         print(result.exit_code)
         print(result.output)
-        result = runner.invoke(show.cli.commands["feature"].commands["status"], ["bgp"])
+        assert result.exit_code == 0
+        result = runner.invoke(show.cli.commands["feature"].commands["status"], ["bgp"], obj=db)
         print(result.output)
         assert result.exit_code == 0
         assert result.output == show_feature_bgp_disabled_status_output
 
     def test_config_bgp_autorestart(self, get_cmd_module):
         (config, show) = get_cmd_module
+        db = Db()
         runner = CliRunner()
-        result = runner.invoke(config.config.commands["feature"].commands["autorestart"], ["bgp", "disabled"])
+        result = runner.invoke(config.config.commands["feature"].commands["autorestart"], ["bgp", "disabled"], obj=db)
         print(result.exit_code)
         print(result.output)
-        result = runner.invoke(show.cli.commands["feature"].commands["autorestart"], ["bgp"])
+        assert result.exit_code == 0
+        result = runner.invoke(show.cli.commands["feature"].commands["autorestart"], ["bgp"], obj=db)
         print(result.output)
         assert result.exit_code == 0
         assert result.output == show_feature_bgp_disabled_autorestart_output

utilities_common/db.py

@@ -0,0 +1,6 @@
+from swsssdk import ConfigDBConnector
+
+class Db(object):
+    def __init__(self):
+        self.cfgdb = ConfigDBConnector()
+        self.cfgdb.connect()