Merge pull request #36 from wobcom/feature/performance2

Splitted, multi-threaded and rewamped some queries to speed up perfor…

Author: johannwagner

.gitignore

@@ -6,8 +6,7 @@ cosmo.yml
 .direnv
 leaf*.json
 
-cosmo_data.yaml
-machines/test0001/*
+machines/**
 
 .coverage
 *~

cosmo/__main__.py

@@ -1,4 +1,3 @@
-import ipaddress
 import json
 import os
 import sys
@@ -8,12 +7,11 @@
 import yaml
 import argparse
 
-from cosmo.netboxclient import NetboxClient
+from cosmo.clients.netbox import NetboxClient
+from cosmo.log import info
 from cosmo.serializer import RouterSerializer, SwitchSerializer, AbstractRecoverableError, RouterSerializerConfig
 
 
-def info(string: str) -> None:
-    print("[INFO] " + string)
 
 
 def main() -> int:
@@ -81,7 +79,7 @@ def noop(*args, **kwargs):
             if device['name'] in cosmo_configuration['devices']['router']:
 
                 router_serializer_cfg = RouterSerializerConfig(cosmo_configuration.get("router_serializer_configuration", {}))
-                serializer = RouterSerializer(router_serializer_cfg, device, cosmo_data['l2vpn_list'], cosmo_data["vrf_list"])
+                serializer = RouterSerializer(router_serializer_cfg, device, cosmo_data['l2vpn_list'], cosmo_data["vrf_list"], cosmo_data["loopbacks"])
                 content = serializer.serialize()
             elif device['name'] in cosmo_configuration['devices']['switch']:
                 serializer = SwitchSerializer(device)

cosmo/clients/__init__.py

@@ -0,0 +1,47 @@
+import time
+from urllib.parse import urljoin
+
+import requests
+
+from cosmo import log
+from cosmo.clients.netbox_v4 import NetboxV4Strategy
+
+
+class NetboxClient:
+    def __init__(self, url, token):
+        self.url = url
+        self.token = token
+
+        self.version = self.query_version()
+
+        if self.version.startswith("4."):
+            log.info("Using version 4.x strategy...")
+            self.child_client = NetboxV4Strategy(url, token)
+        else:
+            raise Exception("Unknown Version")
+
+    def query_version(self):
+        r = requests.get(
+            urljoin(self.url, "/api/status/"),
+            headers={
+                "Authorization": f"Token {self.token}",
+                "Content-Type": "application/json",
+                "Accept": "application/json",
+            },
+        )
+        if r.status_code != 200:
+            raise Exception("Error querying api: " + r.text)
+
+        json = r.json()
+        return json['netbox-version']
+
+    def get_data(self, device_config):
+        start_time = time.perf_counter()
+        data = self.child_client.get_data(device_config)
+        end_time = time.perf_counter()
+        diff_time = end_time - start_time
+        log.info(f"Data fetching took {round(diff_time, 2)} s...")
+
+        return data
+
+

cosmo/clients/netbox.py

@@ -0,0 +1,56 @@
+from urllib.parse import urlencode, urljoin
+
+import requests
+
+
+class NetboxAPIClient:
+    def __init__(self, url, token):
+        self.url = url
+        self.token = token
+
+    def query(self, query):
+        r = requests.post(
+            urljoin(self.url, "/graphql/"),
+            json={"query": query},
+            headers={
+                "Authorization": f"Token {self.token}",
+                "Content-Type": "application/json",
+                "Accept": "application/json",
+            },
+        )
+        if r.status_code != 200:
+            raise Exception("Error querying api: " + r.text)
+
+        json = r.json()
+
+        if 'errors' in json:
+            for e in json['errors']:
+                print(e)
+
+        return json
+
+    def query_rest(self, path, queries):
+        q = urlencode(queries, doseq=True)
+        url = urljoin(self.url, path) + f"?{q}"
+
+        return_array = list()
+
+        while url is not None:
+            r = requests.get(
+                url,
+                headers={
+                    "Authorization": f"Token {self.token}",
+                    "Content-Type": "application/json",
+                    "Accept": "application/json",
+                },
+            )
+
+            if r.status_code != 200:
+                raise Exception("Error querying api: " + r.text)
+
+            data = r.json()
+
+            url = data['next']
+            return_array.extend(data['results'])
+
+        return return_array