[Namespace]: Fix SNMP AgentX socket connection timeout when using Namespace.get_all()

src/sonic_ax_impl/mibs/__init__.py

@@ -212,9 +212,6 @@ def init_sync_d_interface_tables(db_conn):
     :return: tuple(if_name_map, if_id_map, oid_map, if_alias_map)
     """
 
-    # Make sure we're connected to COUNTERS_DB
-    db_conn.connect(COUNTERS_DB)
-
     # { if_name (SONiC) -> sai_id }
     # ex: { "Ethernet76" : "1000000000023" }
     if_name_map, if_id_map = port_util.get_interface_oid_map(db_conn)
@@ -254,7 +251,6 @@ def init_sync_d_interface_tables(db_conn):
                        .format(port_util.SONIC_ETHERNET_RE_PATTERN))
         logger.warning("Port name map:\n" + pprint.pformat(if_name_map, indent=2))
 
-    db_conn.connect(APPL_DB)
 
     if_alias_map = dict()
 
@@ -325,13 +321,11 @@ def init_sync_d_lag_tables(db_conn):
     # { lag_name (SONiC) -> lag_oid (SAI) }
     lag_sai_map = {}
 
-    db_conn.connect(APPL_DB)
     lag_entries = db_conn.keys(APPL_DB, b"LAG_TABLE:*")
 
     if not lag_entries:
         return lag_name_if_name_map, if_name_lag_name_map, oid_lag_name_map, lag_sai_map
 
-    db_conn.connect(COUNTERS_DB)
     lag_sai_map = db_conn.get_all(COUNTERS_DB, b"COUNTERS_LAG_NAME_MAP")
     lag_sai_map = {name: sai_id.lstrip(b"oid:0x") for name, sai_id in lag_sai_map.items()}
 
@@ -363,9 +357,6 @@ def init_sync_d_queue_tables(db_conn):
     :return: tuple(port_queues_map, queue_stat_map)
     """
 
-    # Make sure we're connected to COUNTERS_DB
-    db_conn.connect(COUNTERS_DB)
-
     # { Port index : Queue index (SONiC) -> sai_id }
     # ex: { "1:2" : "1000000000023" }
     queue_name_map = db_conn.get_all(COUNTERS_DB, COUNTERS_QUEUE_NAME_MAP, blocking=True)
@@ -526,8 +517,15 @@ def init_namespace_dbs():
             db = SonicV2Connector(use_unix_socket_path=True, namespace=namespace)
             db_conn.append(db)
 
+        Namespace.connect_namespace_dbs(db_conn)
         return db_conn
 
+    @staticmethod
+    def connect_namespace_dbs(dbs):
+        list_of_dbs = [APPL_DB, COUNTERS_DB, CONFIG_DB, STATE_DB, ASIC_DB, SNMP_OVERLAY_DB]
+        for db_name in list_of_dbs:
+            Namespace.connect_all_dbs(dbs, db_name)
+
     @staticmethod
     def connect_all_dbs(dbs, db_name):
         for db_conn in dbs:
@@ -540,7 +538,6 @@ def dbs_keys(dbs, db_name, pattern='*'):
         """
         result_keys=[]
         for db_conn in dbs:
-            db_conn.connect(db_name)
             keys = db_conn.keys(db_name, pattern)
             if keys is not None:
                 result_keys.extend(keys)
@@ -552,12 +549,18 @@ def dbs_get_all(dbs, db_name, _hash, *args, **kwargs):
         db get_all function executed on global and all namespace DBs.
         """
         result = {}
+        # If there are multiple namespaces, _hash might not be 
+        # present in all namespace, ignore if not present in a
+        # specfic namespace.
+        if len(dbs) > 1:
+            tmp_kwargs = kwargs.copy()
+            tmp_kwargs['blocking'] = False
+        else:
+            tmp_kwargs = kwargs
         for db_conn in dbs:
-            db_conn.connect(db_name)
-            if(db_conn.exists(db_name, _hash)):
-                ns_result = db_conn.get_all(db_name, _hash, *args, **kwargs)
-                if ns_result is not None:
-                    result.update(ns_result)
+            ns_result = db_conn.get_all(db_name, _hash, *args, **tmp_kwargs)
+            if ns_result is not None:
+                result.update(ns_result)
         return result
 
     @staticmethod
@@ -589,6 +592,7 @@ def init_namespace_sync_d_interface_tables(dbs):
         Ignore first global db to get interface tables if there
         are multiple namespaces.
         """
+        Namespace.connect_namespace_dbs(dbs)
         for db_conn in Namespace.get_non_host_dbs(dbs):
             if_name_map_ns, \
             if_alias_map_ns, \
@@ -617,6 +621,7 @@ def init_namespace_sync_d_lag_tables(dbs):
         Ignore first global db to get lag tables if
         there are multiple namespaces.
         """
+        Namespace.connect_namespace_dbs(dbs)
         for db_conn in Namespace.get_non_host_dbs(dbs):
             lag_name_if_name_map_ns, \
             if_name_lag_name_map_ns, \
@@ -634,6 +639,7 @@ def init_namespace_sync_d_rif_tables(dbs):
         rif_port_map = {}
         port_rif_map = {}
 
+        Namespace.connect_namespace_dbs(dbs)
         for db_conn in Namespace.get_non_host_dbs(dbs):
             rif_port_map_ns, \
             port_rif_map_ns = init_sync_d_rif_tables(db_conn)
@@ -648,6 +654,7 @@ def init_namespace_sync_d_vlan_tables(dbs):
         oid_sai_map = {}
         oid_name_map = {}
 
+        Namespace.connect_namespace_dbs(dbs)
         for db_conn in Namespace.get_non_host_dbs(dbs):
             vlan_name_map_ns, \
             oid_sai_map_ns, \
@@ -670,6 +677,7 @@ def init_namespace_sync_d_queue_tables(dbs):
         Ignore first global db to get queue tables if there
         are multiple namespaces.
         """
+        Namespace.connect_namespace_dbs(dbs)
         for db_conn in Namespace.get_non_host_dbs(dbs):
             port_queues_map_ns, \
             queue_stat_map_ns, \

tests/test_mibs.py

@@ -7,7 +7,7 @@
 modules_path = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
 sys.path.insert(0, os.path.join(modules_path, 'src'))
 
-from sonic_ax_impl import mibs
+from sonic_ax_impl.mibs import Namespace
 
 class TestGetNextPDU(TestCase):
     @classmethod
@@ -16,12 +16,12 @@ def setUpClass(cls):
         tests.mock_tables.dbconnector.load_database_config() 
 
     def test_init_sync_d_lag_tables(self):
-        db_conn = mibs.init_db()
+        db_conn = Namespace.init_namespace_dbs()
 
         lag_name_if_name_map, \
         if_name_lag_name_map, \
         oid_lag_name_map, \
-        lag_sai_map = mibs.init_sync_d_lag_tables(db_conn)
+        lag_sai_map = Namespace.init_namespace_sync_d_lag_tables(db_conn)
 
         self.assertTrue(b"PortChannel04" in lag_name_if_name_map)
         self.assertTrue(lag_name_if_name_map[b"PortChannel04"] == [b"Ethernet124"])