Merge pull request #45 from NickLaMuro/rds_support

Adds support for using pg_activity with Amazon RDS

pg_activity

@@ -119,6 +119,13 @@ server activity monitoring.")
                 help = "Filesystem blocksize (default: 4096)",
                 metavar = 'BLOCKSIZE',
                 default = 4096)
+            # --rds
+            parser.add_option(
+                '--rds',
+                dest = 'rds',
+                action = 'store_true',
+                help = "Enable support for AWS RDS",
+                default = 'false')
             group = OptionGroup(
                 parser,
                 "Display Options, you can exclude some columns by using them ")
@@ -232,7 +239,8 @@ server activity monitoring.")
                     port = options.port,
                     user = options.username,
                     password = password,
-                    database = options.dbname)
+                    database = options.dbname,
+                    rds_mode = options.rds)
                 break
             except psycopg2.OperationalError as err:
                 if nb_try < 1 and (err.pgcode == errorcodes.INVALID_PASSWORD or
@@ -278,7 +286,7 @@ server activity monitoring.")
         disp_procs = None
         delta_disk_io = None
         # get DB informations
-        db_info = PGAUI.data.pg_get_db_info(None)
+        db_info = PGAUI.data.pg_get_db_info(None, using_rds = options.rds)
         PGAUI.set_max_db_length(db_info['max_length'])
         # indentation
         indent = PGAUI.get_indent(flag)
@@ -305,7 +313,7 @@ server activity monitoring.")
                 delta_disk_io = PGAUI.data.get_global_io_counters()
             procs = new_procs
             # refresh the winodw
-            db_info = PGAUI.data.pg_get_db_info(db_info)
+            db_info = PGAUI.data.pg_get_db_info(db_info, using_rds = options.rds)
             PGAUI.set_max_db_length(db_info['max_length'])
             # bufferize
             PGAUI.set_buffer({

pgactivity/Data.py

@@ -134,7 +134,8 @@ def pg_connect(self,
         port = 5432,
         user = 'postgres',
         password = None,
-        database = 'postgres'):
+        database = 'postgres',
+        rds_mode = False):
         """
         Connect to a PostgreSQL server and return
         cursor & connector.
@@ -163,11 +164,12 @@ def pg_connect(self,
                 connection_factory = psycopg2.extras.DictConnection
                 )
         self.pg_conn.set_isolation_level(0)
-        cur = self.pg_conn.cursor()
-        cur.execute("SELECT current_setting('is_superuser')")
-        ret = cur.fetchone()
-        if ret[0] != "on":
-            raise Exception("Must be run with database superuser privileges.")
+        if rds_mode != True: # Make sure we are using superuser if not on RDS
+          cur = self.pg_conn.cursor()
+          cur.execute("SELECT current_setting('is_superuser')")
+          ret = cur.fetchone()
+          if ret[0] != "on":
+              raise Exception("Must be run with database superuser privileges.")
 
     def pg_is_local_access(self,):
         """
@@ -256,7 +258,7 @@ def pg_get_num_dev_version(self, text_version):
             return
         raise Exception('Undefined PostgreSQL version.')
 
-    def pg_get_db_info(self, prev_db_infos):
+    def pg_get_db_info(self, prev_db_infos, using_rds = False):
         """
         Get current sum of transactions, total size and  timestamp.
         """
@@ -269,6 +271,7 @@ def pg_get_db_info(self, prev_db_infos):
     FROM
         pg_database
         """
+        query += "\nWHERE datname <> 'rdsadmin'" if using_rds else ''
         cur = self.pg_conn.cursor()
         cur.execute(query,)
         ret = cur.fetchone()