diff --git a/pgdb.py b/pgdb.py index a96cb97..9d0476c 100644 --- a/pgdb.py +++ b/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)+"';"