This repository has been archived by the owner on Aug 6, 2024. It is now read-only.
-
Notifications
You must be signed in to change notification settings - Fork 0
/
csvdump.py
executable file
·106 lines (103 loc) · 4.96 KB
/
csvdump.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
import time
from datetime import datetime
import pymongo.errors
from pymongo.mongo_client import MongoClient
import argparse
import csv
import sys
import os
parser = argparse.ArgumentParser(prog='csvdump.py',
description='A program to visualize the planes collected by airstrik.py',
epilog='Go Pack!')
parser.add_argument('-d', '--database', help='which database in the mongodb to pull from', required=False)
parser.add_argument('-u', '--uri', default='mongodb://localhost:27017', help='The URI to connect to (mongodb)')
parser.add_argument('-o', '--out', default='out.csv', help="The file to output to")
parser.add_argument('-s', '--stats', action='store_true', help='whether to comb for stats or not')
parser.add_argument('-n', '--no-warn', action='store_true', help='warn about deleting files or not')
args = parser.parse_args()
progress_ht_num = 30
print("Connecting to MongoDB...")
client = MongoClient(args.uri)
client.admin.command('ping')
all_databases = client.list_database_names()
if (args.database not in all_databases) or (args.database in ['admin', 'config']) or args.database is None:
print("ERROR: You need to provide a valid database name (-d). Here's a list of those databases:")
for item in all_databases:
if item not in ['admin', 'config']:
print(item)
sys.exit(0)
db = client[args.database]
colnames = db.list_collection_names()
if args.out in os.listdir():
if not args.no_warn:
cont = input("The file " + args.out + ' already exists! Would you like to delete it? (y/n)')
if cont not in ['y', 'yes']:
sys.exit(0)
os.remove(args.out)
if args.stats:
fieldnames = ['date', 'unique_planes', 'total_trips', 'unique_alarm_planes', 'total_alarm_trips']
with open(args.out, 'x', newline='') as csvfile:
writer = csv.DictWriter(csvfile, fieldnames=fieldnames)
writer.writeheader()
print()
lc = [i for i in colnames if i.startswith('stats.')]
for cur, item in enumerate(lc):
pct = (cur + 1) / len(lc) * progress_ht_num
print("Writing", item,
"(" + ("#" * int(pct)) + "." * int(progress_ht_num - pct) + ") (" + str(cur + 1) + "/" + str(
len(lc)) + ")", flush=True)
sys.stdout.flush()
split_item = item.split('.')
dat = list(db['stats'][split_item[1]].find())
for i, data in enumerate(dat):
write_dict = {'date': split_item[1]}
for it in data.keys():
if it == '_id':
continue
write_dict.update({it: data[it]})
writer.writerow(write_dict)
else:
fieldnames = ['name', 'flight_id', 'start_time', 'end_time', 'lat', 'lon', 'nav_heading', 'alt_geom',
'calc_heading',
'calc_speed', 'distance', 'trip', 'filters']
with open(args.out, 'x', newline='') as csvfile:
writer = csv.DictWriter(csvfile, fieldnames=fieldnames)
writer.writeheader()
print()
lc = [i for i in colnames if not i.startswith('stats')]
for cur, item in enumerate(lc):
pct = (cur + 1) / len(lc) * progress_ht_num
print("Writing", item,
"(" + ("#" * int(pct)) + "." * int(progress_ht_num - pct) + ") (" + str(cur + 1) + "/" + str(
len(lc)) + ")", flush=True)
sys.stdout.flush()
dat = list(db[item].find())
for i, data in enumerate(dat):
flight_name = data['flight_name_id']
if flight_name is not None:
flight_name = flight_name[0]
write_dict = {'name': item, 'flight_id': flight_name, 'trip': i + 1}
for it in data.keys():
if it not in ['_id', 'alarm', 'extras', 'flight_name_id', 'filters']:
try:
write_dict.update({it: data[it][0]})
except TypeError:
continue
for it in data['extras'].keys():
if it not in ['alarm_triggered', 'commentary']: # commentary for legacy db
write_dict.update({it: datetime.fromtimestamp(data['extras'][it])})
ftext = ''
try:
for each_filter in data['filters'].keys():
each_filter_data = data['filters'][each_filter]
ftext += (each_filter + ' (' + str(each_filter_data['dist']) + ', ' + str(each_filter_data['alt']) + '), ')
ftext = ftext[:-2]
write_dict.update({'filters': ftext})
except KeyError:
print('this instance is not compatible with filters!')
write_dict.update({'filters': ''})
writer.writerow(write_dict)
print("Dumped to", args.out)
print("Stopping MongoDB...")
client.close()
print('done')