diff --git a/cosmo/clients/netbox_v4.py b/cosmo/clients/netbox_v4.py index 45eaa6f..09a82b2 100644 --- a/cosmo/clients/netbox_v4.py +++ b/cosmo/clients/netbox_v4.py @@ -2,10 +2,12 @@ from abc import ABC, abstractmethod from builtins import map from multiprocessing import Manager -from string import Template +from os import PathLike +from pathlib import Path from cosmo.clients import get_client_mp_context from cosmo.clients.netbox_client import NetboxAPIClient +from cosmo.common import FileTemplate class ParallelQuery(ABC): @@ -16,6 +18,10 @@ def __init__(self, client: NetboxAPIClient, **kwargs): self.data_promise = None self.kwargs = kwargs + @staticmethod + def file_template(relpath: str | PathLike): + return FileTemplate(Path(__file__).parent.joinpath(Path(relpath))) + def fetch_data(self, pool): return pool.apply_async(self._fetch_data, args=(self.kwargs, pool)) @@ -43,43 +49,7 @@ def _fetch_data(self, kwargs, pool): if self.netbox_43_query_syntax else 'tag: "bgp_cpe"' ) - query_template = Template( - """ - query { - interface_list(filters: { $tag_filter }) { - __typename - id, - parent { - __typename - id, - connected_endpoints { - ... on InterfaceType { - __typename - name - device { - name - __typename - primary_ip4 { - __typename - address - } - interfaces { - id - name - __typename - ip_addresses { - __typename - address - } - } - } - } - } - } - } - } - """ - ) + query_template = self.file_template("queries/connected_devices.graphql") return self.client.query( query_template.substitute(tag_filter=tag_filter), "connected_devices_query" @@ -123,49 +93,7 @@ def _fetch_data(self, kwargs, pool): # Note: This does not use the device list, because we can have other participating devices # which are not in the same repository and thus are not appearing in device list. - query_template = Template( - """ - query{ - interface_list(filters: { - name: {starts_with: "lo"} - }) { - __typename - name, - child_interfaces { - __typename - name, - vrf { - __typename - id - name - description - rd - export_targets { - __typename - name - } - import_targets { - __typename - name - } - }, - ip_addresses { - __typename - address, - family { - __typename - value, - } - } - } - device{ - __typename - name, - } - } - } - """ - ) + query_template = self.file_template("queries/loopback.graphql") return self.client.query(query_template.substitute(), "loopback_query")["data"] @@ -205,74 +133,7 @@ def _merge_into(self, data: dict, query_data): class L2VPNDataQuery(ParallelQuery): def _fetch_data(self, kwargs, pool): - query_template = Template( - """ - query { - l2vpn_list (filters: {name: {starts_with: "WAN: "}}) { - __typename - id - name - type - identifier - terminations { - __typename - id - assigned_object { - __typename - ... on VLANType { - __typename - id - name - interfaces_as_tagged { - id - name - __typename - device { - __typename - id - name - } - } - interfaces_as_untagged { - id - name - __typename - device { - __typename - id - name - } - } - } - ... on InterfaceType { - __typename - id - name - custom_fields - untagged_vlan { - __typename - id - name - vid - } - tagged_vlans { - __typename - id - name - vid - } - device { - __typename - id - name - } - } - } - } - } - } - """ - ) + query_template = self.file_template("queries/l2vpn.graphql") return self.client.query(query_template.substitute(), "l2vpn_query")["data"] @@ -428,197 +289,7 @@ def __init__(self, *args, multiple_mac_addresses=False, **kwargs): def _fetch_data(self, kwargs, pool): device = kwargs.get("device") - query_template = Template( - """ - query { - device_list(filters: { - name: { i_exact: $device }, - }) { - __typename - id - name - custom_fields - - device_type { - __typename - slug - } - platform { - __typename - manufacturer { - __typename - slug - } - slug - } - primary_ip4 { - __typename - address - } - - interfaces { - __typename - id - name - enabled - type - mode - mtu - description - connected_endpoints { - ... on ProviderNetworkType { - __typename - display - } - ... on CircuitTerminationType { - __typename - display - } - ... on VirtualCircuitTerminationType { - __typename - display - } - ... on InterfaceType { - __typename - name - device { - __typename - name - } - } - ... on FrontPortType { - __typename - name - device { - __typename - name - } - } - ... on RearPortType { - __typename - name - device { - __typename - name - } - } - ... on ConsolePortType { - __typename - name - device { - __typename - name - } - } - ... on ConsoleServerPortType { - __typename - name - device { - __typename - name - } - } - } - link_peers { - ... on CircuitTerminationType { - __typename - display - } - ... on FrontPortType { - __typename - name - device { - __typename - name - } - } - ... on RearPortType { - __typename - name - device { - __typename - name - } - } - ... on ConsolePortType { - __typename - name - device { - __typename - name - } - } - ... on ConsoleServerPortType { - __typename - name - device { - __typename - name - } - } - ... on InterfaceType { - __typename - name - device { - __typename - name - } - } - } - vrf { - __typename - id - name - description - rd - export_targets { - __typename - name - } - import_targets { - __typename - name - } - } - lag { - __typename - id - name - } - ip_addresses { - __typename - address - role - } - untagged_vlan { - __typename - id - name - vid - } - tagged_vlans { - __typename - id - name - vid - } - tags { - __typename - id - name - slug - } - parent { - __typename - id - mtu - name - } - custom_fields - } - } - }""" - ) + query_template = self.file_template("queries/device.graphql") query = query_template.substitute( device=json.dumps(device), diff --git a/cosmo/clients/queries/connected_devices.graphql b/cosmo/clients/queries/connected_devices.graphql new file mode 100644 index 0000000..05f9af2 --- /dev/null +++ b/cosmo/clients/queries/connected_devices.graphql @@ -0,0 +1,33 @@ +query { + interface_list(filters: { $tag_filter }) { + __typename + id, + parent { + __typename + id, + connected_endpoints { + ... on InterfaceType { + __typename + name + device { + name + __typename + primary_ip4 { + __typename + address + } + interfaces { + id + name + __typename + ip_addresses { + __typename + address + } + } + } + } + } + } + } +} \ No newline at end of file diff --git a/cosmo/clients/queries/device.graphql b/cosmo/clients/queries/device.graphql new file mode 100644 index 0000000..40de4f7 --- /dev/null +++ b/cosmo/clients/queries/device.graphql @@ -0,0 +1,188 @@ +query { + device_list(filters: { + name: { i_exact: $device }, + }) { + __typename + id + name + custom_fields + + device_type { + __typename + slug + } + platform { + __typename + manufacturer { + __typename + slug + } + slug + } + primary_ip4 { + __typename + address + } + + interfaces { + __typename + id + name + enabled + type + mode + mtu + description + connected_endpoints { + ... on ProviderNetworkType { + __typename + display + } + ... on CircuitTerminationType { + __typename + display + } + ... on VirtualCircuitTerminationType { + __typename + display + } + ... on InterfaceType { + __typename + name + device { + __typename + name + } + } + ... on FrontPortType { + __typename + name + device { + __typename + name + } + } + ... on RearPortType { + __typename + name + device { + __typename + name + } + } + ... on ConsolePortType { + __typename + name + device { + __typename + name + } + } + ... on ConsoleServerPortType { + __typename + name + device { + __typename + name + } + } + } + link_peers { + ... on CircuitTerminationType { + __typename + display + } + ... on FrontPortType { + __typename + name + device { + __typename + name + } + } + ... on RearPortType { + __typename + name + device { + __typename + name + } + } + ... on ConsolePortType { + __typename + name + device { + __typename + name + } + } + ... on ConsoleServerPortType { + __typename + name + device { + __typename + name + } + } + ... on InterfaceType { + __typename + name + device { + __typename + name + } + } + } + vrf { + __typename + id + name + description + rd + export_targets { + __typename + name + } + import_targets { + __typename + name + } + } + lag { + __typename + id + name + } + ip_addresses { + __typename + address + role + } + untagged_vlan { + __typename + id + name + vid + } + tagged_vlans { + __typename + id + name + vid + } + tags { + __typename + id + name + slug + } + parent { + __typename + id + mtu + name + } + custom_fields + } + } +} \ No newline at end of file diff --git a/cosmo/clients/queries/l2vpn.graphql b/cosmo/clients/queries/l2vpn.graphql new file mode 100644 index 0000000..1b1fbb4 --- /dev/null +++ b/cosmo/clients/queries/l2vpn.graphql @@ -0,0 +1,64 @@ +query { + l2vpn_list (filters: {name: {starts_with: "WAN: "}}) { + __typename + id + name + type + identifier + terminations { + __typename + id + assigned_object { + __typename + ... on VLANType { + __typename + id + name + interfaces_as_tagged { + id + name + __typename + device { + __typename + id + name + } + } + interfaces_as_untagged { + id + name + __typename + device { + __typename + id + name + } + } + } + ... on InterfaceType { + __typename + id + name + custom_fields + untagged_vlan { + __typename + id + name + vid + } + tagged_vlans { + __typename + id + name + vid + } + device { + __typename + id + name + } + } + } + } + } +} \ No newline at end of file diff --git a/cosmo/clients/queries/loopback.graphql b/cosmo/clients/queries/loopback.graphql new file mode 100644 index 0000000..b65f4f8 --- /dev/null +++ b/cosmo/clients/queries/loopback.graphql @@ -0,0 +1,39 @@ +query{ + interface_list(filters: { + name: {starts_with: "lo"} + }) { + __typename + name, + child_interfaces { + __typename + name, + vrf { + __typename + id + name + description + rd + export_targets { + __typename + name + } + import_targets { + __typename + name + } + }, + ip_addresses { + __typename + address, + family { + __typename + value, + } + } + } + device{ + __typename + name, + } + } +} \ No newline at end of file diff --git a/cosmo/common.py b/cosmo/common.py index c5a164a..77ca4b5 100644 --- a/cosmo/common.py +++ b/cosmo/common.py @@ -1,3 +1,6 @@ +from os import PathLike +from pathlib import Path +from string import Template from abc import ABC, abstractmethod from typing import Optional, Union, Protocol, TypeVar, Sequence @@ -70,3 +73,10 @@ def without_keys(d, keys) -> dict: if type(keys) != list: keys = [keys] return {k: v for k, v in d.items() if k not in keys} + + +class FileTemplate(Template): + def __init__(self, template_file_path: str | bytes | PathLike): + with open(template_file_path, "r") as template_file: + template = template_file.read() + super().__init__(template) diff --git a/cosmo/features.py b/cosmo/features.py index 2d53a03..bbc80d2 100644 --- a/cosmo/features.py +++ b/cosmo/features.py @@ -1,5 +1,6 @@ # implementation guide # https://martinfowler.com/articles/feature-toggles.html +import functools from argparse import Action, ArgumentParser from typing import Never, Self, Optional, TextIO, Sequence, Any, Callable @@ -83,11 +84,14 @@ def __str__(self): return ", ".join(features_desc) -def with_feature(instance: FeatureToggle, feature_name: str): +def _feature_toggler_decorator_gen( + instance: FeatureToggle, feature_name: str, target_state: bool +): def decorator_with_feature(func: Callable): + @functools.wraps(func) def exe_with_feature(*args, **kwargs): previous_state = instance.featureIsEnabled(feature_name) - instance.setFeature(feature_name, True) + instance.setFeature(feature_name, target_state) func(*args, **kwargs) instance.setFeature(feature_name, previous_state) @@ -96,6 +100,14 @@ def exe_with_feature(*args, **kwargs): return decorator_with_feature +def with_feature(instance: FeatureToggle, feature_name: str): + return _feature_toggler_decorator_gen(instance, feature_name, True) + + +def without_feature(instance: FeatureToggle, feature_name: str): + return _feature_toggler_decorator_gen(instance, feature_name, False) + + features = FeatureToggle( { "interface-auto-descriptions": True, diff --git a/cosmo/tests/test_features.py b/cosmo/tests/test_features.py index 642fc82..6dc7016 100644 --- a/cosmo/tests/test_features.py +++ b/cosmo/tests/test_features.py @@ -7,6 +7,7 @@ NonExistingFeatureToggleException, FeatureToggle, with_feature, + without_feature, ) @@ -72,6 +73,17 @@ def execute_with_decorator(): assert not ft.featureIsEnabled("feature_a") +def test_without_feature_decorator(): + ft = FeatureToggle({"feature_a": True}) + + @without_feature(ft, "feature_a") + def execute_with_decorator(): + assert not ft.featureIsEnabled("feature_a") + + execute_with_decorator() + assert ft.featureIsEnabled("feature_a") + + def test_argparse_integration(): ft = FeatureToggle({"feature_a": False, "feature_b": False, "feature_c": True})