Merge pull request UCL-ShippingGroup#29 from willu47/bug_count

Changed status to approximate count of rows to allow scaling to large tables

pyrate/repositories/aisdb.py

@@ -136,7 +136,7 @@ def __init__(self, options, readonly=False):
         self.dirty = sql.Table(self, 'ais_dirty', self.dirty_db_spec['cols'],
                                self.dirty_db_spec['indices'])
         self.sources = sql.Table(self, 'ais_sources', self.sources_db_spec['cols'])
-        self.imolist = sql.Table(self, 'imo_list', self.imolist_db_spec['cols'], 
+        self.imolist = sql.Table(self, 'imo_list', self.imolist_db_spec['cols'],
                                  constraint=self.imolist_db_spec['constraint'])
         self.extended = AISExtendedTable(self)
         self.clean_imolist = sql.Table(self, 'imo_list_clean', self.clean_imo_list['cols'], constraint=self.clean_imo_list['constraint'])
@@ -215,7 +215,7 @@ def get_message_stream(self, mmsi, from_ts=None, to_ts=None, use_clean_db=False,
         if not to_ts is None:
             where.append("time <= %s")
             params.append(to_ts)
-        
+
         cols_list = ','.join([c[0].lower() for c in db.cols])
         where_clause = ' AND '.join(where)
         sql = "SELECT {} FROM {} WHERE {} ORDER BY time ASC".format(cols_list,
@@ -245,10 +245,10 @@ def get_message_stream(self, mmsi, from_ts=None, to_ts=None, use_clean_db=False,
 class AISExtendedTable(sql.Table):
 
     def __init__(self, db):
-        super(AISExtendedTable, self).__init__(db, 'ais_extended', 
+        super(AISExtendedTable, self).__init__(db, 'ais_extended',
             AISdb.clean_db_spec['cols'] + [('location', 'geography(POINT, 4326)')],
             AISdb.clean_db_spec['indices'])
-    
+
     def create(self):
         with self.db.conn.cursor() as cur:
             cur.execute("CREATE EXTENSION IF NOT EXISTS postgis")
@@ -262,7 +262,7 @@ def create(self):
                             RETURN NEW;
                         END;
                         ' LANGUAGE plpgsql;
-                        CREATE TRIGGER {0}_gis_insert 
+                        CREATE TRIGGER {0}_gis_insert
                         BEFORE INSERT OR UPDATE ON {0} FOR EACH ROW EXECUTE PROCEDURE location_insert();
                         """.format(self.name))
             except psycopg2.ProgrammingError:

pyrate/repositories/sql.py

@@ -1,5 +1,17 @@
-import psycopg2
+"""Classes for connection to and management of database tables
+
+PgsqlRepository
+---------------
+Sets up a connection to a pyrate database repository
+
+Table
+-----
+Used to encapsulate a pyrate database table
+
+"""
 import logging
+import psycopg2
+
 
 def load(options, readonly=False):
     return PgsqlRepository(options)
@@ -19,7 +31,7 @@ def __init__(self, options, readonly=False):
         self.conn = None
 
     def connection(self):
-        return psycopg2.connect(host=self.host, database=self.db, user=self.user, password=self.password)
+        return psycopg2.connect(host=self.host, database=self.db, user=self.user, password=self.password, connect_timeout=3)
 
     def __enter__(self):
         self.conn = self.connection()
@@ -28,6 +40,8 @@ def __exit__(self, exc_type, exc_value, traceback):
         self.conn.close()
 
 class Table(object):
+    """A database table
+    """
 
     def __init__(self, db, name, cols, indices=None, constraint=None,
                  foreign_keys=None):
@@ -101,13 +115,18 @@ def truncate(self):
             self.db.conn.commit()
 
     def status(self):
-        """ Returns the number of records in the table
+        """ Returns the approximate number of records in the table
+
+        Returns
+        -------
+        integer
+
         """
         with self.db.conn.cursor() as cur:
             try:
-                cur.execute("SELECT COUNT(*) FROM \""+ self.name +"\"")
+                cur.execute("SELECT reltuples FROM pg_class WHERE oid = %s::regclass::oid", [self.name])
                 self.db.conn.commit()
-                return cur.fetchone()[0]
+                return int(cur.fetchone()[0])
             except psycopg2.ProgrammingError:
                 self.db.conn.rollback()
                 return -1