ecn-robots: reimplemented pg robot with psycopg2 and solved whitespace issues

pgdb.py

@@ -1,121 +1,125 @@
 'Python interface to PostgreSQL for managing database nodes.'
 
 from subprocess import Popen, PIPE
+import psycopg2
+import psycopg2.extensions
 
 class Db(object):
     
     def __init__(self,name):
         self._name = name
-        self._p = Popen('/usr/bin/psql -Aqt -U postgres', stdin=PIPE, stdout=PIPE, stderr=PIPE, shell=True)
+        self.conn = psycopg2.connect("dbname='postgres' user='postgres' host='localhost' password=''");
+        self.conn.set_isolation_level(psycopg2.extensions.ISOLATION_LEVEL_AUTOCOMMIT)
+        self.cursor = self.conn.cursor()
 
     def _run(self, command,args=[]):
         execfn = [command] + list(args)
-	try:
+        try:
             p = Popen(execfn, stdout=PIPE, stderr=PIPE)
-	    return p.communicate()
+            return p.communicate()
         except Exception,e:
-	    print str(e)
-	return -1
+            print str(e)
+        return -1
 
     def _runsql(self, sql, db='postgres'):
-	given_sql = sql
-        sql = sql +'\echo this is the end' +'\n'
-        self._p.stdout.flush()
-        self._p.stdin.write(str(sql))
-	out = ""
-	out  = self._p.stdout.readline()
-	while out.find("this is the end") == -1:
-	    out = out +self._p.stdout.readline()
-        lout = out.replace('this is the end','')
-	return lout.strip()
+        given_sql = sql
+        self.cursor.execute(given_sql)
+        try:
+            out = self.cursor.fetchall()
+        except Exception,e: 
+            out = "Exception " +str(e)+" ocured"
+        return out
     
     def _get_owner(self):
         sql = "SELECT pg_get_userbyid(datdba) FROM pg_database WHERE datname ='"+self.sanitize(self._name)+"';"
-	own = self._runsql(sql)
+        own = self._runsql(sql)
         return own
 
     def _set_owner(self, owner):
         sql = "ALTER DATABASE "+self._name+" OWNER TO "+self.sanitize(owner)+";"
         own = self._runsql(sql)
-	return own
+        return own
     
     db_owner = property(_get_owner, _set_owner)
 
     @property
     def db_OID(self):
         sql = "SELECT oid FROM pg_database WHERE datname = '"+self.sanitize(self._name)+"';"
-	oid = self._runsql(sql)
-	return oid
+        oid = self._runsql(sql)
+        return oid
 
     @property
     def db_info(self):
         information = {'size':'', 'encoding':'', 'collation':'','ctype':''}
         information['size'] = self._runsql("SELECT pg_size_pretty(pg_database_size('"+self.sanitize(self._name)+"'));")
-        information['encoding'], information['collation'], \
-	    information['ctype'] = self._runsql("SELECT pg_encoding_to_char(encoding), datcollate, datctype FROM pg_database WHERE datname='"+self.sanitize(self._name)+"';").split('|')
+        inf = self._runsql("""SELECT pg_encoding_to_char(encoding), 
+                        datcollate, datctype FROM pg_database WHERE datname='"""+self._name+"';")
+        information['encoding'] = inf[0][0]
+        information['collation'] = inf[0][1]
+        information['ctype'] = inf[0][2]
         return information
 
     @property
     def db_connections(self):
         sql = "SELECT numbackends from pg_stat_database WHERE datname = '"+self.sanitize(self._name)+"';"
-	cncs = self._runsql(sql)
-	return cncs
+        cncs = self._runsql(sql)
+        return cncs
     
     def user_exists(self, user):
         sql = "SELECT rolname FROM pg_authid WHERE rolname = '"+self.sanitize(user)+"';"
-	u = self._runsql(sql)
-	if (u == ""):
-	    return False
-	return True
+        u = self._runsql(sql)
+        if (len(u) == 0):
+            return False
+        return True
 
     def db_exists(self, xdb):
         sql = "SELECT datname FROM pg_database WHERE datname = '"+self.sanitize(xdb)+"';"
-	d = self._runsql(sql)
-        if (d == ""):
-	    return False
-	return True
+        d = self._runsql(sql)
+        if (len(d) == 0):
+            return False
+        return True
 
     def db_delete(self):
         if self.db_exists(self._name) == True:
             sql = "DROP DATABASE "+self.sanitize(self._name)+";"
-	    drop = self._runsql(sql)
-	    return drop
+            drop = self._runsql(sql)
+            return drop
         return "Failed"
 
     def db_create(self, own, coll, ctyp, enc=u'UTF8'):
         if self.db_exists(self._name) == False:
-	    sql = "CREATE DATABASE "+self.sanitize(self._name)+" WITH OWNER = "+self.sanitize(own)+" ENCODING = '"+self.sanitize(enc)+"' LC_COLLATE = '"+self.sanitize(coll)+"' LC_CTYPE = '"+self.sanitize(ctyp)+"';"
-	    create = self._runsql(sql)
-	    return create
+            sql = "CREATE DATABASE "+self.sanitize(self._name)+" WITH OWNER = "+self.sanitize(own)+" ENCODING = '"+self.sanitize(enc)+"' LC_COLLATE = '"+self.sanitize(coll)+"' LC_CTYPE = '"+self.sanitize(ctyp)+"';"
+            create = self._runsql(sql)
+            return create
         return "Failed"
 
     def db_dump(self, path, method):
         dump = Popen(['/usr/bin/pg_dump', '-U','postgres','-F'+ method, self._name], stdout=PIPE)
-	fl = open(path,"wb")
-	gz = Popen(['gzip'], stdin = dump.stdout, stdout = fl)
-	fl.close
-	return "Finished dumping "+self._name
+        fl = open(path,"wb")
+        gz = Popen(['gzip'], stdin = dump.stdout, stdout = fl)
+        fl.close
+        return "Finished dumping "+self._name
 
 
     def db_rename(self,old, new):
         if self.db_exists(new) == True or self.db_exists(old) == False:
-	    return "Cannot"
-	sql = "ALTER DATABASE "+self.sanitize(old)+" RENAME TO "+self.sanitize(new)+";"
-	rename = self._runsql(sql)
-	return rename
+            return "Cannot"
+        sql = "ALTER DATABASE "+self.sanitize(old)+" RENAME TO "+self.sanitize(new)+";"
+        rename = self._runsql(sql)
+        return rename
 
     def copy():
         pass
 
     def db_list(self):
         sql = "SELECT datname FROM pg_database WHERE datname NOT IN ('template0', 'template1', 'postgres');"
-	dbl = self._runsql(sql)
-	return dbl
+        dbl = self._runsql(sql)
+        return dbl
 
     def usr_list(self):
         sql = "SELECT rolname FROM pg_authid WHERE rolcanlogin=true;"
-	usrl = self._runsql(sql)
-	return usrl
+        usrl = self._runsql(sql)
+        return usrl
 
     def usr_add(self, us, passw):
         sql = "CREATE ROLE "+self.sanitize(us)+" WITH NOSUPERUSER NOCREATEDB NOCREATEROLE NOCREATEUSER LOGIN PASSWORD '"+self.sanitize(passw)+"';"