annet 0.11__py3-none-any.whl → 0.12__py3-none-any.whl

annet/cli_args.py

@@ -3,6 +3,7 @@
 import abc
 import argparse
 import enum
+import logging
 import os
 
 from valkit.common import valid_string_list
@@ -360,15 +361,14 @@ class QueryOptionsBase(CacheOptions):
     def __init__(self, *args, **kwargs):
         super().__init__(*args, **kwargs)
         if not isinstance(self.query, Query):
-            query_type = storage_connector.get().query()
-            self.query = query_type.new(self.query, hosts_range=self.hosts_range)
-
-    def validate_stdin(self, arg, val, **kwargs):
-        if "storage" in kwargs and arg == "config":
-            storage = kwargs["storage"]
-            if len(storage.resolve_object_ids_by_query(self.query)) > 1:
-                raise ValueError("stdin config can not be used with multiple devices")
-        super().validate_stdin(arg, val, **kwargs)
+            connectors = storage_connector.get_all()
+            if not connectors:
+                pass
+            elif len(connectors) == 1:
+                query_type = connectors[0].query()
+                self.query = query_type.new(self.query, hosts_range=self.hosts_range)
+            else:
+                logging.warning("Multiple connectors found, skip parsing query")
 
 
 class QueryOptions(QueryOptionsBase):

annet/connectors.py

@@ -2,8 +2,7 @@ import sys
 from abc import ABC
 from functools import cached_property
 from importlib.metadata import entry_points
-from typing import Generic, Optional, Type, TypeVar
-
+from typing import Generic, Optional, Type, TypeVar, List
 
 T = TypeVar("T")
 
@@ -12,29 +11,46 @@ class Connector(ABC, Generic[T]):
     name: str
     ep_name: str
     ep_group: str = "annet.connectors"
-    _cls: Optional[Type[T]] = None
+    _classes: Optional[List[Type[T]]] = None
 
     def _get_default(self) -> Type[T]:
         raise RuntimeError(f"{self.name} is not set")
 
     @cached_property
-    def _entry_point(self) -> Optional[Type[T]]:
+    def _entry_point(self) -> List[Type[T]]:
         return load_entry_point(self.ep_group, self.ep_name)
 
     def get(self, *args, **kwargs) -> T:
-        if self._cls is not None:
-            res = self._cls
-        else:
-            res = self._entry_point or self._get_default()
+        if self._classes is None:
+            self._classes = self._entry_point or [self._get_default()]
+        if len(self._classes) > 1:
+            raise RuntimeError(
+                f"Multiple classes are registered with the same "
+                f"group={self.ep_group} and name={self.ep_name}: "
+                f"{[cls for cls in self._classes]}",
+            )
+
+        res = self._classes[0]
         return res(*args, **kwargs)
 
+    def get_all(self, *args, **kwargs) -> T:
+        if self._classes is None:
+            self._classes = self._entry_point or [self._get_default()]
+
+        return [cls(*args, **kwargs) for cls in self._classes]
+
     def set(self, cls: Type[T]):
-        if self._cls is not None:
+        if self._classes is not None:
+            raise RuntimeError(f"Cannot reinitialize value of {self.name}")
+        self._classes = [cls]
+
+    def set_all(self, classes: List[Type[T]]):
+        if self._classes is not None:
             raise RuntimeError(f"Cannot reinitialize value of {self.name}")
-        self._cls = cls
+        self._classes = list(classes)
 
     def is_default(self) -> bool:
-        return self._cls is self._entry_point is None
+        return self._classes is self._entry_point is None
 
 
 class CachedConnector(Connector[T], ABC):
@@ -58,7 +74,4 @@ def load_entry_point(group: str, name: str):
         ep = entry_points(group=group, name=name)  # pylint: disable=unexpected-keyword-arg
     if not ep:
         return None
-    if len(ep) > 1:
-        raise RuntimeError(f"Multiple entry points with the same {group=} and {name=}: {[item.value for item in ep]}")
-    for item in ep:
-        return item.load()
+    return [item.load() for item in ep]

annet/gen.py

@@ -39,6 +39,7 @@ from annet.generators import (
     JSONFragment,
     NotSupportedDevice,
     PartialGenerator,
+    RefGenerator,
 )
 from annet.lib import merge_dicts, percentile
 from annet.output import output_driver_connector
@@ -60,13 +61,16 @@ class DeviceGenerators:
     """Collections of various types of generators found for devices."""
 
     # map device fqdn to found partial generators
-    partial: Dict[str, List[PartialGenerator]] = dataclasses.field(default_factory=dict)
+    partial: Dict[Any, List[PartialGenerator]] = dataclasses.field(default_factory=dict)
+
+    # ref generators
+    ref: Dict[Any, List[RefGenerator]] = dataclasses.field(default_factory=dict)
 
     # map device fqdn to found entire generators
-    entire: Dict[str, List[Entire]] = dataclasses.field(default_factory=dict)
+    entire: Dict[Any, List[Entire]] = dataclasses.field(default_factory=dict)
 
     # map device fqdn to found json fragment generators
-    json_fragment: Dict[str, List[JSONFragment]] = dataclasses.field(default_factory=dict)
+    json_fragment: Dict[Any, List[JSONFragment]] = dataclasses.field(default_factory=dict)
 
     def iter_gens(self) -> Iterator[BaseGenerator]:
         """Iterate over generators."""
@@ -75,19 +79,24 @@ class DeviceGenerators:
                 for gen in gen_list:
                     yield gen
 
-    def file_gens(self, device_fqdn: str) -> Iterator[Union[Entire, JSONFragment]]:
+    def file_gens(self, device: Any) -> Iterator[Union[Entire, JSONFragment]]:
         """Iterate over generators that generate files or file parts."""
         yield from itertools.chain(
-            self.entire.get(device_fqdn, []),
-            self.json_fragment.get(device_fqdn, []),
+            self.entire.get(device, []),
+            self.json_fragment.get(device, []),
         )
 
+    def update(self, other: "DeviceGenerators") -> None:
+        self.partial.update(other.partial)
+        self.ref.update(other.ref)
+        self.entire.update(other.entire)
+        self.json_fragment.update(other.json_fragment)
+
 
 @dataclasses.dataclass
 class OldNewDeviceContext:
     config: str
     args: GenOptions
-    storage: Storage
     downloaded_files: Dict[Device, DeviceDownloadedFiles]
     failed_files: Dict[Device, Exception]
     running: Dict[Device, Dict[str, str]]
@@ -143,21 +152,24 @@ def _old_new_per_device(ctx: OldNewDeviceContext, device: Device, filterer: Filt
                 splitter=tabparser.make_formatter(device.hw).split,
             )
         if not old:
-            res = generators.run_partial_initial(device, ctx.storage)
+            res = generators.run_partial_initial(device)
             old = res.config_tree()
             perf = res.perf_mesures()
             if ctx.args.profile and ctx.do_print_perf:
                 _print_perf("INITIAL", perf)
         run_args = generators.GeneratorPartialRunArgs(
             device=device,
-            storage=ctx.storage,
             use_acl=not ctx.args.no_acl,
             use_acl_safe=ctx.args.acl_safe,
             annotate=ctx.add_annotations,
             generators_context=ctx.args.generators_context,
             no_new=ctx.no_new,
         )
-        res = generators.run_partial_generators(ctx.gens.partial[device.fqdn], run_args)
+        res = generators.run_partial_generators(
+            ctx.gens.partial[device],
+            ctx.gens.ref[device],
+            run_args,
+        )
         partial_results = res.partial_results
         perf = res.perf_mesures()
         if ctx.no_new:
@@ -227,9 +239,8 @@ def _old_new_per_device(ctx: OldNewDeviceContext, device: Device, filterer: Filt
                     get_logger(host=device.hostname).error(error_msg)
                     return OldNewResult(device=device, err=Exception(error_msg))
         res = generators.run_file_generators(
-            ctx.gens.file_gens(device.fqdn),
+            ctx.gens.file_gens(device),
             device,
-            ctx.storage,
         )
 
         entire_results = res.entire_results
@@ -324,7 +335,7 @@ def split_downloaded_files(
     """Split downloaded files per generator type: entire/json_fragment."""
     ret = DeviceDownloadedFiles()
 
-    for gen in gens.file_gens(device.fqdn):
+    for gen in gens.file_gens(device):
         filepath = gen.path(device)
         if filepath in device_flat_files:
             if isinstance(gen, Entire):
@@ -355,7 +366,6 @@ def split_downloaded_files_multi_device(
 @tracing.function
 def old_new(
     args: GenOptions,
-    storage: Storage,
     config: str,
     loader: "Loader",
     filterer: Filterer,
@@ -367,13 +377,13 @@ def old_new(
     do_files_download=False,
     do_print_perf=True,
 ):
-    devices = loader.resolve_devices(storage, device_ids)
-    gens = loader.resolve_gens(storage, devices)
+    devices = loader.devices
+    gens = loader.resolve_gens(devices)
     running, failed_running = _old_resolve_running(config, devices)
     downloaded_files, failed_files = _old_resolve_files(config, devices, gens, do_files_download)
 
     if stdin is None:
-        stdin = args.stdin(storage=storage, filter_acl=args.filter_acl, config=config)
+        stdin = args.stdin(filter_acl=args.filter_acl, config=config)
 
     fetched_packages, failed_packages = {}, {}
     if do_files_download and config == "running":
@@ -385,7 +395,6 @@ def old_new(
     ctx = OldNewDeviceContext(
         config=config,
         args=args,
-        storage=storage,
         downloaded_files=split_downloaded_files_multi_device(downloaded_files, gens, devices),
         failed_files=failed_files,
         running=running,
@@ -413,26 +422,25 @@ def old_new(
 
 
 @tracing.function
-def old_raw(args: GenOptions, storage, config, stdin=None,
-            do_files_download=False, use_mesh=True,
-            ) -> Iterable[Tuple[Device, Union[str, Dict[str, str]]]]:
-    devices = storage.make_devices(args.query, preload_neighbors=True, use_mesh=use_mesh)
-    device_gens = _old_resolve_gens(args, storage, devices)
-    running, failed_running = _old_resolve_running(config, devices)
-    downloaded_files, failed_files = _old_resolve_files(config, devices, device_gens, do_files_download)
+def old_raw(
+        args: GenOptions, loader: Loader, config, stdin=None,
+        do_files_download=False, use_mesh=True,
+) -> Iterable[Tuple[Device, Union[str, Dict[str, str]]]]:
+    device_gens = loader.resolve_gens(loader.devices)
+    running, failed_running = _old_resolve_running(config, loader.devices)
+    downloaded_files, failed_files = _old_resolve_files(config, loader.devices, device_gens, do_files_download)
     if stdin is None:
-        stdin = args.stdin(storage=storage, filter_acl=args.filter_acl, config=config)
+        stdin = args.stdin(filter_acl=args.filter_acl, config=config)
     ctx = OldNewDeviceContext(
         config=config,
         args=args,
-        storage=storage,
-        downloaded_files=split_downloaded_files_multi_device(downloaded_files, device_gens, devices),
+        downloaded_files=split_downloaded_files_multi_device(downloaded_files, device_gens, loader.devices),
         failed_files=failed_files,
         running=running,
         failed_running=failed_running,
         stdin=stdin,
         do_files_download=do_files_download,
-        device_count=len(devices),
+        device_count=len(loader.devices),
         no_new=True,
         add_annotations=False,
         add_implicit=False,
@@ -441,7 +449,7 @@ def old_raw(args: GenOptions, storage, config, stdin=None,
         failed_packages={},
         do_print_perf=True,
     )
-    for device in devices:
+    for device in loader.devices:
         if not device.is_pc():
             config = _old_new_get_config_cli(ctx, device)
             config = scrub_config(config, device.breed)
@@ -463,68 +471,59 @@ def worker(device_id, args: ShowGenOptions, stdin, loader: "Loader", filterer: F
     if span:
         span.set_attribute("device.id", device_id)
 
-    connector = storage_connector.get()
-    storage_opts = connector.opts().from_cli_opts(args)
-    with connector.storage()(storage_opts) as storage:
-        for res in old_new(
-            args,
-            storage,
-            config="/dev/null",
-            loader=loader,
-            filterer=filterer,
-            add_implicit=False,
-            add_annotations=args.annotate,
-            stdin=stdin,
-            device_ids=[device_id],
-        ):
-            new = res.get_new(args.acl_safe)
-            new_files = res.get_new_files(args.acl_safe)
-            new_file_fragments = res.get_new_file_fragments(args.acl_safe)
-            output_driver = output_driver_connector.get()
-            device = res.device
-            if new is None:
-                continue
-            for (entire_path, (entire_data, _)) in sorted(new_files.items(), key=itemgetter(0)):
-                yield (output_driver.entire_config_dest_path(device, entire_path), entire_data, False)
-
-            for (path, (data, _)) in sorted(new_file_fragments.items(), key=itemgetter(0)):
-                dumped_data = json.dumps(data, indent=4, sort_keys=True, ensure_ascii=False)
-                yield (output_driver.entire_config_dest_path(device, path), dumped_data, False)
-
-            has_file_result = new_files or new_file_fragments
-            has_partial_result = new or not has_file_result
-            if device.hw.vendor in platform.VENDOR_REVERSES and has_partial_result:
-                orderer = patching.Orderer.from_hw(device.hw)
-                yield (output_driver.cfg_file_names(device)[0],
-                       format_config_blocks(
-                           orderer.order_config(new),
-                           device.hw,
-                           args.indent
-                       ),
-                       False)
+    for res in old_new(
+        args,
+        config="/dev/null",
+        loader=loader,
+        filterer=filterer,
+        add_implicit=False,
+        add_annotations=args.annotate,
+        stdin=stdin,
+        device_ids=[device_id],
+    ):
+        new = res.get_new(args.acl_safe)
+        new_files = res.get_new_files(args.acl_safe)
+        new_file_fragments = res.get_new_file_fragments(args.acl_safe)
+        output_driver = output_driver_connector.get()
+        device = res.device
+        if new is None:
+            continue
+        for (entire_path, (entire_data, _)) in sorted(new_files.items(), key=itemgetter(0)):
+            yield (output_driver.entire_config_dest_path(device, entire_path), entire_data, False)
+
+        for (path, (data, _)) in sorted(new_file_fragments.items(), key=itemgetter(0)):
+            dumped_data = json.dumps(data, indent=4, sort_keys=True, ensure_ascii=False)
+            yield (output_driver.entire_config_dest_path(device, path), dumped_data, False)
+
+        has_file_result = new_files or new_file_fragments
+        has_partial_result = new or not has_file_result
+        if device.hw.vendor in platform.VENDOR_REVERSES and has_partial_result:
+            orderer = patching.Orderer.from_hw(device.hw)
+            yield (output_driver.cfg_file_names(device)[0],
+                   format_config_blocks(
+                       orderer.order_config(new),
+                       device.hw,
+                       args.indent
+                   ),
+                   False)
 
 
 def old_new_worker(device_id, args: DeployOptions, config, stdin, loader: "Loader", filterer: Filterer):
-    connector = storage_connector.get()
-    storage_opts = connector.opts().from_cli_opts(args)
-    with connector.storage()(storage_opts) as storage:
-        yield from old_new(
-            args,
-            storage,
-            config=config,
-            loader=loader,
-            filterer=filterer,
-            stdin=stdin,
-            device_ids=[device_id],
-            no_new=args.clear,
-            do_files_download=True,
-        )
+    yield from old_new(
+        args,
+        config=config,
+        loader=loader,
+        filterer=filterer,
+        stdin=stdin,
+        device_ids=[device_id],
+        no_new=args.clear,
+        do_files_download=True,
+    )
 
 
 class OldNewParallel(Parallel):
-    def __init__(self, storage: Storage, args: DeployOptions, loader: "Loader", filterer: Filterer):
-        self.storage = storage
-        stdin = args.stdin(storage=storage, filter_acl=args.filter_acl, config=args.config)
+    def __init__(self, args: DeployOptions, loader: "Loader", filterer: Filterer):
+        stdin = args.stdin(filter_acl=args.filter_acl, config=args.config)
         super().__init__(
             old_new_worker,
             args,
@@ -535,20 +534,19 @@ class OldNewParallel(Parallel):
         )
         self.tune_args(args)
 
-    def generated_configs(self, device_ids: List[int]) -> Generator[OldNewResult, None, None]:
-        skipped = set(device_ids)
+    def generated_configs(self, devices: List[Device]) -> Generator[OldNewResult, None, None]:
+        devices_by_id = {device.id: device for device in devices}
+        device_ids = list(devices_by_id)
 
         for task_result in self.irun(device_ids):
             if task_result.exc is not None:
-                device = self.storage.get_device(task_result.device_id, use_mesh=False, preload_neighbors=False)
+                device = devices_by_id.pop(task_result.device_id)
                 yield OldNewResult(device=device, err=task_result.exc)
-                skipped.discard(task_result.device_id)
             elif task_result.result is not None:
                 yield from task_result.result
-                skipped.discard(task_result.device_id)
+                devices_by_id.pop(task_result.device_id)
 
-        for device_id in skipped:
-            device = self.storage.get_device(device_id, use_mesh=False, preload_neighbors=False)
+        for device in devices_by_id.values():
             yield OldNewResult(device=device, err=Exception(f"No config returned for {device.hostname}"))
 
 
@@ -564,7 +562,7 @@ def _get_files_to_download(devices: List[Device], gens: DeviceGenerators) -> Dic
     for device in devices:
         paths = set()
         try:
-            for generator in gens.file_gens(device.fqdn):
+            for generator in gens.file_gens(device):
                 try:
                     path = generator.path(device)
                     if path:
@@ -713,6 +711,8 @@ def _old_new_get_config_cli(ctx: OldNewDeviceContext, device: Device) -> str:
             raise exc
     elif ctx.config == "-":
         text = ctx.stdin["config"]
+        if ctx.device_count > 1:
+            raise ValueError("stdin config can not be used with multiple devices")
     else:
         if os.path.isdir(ctx.config):
             filename = _existing_cfg_file_name(ctx.config, device)
@@ -766,13 +766,15 @@ def _old_new_get_config_files(ctx: OldNewDeviceContext, device: Device) -> Devic
 
 
 @tracing.function
-def _old_resolve_gens(args: GenOptions, storage: Storage, devices: List[Device]) -> DeviceGenerators:
+def _old_resolve_gens(args: GenOptions, storage: Storage, devices: Iterable[Device]) -> DeviceGenerators:
     per_device_gens = DeviceGenerators()
+    devices = devices or [None]  # get all generators if no devices provided
     for device in devices:
         gens = generators.build_generators(storage, gens=args, device=device)
-        per_device_gens.partial[device.fqdn] = gens.partial
-        per_device_gens.entire[device.fqdn] = gens.entire
-        per_device_gens.json_fragment[device.fqdn] = gens.json_fragment
+        per_device_gens.partial[device] = gens.partial
+        per_device_gens.entire[device] = gens.entire
+        per_device_gens.json_fragment[device] = gens.json_fragment
+        per_device_gens.ref[device] = gens.ref
     return per_device_gens
 
 
@@ -808,37 +810,46 @@ def _old_resolve_files(config: str,
 
 
 class Loader:
-    def __init__(self, storage: Storage, args: GenOptions, no_empty_warning: bool = False) -> None:
+    def __init__(
+            self, *storages: Storage,
+            args: GenOptions,
+            no_empty_warning: bool = False,
+    ) -> None:
         self._args = args
-        self._storage = storage
+        self._storages = storages
         self._no_empty_warning = no_empty_warning
-        self._devices_map: Optional[Dict[int, Device]] = None
-        self._gens: Optional[DeviceGenerators] = None
+        self._devices_map: Dict[int, Device] = {}
+        self._gens: DeviceGenerators = DeviceGenerators()
+        self._counter = itertools.count()
 
         self._preload()
 
     def _preload(self) -> None:
         with tracing_connector.get().start_as_current_span("Resolve devices"):
-            devices = self._storage.make_devices(
-                self._args.query,
-                preload_neighbors=True,
-                use_mesh=not self._args.no_mesh,
-                preload_extra_fields=True,
-            )
+            for storage in self._storages:
+                devices = storage.make_devices(
+                    self._args.query,
+                    preload_neighbors=True,
+                    use_mesh=not self._args.no_mesh,
+                    preload_extra_fields=True,
+                )
+                for device in devices:
+                    self._devices_map[next(self._counter)] = device
+                self._gens.update(_old_resolve_gens(self._args, storage, devices))
         if not devices and not self._no_empty_warning:
             get_logger().error("No devices found for %s", self._args.query)
             return
 
-        self._devices_map = {d.id: d for d in devices}
-        self._gens = _old_resolve_gens(self._args, self._storage, devices)
-
     @property
     def device_fqdns(self):
-        return {d.id: d.fqdn for d in self._devices_map.values()} if self._devices_map else {}
+        return {
+            device_id: d.fqdn
+            for device_id, d in self._devices_map.items()
+        }
 
     @property
     def device_ids(self):
-        return list(self.device_fqdns)
+        return list(self._devices_map)
 
     @property
     def devices(self) -> List[Device]:
@@ -846,23 +857,9 @@ class Loader:
             return list(self._devices_map.values())
         return []
 
-    def resolve_devices(self, storage: Storage, device_ids: Iterable[int]) -> List[Device]:
-        devices = []
-        for device_id in device_ids:
-            device = self._devices_map[device_id]
-
-            # can not use self._storage here, we can be in another process
-            device.storage = storage
-
-            devices.append(device)
-        return devices
-
-    def resolve_gens(self, storage: Storage, devices: Iterable[Device]) -> DeviceGenerators:
+    def resolve_gens(self, devices: Iterable[Device]) -> DeviceGenerators:
         if self._gens is not None:
-            for gen in self._gens.iter_gens():
-                # can not use self._storage here, we can be in another process
-                gen.storage = storage
             return self._gens
 
         with tracing_connector.get().start_as_current_span("Resolve gens"):
-            return _old_resolve_gens(self._args, storage, devices)
+            return _old_resolve_gens(self._args, self._storage, devices)