# -*- coding: utf-8 -*- """ PostgreSQL client """ import datetime import logging import sys import psycopg2 log = logging.getLogger(__name__) class PgDB: """ PostgreSQL client """ host = "" user = "" pwd = "" db = "" con = 0 date_format = '%Y-%m-%d' datetime_format = '%Y-%m-%d %H:%M:%S' def __init__(self, host, user, pwd, db, just_try=False): self.host = host self.user = user self.pwd = pwd self.db = db self.just_try = just_try def connect(self): """ Connect to PostgreSQL server """ if self.con == 0: try: con = psycopg2.connect( dbname=self.db, user=self.user, host=self.host, password=self.pwd ) self.con = con except Exception: logging.fatal( 'An error occured during Postgresql database connection (%s@%s, database=%s).', self.user, self.host, self.db, exc_info=1 ) sys.exit(1) return True def close(self): """ Close connection with PostgreSQL server (if opened) """ if self.con: self.con.close() def setEncoding(self, enc): """ Set connection encoding """ if self.con: try: self.con.set_client_encoding(enc) return True except Exception: log.error('An error occured setting Postgresql database connection encoding to "%s"', enc, exc_info=1) return False def doSQL(self, sql, params=None): """ Run SELECT SQL query and return result as dict """ if self.just_try: log.debug(u"Just-try mode : do not really execute SQL query '%s'", sql) return True cursor = self.con.cursor() try: if params is None: cursor.execute(sql) else: cursor.execute(sql, params) self.con.commit() return True except Exception: log.error(u'Error during SQL request "%s"', sql.decode('utf-8', 'ignore'), exc_info=1) self.con.rollback() return False def doSelect(self, sql): """ Run SELECT SQL query and return result as dict """ cursor = self.con.cursor() try: cursor.execute(sql) results = cursor.fetchall() return results except Exception: log.error(u'Error during SQL request "%s"', sql.decode('utf-8', 'ignore'), exc_info=1) return False # # SQL helpers # def _quote_value(self, value): """ Quote a value for SQL query """ if isinstance(value, (int, float)): return str(value) if isinstance(value, datetime.datetime): value = self._format_datetime(value) elif isinstance(value, datetime.date): value = self._format_date(value) return u"'%s'" % value.replace(u"'", u"''") def _format_where_clauses(self, where_clauses, where_op=u'AND'): """ Format WHERE clauses """ if isinstance(where_clauses, str): return where_clauses if isinstance(where_clauses, list): return (u" %s " % where_op).join(where_clauses) if isinstance(where_clauses, dict): return (u" %s " % where_op).join(map(lambda x: "%s=%s" % (x, self._quote_value(where_clauses[x])), where_clauses)) log.error('Unsupported where clauses type %s', type(where_clauses)) return False def _format_datetime(self, value): """ Format datetime object as string """ assert isinstance(value, datetime.datetime) return value.strftime(self.datetime_format) def _format_date(self, value): """ Format date object as string """ assert isinstance(value, (datetime.date, datetime.datetime)) return value.strftime(self.date_format) def time2datetime(self, time): """ Convert timestamp to datetime string """ return self._format_datetime(datetime.datetime.fromtimestamp(int(time))) def time2date(self, time): """ Convert timestamp to date string """ return self._format_date(datetime.date.fromtimestamp(int(time))) def insert(self, table, values, just_try=False): """ Run INSERT SQL query """ sql = u"INSERT INTO %s (%s) VALUES (%s)" % (table, u', '.join(values.keys()), u", ".join(map(lambda x: self._quote_value(values[x]), values))) if just_try: log.debug(u"Just-try mode : execute INSERT query : %s", sql) return True log.debug(sql) if not self.doSQL(sql): log.error(u"Fail to execute INSERT query (SQL : %s)", sql) return False return True def update(self, table, values, where_clauses, where_op=u'AND', just_try=False): """ Run UPDATE SQL query """ where = self._format_where_clauses(where_clauses, where_op=where_op) if not where: return False sql = u"UPDATE %s SET %s WHERE %s" % (table, u", ".join(map(lambda x: "%s=%s" % (x, self._quote_value(values[x])), values)), where) if just_try: log.debug(u"Just-try mode : execute UPDATE query : %s", sql) return True log.debug(sql) if not self.doSQL(sql): log.error(u"Fail to execute UPDATE query (SQL : %s)", sql) return False return True def delete(self, table, where_clauses, where_op=u'AND', just_try=False): """ Run DELETE SQL query """ where = self._format_where_clauses(where_clauses, where_op=where_op) if not where: return False sql = u"DELETE FROM %s WHERE %s" % (table, where) if just_try: log.debug(u"Just-try mode : execute DELETE query : %s", sql) return True log.debug(sql) if not self.doSQL(sql): log.error(u"Fail to execute DELETE query (SQL : %s)", sql) return False return True def select(self, table, where_clauses=None, fields=None, where_op=u'AND', order_by=None): """ Run SELECT SQL query """ sql = u"SELECT " if fields is None: sql += "*" elif isinstance(fields, str): sql += fields else: sql += u", ".join(fields) sql += u" FROM " + table if where_clauses: where = self._format_where_clauses(where_clauses, where_op=where_op) if not where: return False sql += u" WHERE " + where if order_by: sql += u"ORDER %s" % order_by return self.doSelect(sql)