diff --git a/main.py b/main.py index 8e05460..5d5caee 100755 --- a/main.py +++ b/main.py @@ -3,7 +3,7 @@ from os import path from optparse import OptionParser from query import ESQueryPeers -from graph import PDGraphPeers +from postgres import PGDatabase HELP_DESCRIPTION = 'This generates a CSV with buckets of peer_ids for every day.' HELP_EXAMPLE = 'Example: ./unique_count.py -i "logstash-2019.11.*" -f peer_id' @@ -11,16 +11,20 @@ HELP_EXAMPLE = 'Example: ./unique_count.py -i "logstash-2019.11.*" -f peer_id' def parse_opts(): parser = OptionParser(description=HELP_DESCRIPTION, epilog=HELP_EXAMPLE) - parser.add_option('-H', '--es-host', dest='es_host', default='localhost', + parser.add_option('-H', '--es-host', default='localhost', help='ElasticSearch host.') - parser.add_option('-P', '--es-port', dest='es_port', default=9200, + parser.add_option('-P', '--es-port', default=9200, help='ElasticSearch port.') - parser.add_option('-h', '--db-host', dest='db_host', default='localhost', + parser.add_option('-d', '--db-host', default='localhost', help='PostgreSQL host.') - parser.add_option('-p', '--db-port', dest='db_port', default=5432, + parser.add_option('-b', '--db-port', default=5432, help='PostgreSQL port.') - parser.add_option('-f', '--field', type='str', default='peer_id', - help='Name of the field to count.') + parser.add_option('-u', '--db-user', default='postgres', + help='PostgreSQL user.') + parser.add_option('-p', '--db-pass', default='postgres', + help='PostgreSQL password.') + parser.add_option('-n', '--db-name', default='postgres', + help='PostgreSQL database name.') parser.add_option('-i', '--index-pattern', default='logstash-*', help='Patter for matching indices.') parser.add_option('-f', '--field', type='str', default='peer_id', @@ -38,13 +42,26 @@ def parse_opts(): def main(): (opts, args) = parse_opts() - esq = ESQueryPeers(opts.es_host, opts.es_port) + esq = ESQueryPeers( + opts.es_host, + opts.es_port + ) + psg = PGDatabase( + opts.db_name, + opts.db_user, + opts.db_pass, + opts.db_host, + opts.db_port + ) data = [] for index in esq.get_indices(opts.index_pattern): print('Index: {}'.format(index)) data.extend(esq.get_peers(index, opts.field, opts.max_size)) + rval = psg.get_most_recent_day() + print(rval) + if __name__ == '__main__': main() diff --git a/postgres.py b/postgres.py new file mode 100644 index 0000000..19d38c6 --- /dev/null +++ b/postgres.py @@ -0,0 +1,25 @@ +import psycopg2 + +class PGDatabase: + _SCHEMA = """ + CREATE TABLE IF NOT EXISTS peers (date date, peer varchar(64), count int); + """ + + def __init__(self, name, user, password=None, host='localhost', port=5432): + self.db = psycopg2.connect( + user = user, + password = password, + host = host, + port = port, + database = name + ) + self.c = self.db.cursor() + self._create_schema() + + def _create_schema(self): + self.c.execute(self._SCHEMA) + self.db.commit() + + def get_most_recent_day(self): + rval = self.c.execute('SELECT date FROM peers ORDER BY date LIMIT 1;') + return rval