From 90cbce9550088a4099f1856330a75e678b4bf4fe Mon Sep 17 00:00:00 2001 From: Sebastian Lohff Date: Sun, 14 Jun 2020 03:36:42 +0200 Subject: [PATCH] Improve conv tool to convert all old graphs --- bin/conv.py | 43 +++++++++++++++++++++++++++++++------------ 1 file changed, 31 insertions(+), 12 deletions(-) diff --git a/bin/conv.py b/bin/conv.py index 526d71b..36c4545 100755 --- a/bin/conv.py +++ b/bin/conv.py @@ -11,23 +11,42 @@ from backend.crawler import convert_crawl, net_to_json from bgpdata.models import CrawlRun -def main(): - parser = argparse.ArgumentParser() - parser.add_argument("-c", "--crawl-id", type=int) - args = parser.parse_args() - - try: - crawl = CrawlRun.objects.get(pk=args.crawl_id) - except CrawlRun.DoesNotExist: - parser.error("CrawlRun with id {} does not exist".format(args.crawl_id)) - +def _convert_crawl(crawl): net = convert_crawl(crawl) if net.nodes and net.edges: crawl.graph = net_to_json(net) crawl.save() - print("Crawl updated") + print("Crawl {} updated".format(crawl.id)) else: - print("Crawl had no nodes or edges, abort") + print("Crawl {} had no nodes or edges, abort".format(crawl.id)) + + +def main(): + parser = argparse.ArgumentParser() + parser.add_argument("-c", "--crawl-id", type=int) + parser.add_argument("-a", "--all", default=False, action="store_true") + parser.add_argument("-e", "--empty-graph-only", default=False, action="store_true") + args = parser.parse_args() + + if args.crawl_id and args.all: + parser.error("-c and -a don't work together") + + if args.crawl_id: + try: + crawl = CrawlRun.objects.get(pk=args.crawl_id) + except CrawlRun.DoesNotExist: + parser.error("CrawlRun with id {} does not exist".format(args.crawl_id)) + _convert_crawl(crawl) + elif args.all: + if args.empty_graph_only: + crawls = CrawlRun.objects.filter(graph='') + else: + crawls = CrawlRun.objects.all() + + for crawl in crawls: + _convert_crawl(crawl) + else: + parser.error("Either specify a crawl with -c or use -a for all") if __name__ == '__main__':