add postgres flags and initialization

Signed-off-by: Jakub Sokołowski <jakub@status.im>

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',
+    parser.add_option('-u', '--db-user', default='postgres',
-                      help='Name of the field to count.')
+                      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()

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