make object types configurable

Signed-off-by: Asbjørn Sloth Tønnesen <ast@a5n.dk>

1 files changed, 21 insertions, 13 deletions

diff --git a/sync.py b/sync.py
index 5543265..2f9e126 100755
--- a/sync.py
+++ b/sync.py
@@ -157,6 +157,22 @@ def find_spec(s, url):
 
     return urljoin(url, m[2])
 
+def find_object_types_via_apidocs(s):
+    ret = []
+    spec_url = find_spec(s, f'https://{domain_name:s}/apidocs/')
+    s.headers.update({'Accept': 'application/x-yaml'})
+    r = s.get(spec_url)
+    ignored_types = [ 'as_set' ]
+    apidoc = yaml.safe_load(r.text)
+    p = re.compile(r'^/api/([a-z_]+)$')
+    for path in apidoc['paths']:
+        m = p.match(path)
+        if not m: continue
+        key = m[1]
+        if key in ignored_types: continue
+        ret.append(key)
+    return ret
+
 def handle_auth(s):
     auth = get_config().get('auth')
     if type(auth) == str:
@@ -179,23 +195,15 @@ def main():
     req_agent = s.headers.get('User-Agent')
     s.headers.update({'User-Agent': f'peeringdb-simplesync/0.1 {req_agent:s}'})
 
-    spec_url = find_spec(s, f'https://{domain_name:s}/apidocs/')
-
-    s.headers.update({'Accept': 'application/x-yaml'})
-    r = s.get(spec_url)
+    keys = get_config().get('object_types')
+    if not keys:
+        keys = find_object_types_via_apidocs(s)
+        print('Consider setting \'object_types\':', keys)
 
     # subsequent requests are going to be JSON
     s.headers.update({'Accept': 'application/json'})
 
-    ignored = [ 'as_set' ]
-
-    apidoc = yaml.safe_load(r.text)
-    p = re.compile(r'^/api/([a-z_]+)$')
-    for path in apidoc['paths']:
-        m = p.match(path)
-        if not m: continue
-        key = m[1]
-        if key in ignored: continue
+    for key in keys:
         try:
             sync_table(s, key)
         except AssertionError: