python-mylib/PgDB.py

#!/usr/bin/python
# -*- coding: utf-8 -*-

import psycopg2
import logging
import sys
import traceback
import datetime

class PgDB(object):

    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):
        self.host = host
        self.user = user
        self.pwd  = pwd
        self.db   = db

    def connect(self):
        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):
        if self.con:
            self.con.close()

    def setEncoding(self,enc):
        if self.con:
            try:
                self.con.set_client_encoding(enc)
                return True
            except Exception:
                logging.error('An error occured setting Postgresql database connection encoding to "%s"', enc, exc_info=1)
        return False

    def doSQL(self,sql,params=None):
        cursor = self.con.cursor()
        try:
            if params is None:
                cursor.execute(sql)
            else:
                cursor.execute(sql,params)
            self.con.commit()
            return True
        except Exception:
            logging.error(u'Error during SQL request "%s"', sql.decode('utf-8', 'ignore'), exc_info=1)
            self.con.rollback()
            return False

    def doSelect(self,sql,params):
        cursor = self.con.cursor()
        try:
            if params is None:
                cursor.execute(sql)
            else:
                cursor.execute(sql,params)
            results = cursor.fetchall()
            return results
        except Exception:
            logging.error(u'Error during SQL request "%s"', sql.decode('utf-8', 'ignore'), exc_info=1)
            return False

    #
    # SQL helpers
    #
    def _quote_value(self, value):
        if isinstance(value, int) or isinstance(value, float):
            return unicode(value)

        if isinstance(value, str):
            value = unicode(value)
        elif isinstance(value, datetime.datetime):
            value = unicode(self._format_datetime(value))
        elif isinstance(value, datetime.date):
            value = unicode(self._format_date(value))

        return u"'%s'" % value.replace(u"'",u"''")

    def _format_where_clauses(self, where_clauses, where_op=u'AND'):
        if isinstance(where_clauses, str):
            return where_clauses
        elif isinstance(where_clauses, list):
            return (u" %s " % where_op).join(where_clauses)
        elif isinstance(where_clauses, dict):
            return (u" %s " % where_op).join(map(lambda x: "%s=%s" % (x, self._quote_value(where_clauses[x])), where_clauses))
        logging.error('Unsupported where clauses type %s', type(where_clauses))
        return False

    def _format_datetime(self, datetime):
        return datetime.strftime(self.datetime_format)

    def _format_date(self, date):
        return date.strftime(self.date_format)

    def time2datetime(self, time):
        return self._format_datetime(datetime.fromtimestamp(int(time)))

    def time2date(self, time):
        return self._format_date(datetime.fromtimestamp(int(time)))

    def insert(self, table, values, just_try=False):
        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:
            logging.debug(u"Just-try mode : execute INSERT query : %s", sql)
            return True

        logging.debug(sql)
        if not self.doSQL(sql):
            logging.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):
        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:
            logging.debug(u"Just-try mode : execute UPDATE query : %s", sql)
            return True

        logging.debug(sql)
        if not self.doSQL(sql):
            logging.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):
        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:
            logging.debug(u"Just-try mode : execute DELETE query : %s", sql)
            return True

        logging.debug(sql)
        if not self.doSQL(sql):
            logging.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, just_try=False):
        sql = u"SELECT "
        if fields is None:
            sql += "*"
        elif isinstance(fields, str) or isinstance(fields, unicode):
            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)