[sonic-package-manager] support extension with multiple YANG modules

sonic_package_manager/manifest.py

@@ -212,6 +212,8 @@ def unmarshal(self, value):
             ManifestField('clear', DefaultMarshaller(str), ''),
             ManifestField('auto-generate-show', DefaultMarshaller(bool), False),
             ManifestField('auto-generate-config', DefaultMarshaller(bool), False),
+            ManifestArray('auto-generate-show-source-yang-modules', DefaultMarshaller(str)),
+            ManifestArray('auto-generate-config-source-yang-modules', DefaultMarshaller(str)),
         ])
     ])
 

sonic_package_manager/metadata.py

@@ -4,10 +4,11 @@
 
 import json
 import tarfile
-from typing import Dict, Optional
+from typing import Dict, List
 
 from sonic_package_manager import utils
 from sonic_package_manager.errors import MetadataError
+from sonic_package_manager.logger import log
 from sonic_package_manager.manifest import Manifest
 from sonic_package_manager.version import Version
 
@@ -54,7 +55,7 @@ class Metadata:
 
     manifest: Manifest
     components: Dict[str, Version] = field(default_factory=dict)
-    yang_module_str: Optional[str] = None
+    yang_modules: List[str] = field(default_factory=list)
 
 
 class MetadataResolver:
@@ -164,6 +165,16 @@ def from_labels(cls, labels: Dict[str, str]) -> Metadata:
                 except ValueError as err:
                     raise MetadataError(f'Failed to parse component version: {err}')
 
-        yang_module_str = sonic_metadata.get('yang-module')
+        labels_yang_modules = sonic_metadata.get('yang-module')
+        yang_modules = []
 
-        return Metadata(Manifest.marshal(manifest_dict), components, yang_module_str)
+        if isinstance(labels_yang_modules, str):
+            yang_modules.append(labels_yang_modules)
+            log.debug("Found one YANG module")
+        elif isinstance(labels_yang_modules, dict):
+            yang_modules.extend(labels_yang_modules.values())
+            log.debug(f"Found YANG modules: {labels_yang_modules.keys()}")
+        else:
+            log.debug("No YANG modules found")
+
+        return Metadata(Manifest.marshal(manifest_dict), components, yang_modules)

sonic_package_manager/service_creator/creator.py

@@ -518,18 +518,19 @@ def remove_config(self, package):
             None
         """
 
-        if not package.metadata.yang_module_str:
+        if not package.metadata.yang_modules:
             return
 
-        module_name = self.cfg_mgmt.get_module_name(package.metadata.yang_module_str)
-        for tablename, module in self.cfg_mgmt.sy.confDbYangMap.items():
-            if module.get('module') != module_name:
-                continue
+        for module in package.metadata.yang_modules:
+            module_name = self.cfg_mgmt.get_module_name(module)
+            for tablename, module in self.cfg_mgmt.sy.confDbYangMap.items():
+                if module.get('module') != module_name:
+                    continue
 
-            for conn in self.sonic_db.get_connectors():
-                keys = conn.get_table(tablename).keys()
-                for key in keys:
-                    conn.set_entry(tablename, key, None)
+                for conn in self.sonic_db.get_connectors():
+                    keys = conn.get_table(tablename).keys()
+                    for key in keys:
+                        conn.set_entry(tablename, key, None)
 
     def validate_config(self, config):
         """ Validate configuration through YANG.
@@ -560,10 +561,11 @@ def install_yang_module(self, package: Package):
             None
         """
 
-        if not package.metadata.yang_module_str:
+        if not package.metadata.yang_modules:
             return
 
-        self.cfg_mgmt.add_module(package.metadata.yang_module_str)
+        for module in package.metadata.yang_modules:
+            self.cfg_mgmt.add_module(module)
 
     def uninstall_yang_module(self, package: Package):
         """ Uninstall package's yang module in the system.
@@ -574,11 +576,12 @@ def uninstall_yang_module(self, package: Package):
             None
         """
 
-        if not package.metadata.yang_module_str:
+        if not package.metadata.yang_modules:
             return
 
-        module_name = self.cfg_mgmt.get_module_name(package.metadata.yang_module_str)
-        self.cfg_mgmt.remove_module(module_name)
+        for module in package.metadata.yang_modules:
+            module_name = self.cfg_mgmt.get_module_name(module)
+            self.cfg_mgmt.remove_module(module_name)
 
     def install_autogen_cli_all(self, package: Package):
         """ Install autogenerated CLI plugins for package.
@@ -614,15 +617,16 @@ def install_autogen_cli(self, package: Package, command: str):
             None
         """
 
-        if package.metadata.yang_module_str is None:
+        if not package.metadata.yang_modules:
             return
         if f'auto-generate-{command}' not in package.manifest['cli']:
             return
         if not package.manifest['cli'][f'auto-generate-{command}']:
             return
-        module_name = self.cfg_mgmt.get_module_name(package.metadata.yang_module_str)
-        self.cli_gen.generate_cli_plugin(command, module_name)
-        log.debug(f'{command} command line interface autogenerated for {module_name}')
+
+        for module_name in self._get_yang_modules_for_auto_gen(command, package):
+            self.cli_gen.generate_cli_plugin(command, module_name)
+            log.debug(f'{command} command line interface autogenerated for {module_name}')
 
     def uninstall_autogen_cli(self, package: Package, command: str):
         """ Uninstall autogenerated CLI plugins for package for particular command.
@@ -634,18 +638,33 @@ def uninstall_autogen_cli(self, package: Package, command: str):
             None
         """
 
-        if package.metadata.yang_module_str is None:
+        if not package.metadata.yang_modules:
             return
         if f'auto-generate-{command}' not in package.manifest['cli']:
             return
         if not package.manifest['cli'][f'auto-generate-{command}']:
             return
-        module_name = self.cfg_mgmt.get_module_name(package.metadata.yang_module_str)
-        self.cli_gen.remove_cli_plugin(command, module_name)
-        log.debug(f'{command} command line interface removed for {module_name}')
+
+        for module_name in self._get_yang_modules_for_auto_gen(command, package):
+            self.cli_gen.remove_cli_plugin(command, module_name)
+            log.debug(f'{command} command line interface removed for {module_name}')
 
     def _post_operation_hook(self):
         """ Common operations executed after service is created/removed. """
 
         if not in_chroot():
             run_command(['systemctl', 'daemon-reload'])
+
+    def _get_yang_modules_for_auto_gen(self, command: str, package: Package):
+        source_yang_modules = package.manifest['cli'][f'auto-generate-{command}-source-yang-modules']
+
+        def filter_yang_modules_for_auto_gen(module_name):
+            if not source_yang_modules:
+                return True
+            if module_name in source_yang_modules:
+                return True
+            return False
+
+        filtered_yang_modules = filter(filter_yang_modules_for_auto_gen,
+                map(self.cfg_mgmt.get_module_name, package.metadata.yang_modules))
+        return list(filtered_yang_modules)

tests/sonic_package_manager/test_metadata.py

@@ -1,14 +1,36 @@
 #!/usr/bin/env python
 
+import json
 import contextlib
 from unittest.mock import Mock, MagicMock
 
+import pytest
+
 from sonic_package_manager.database import PackageEntry
 from sonic_package_manager.errors import MetadataError
+from sonic_package_manager.manifest import Manifest
 from sonic_package_manager.metadata import MetadataResolver
 from sonic_package_manager.version import Version
 
 
+@pytest.fixture
+def manifest_str():
+    return json.dumps({
+        'package': {
+            'name': 'test',
+            'version': '1.0.0',
+        },
+        'service': {
+            'name': 'test',
+            'asic-service': False,
+            'host-service': True,
+        },
+        'container': {
+            'privileged': True,
+        },
+    })
+
+
 def test_metadata_resolver_local(mock_registry_resolver, mock_docker_api):
     metadata_resolver = MetadataResolver(mock_docker_api, mock_registry_resolver)
     # it raises exception because mock manifest is not a valid manifest
@@ -35,3 +57,33 @@ def return_mock_registry(repository):
     mock_registry.manifest.assert_called_once_with('test-repository', '1.2.0')
     mock_registry.blobs.assert_called_once_with('test-repository', 'some-digest')
     mock_docker_api.labels.assert_not_called()
+
+
+def test_metadata_construction(manifest_str):
+    metadata = MetadataResolver.from_labels({
+        'com': {
+            'azure': {
+                'sonic': {
+                    'manifest': manifest_str,
+                    'yang-module': 'TEST'
+                }
+            }
+        }
+    })
+    assert metadata.yang_modules == ['TEST']
+
+    metadata = MetadataResolver.from_labels({
+        'com': {
+            'azure': {
+                'sonic': {
+                    'manifest': manifest_str,
+                    'yang-module': {
+                        'sonic-test': 'TEST',
+                        'sonic-test-2': 'TEST 2',
+                    },
+                },
+            },
+        },
+    })
+    assert metadata.yang_modules == ['TEST', 'TEST 2']
+

tests/sonic_package_manager/test_service_creator.py

@@ -157,7 +157,7 @@ def test_service_creator_yang(sonic_fs, manifest, mock_sonic_db,
     })
 
     entry = PackageEntry('test', 'azure/sonic-test')
-    package = Package(entry, Metadata(manifest, yang_module_str=test_yang))
+    package = Package(entry, Metadata(manifest, yang_modules=[test_yang]))
     service_creator.create(package)
 
     mock_config_mgmt.add_module.assert_called_with(test_yang)
@@ -171,7 +171,7 @@ def test_service_creator_yang(sonic_fs, manifest, mock_sonic_db,
             },
         },
     }
-    package = Package(entry, Metadata(manifest, yang_module_str=test_yang))
+    package = Package(entry, Metadata(manifest, yang_modules=[test_yang]))
 
     service_creator.create(package)
 
@@ -190,6 +190,42 @@ def test_service_creator_yang(sonic_fs, manifest, mock_sonic_db,
     mock_config_mgmt.remove_module.assert_called_with(test_yang_module)
 
 
+def test_service_creator_multi_yang(sonic_fs, manifest, mock_config_mgmt, service_creator):
+    test_yang = 'TEST YANG'
+    test_yang_2 = 'TEST YANG 2'
+
+    def get_module_name(module_src):
+        if module_src == test_yang:
+            return 'sonic-test'
+        elif module_src == test_yang_2:
+            return 'sonic-test-2'
+        else:
+            raise ValueError(f'Unknown module {module_src}')
+
+    entry = PackageEntry('test', 'azure/sonic-test')
+    package = Package(entry, Metadata(manifest, yang_modules=[test_yang, test_yang_2]))
+    service_creator.create(package)
+
+    mock_config_mgmt.add_module.assert_has_calls(
+        [
+            call(test_yang),
+            call(test_yang_2)
+        ],
+        any_order=True,
+    )
+
+    mock_config_mgmt.get_module_name = Mock(side_effect=get_module_name)
+
+    service_creator.remove(package)
+    mock_config_mgmt.remove_module.assert_has_calls(
+        [
+            call(get_module_name(test_yang)),
+            call(get_module_name(test_yang_2))
+        ],
+        any_order=True,
+    )
+
+
 def test_service_creator_autocli(sonic_fs, manifest, mock_cli_gen,
                                  mock_config_mgmt, service_creator):
     test_yang = 'TEST YANG'
@@ -199,7 +235,7 @@ def test_service_creator_autocli(sonic_fs, manifest, mock_cli_gen,
     manifest['cli']['auto-generate-config'] = True
 
     entry = PackageEntry('test', 'azure/sonic-test')
-    package = Package(entry, Metadata(manifest, yang_module_str=test_yang))
+    package = Package(entry, Metadata(manifest, yang_modules=[test_yang]))
     mock_config_mgmt.get_module_name = Mock(return_value=test_yang_module)
     service_creator.create(package)
 
@@ -226,6 +262,58 @@ def test_service_creator_post_operation_hook(sonic_fs, manifest, mock_sonic_db,
             service_creator._post_operation_hook()
             run_command.assert_called_with(['systemctl', 'daemon-reload'])
 
+def test_service_creator_multi_yang_filter_auto_cli_modules(sonic_fs, manifest, mock_cli_gen,
+                                                            mock_config_mgmt, service_creator):
+    test_yang = 'TEST YANG'
+    test_yang_2 = 'TEST YANG 2'
+    test_yang_3 = 'TEST YANG 3'
+    test_yang_4 = 'TEST YANG 4'
+
+    def get_module_name(module_src):
+        if module_src == test_yang:
+            return 'sonic-test'
+        elif module_src == test_yang_2:
+            return 'sonic-test-2'
+        elif module_src == test_yang_3:
+            return 'sonic-test-3'
+        elif module_src == test_yang_4:
+            return 'sonic-test-4'
+        else:
+            raise ValueError(f'Unknown module {module_src}')
+
+    manifest['cli']['auto-generate-show'] = True
+    manifest['cli']['auto-generate-config'] = True
+    manifest['cli']['auto-generate-show-source-yang-modules'] = ['sonic-test-2', 'sonic-test-4']
+    manifest['cli']['auto-generate-config-source-yang-modules'] = ['sonic-test-2', 'sonic-test-4']
+
+    entry = PackageEntry('test', 'azure/sonic-test')
+    package = Package(entry, Metadata(manifest, yang_modules=[test_yang, test_yang_2, test_yang_3, test_yang_4]))
+    mock_config_mgmt.get_module_name = Mock(side_effect=get_module_name)
+    service_creator.create(package)
+
+    assert mock_cli_gen.generate_cli_plugin.call_count == 4
+    mock_cli_gen.generate_cli_plugin.assert_has_calls(
+        [
+            call('show', get_module_name(test_yang_2)),
+            call('show', get_module_name(test_yang_4)),
+            call('config', get_module_name(test_yang_2)),
+            call('config', get_module_name(test_yang_4)),
+        ],
+        any_order=True
+    )
+
+    service_creator.remove(package)
+    assert mock_cli_gen.remove_cli_plugin.call_count == 4
+    mock_cli_gen.remove_cli_plugin.assert_has_calls(
+        [
+            call('show', get_module_name(test_yang_2)),
+            call('show', get_module_name(test_yang_4)),
+            call('config', get_module_name(test_yang_2)),
+            call('config', get_module_name(test_yang_4)),
+        ],
+        any_order=True
+    )
+
 def test_feature_registration(mock_sonic_db, manifest):
     mock_connector = Mock()
     mock_connector.get_entry = Mock(return_value={})