make object types configurable

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

2 files changed, 28 insertions, 13 deletions

diff --git a/config.py.sample b/config.py.sample
index 05a0650..9e3ba44 100644
--- a/config.py.sample
+++ b/config.py.sample
@@ -6,6 +6,13 @@
 #   Defaults to "www.peeringdb.com",
 #   but can eg. be set to "beta.peeringdb.com" for testing
 #
+# object_types: array of object types to sync
+#   defaults to finding them through the machine-readable
+#   API documentaion, but if that is too slow (parsing a big YAML file),
+#   then you can just set it statically.
+#   This is also useful if you only want a subset of object types:
+#   Example: 'object_types': [ 'org', 'net', 'poc' ]
+#
 # db_conn_str: libpq database connection string
 #   https://www.postgresql.org/docs/current/libpq-connect.html#LIBPQ-CONNSTRING
 #
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: