more ironing in the corner of that 'types' modules
[plcapi.git] / PLC / PostgreSQL.py
index 076a864..2368419 100644 (file)
@@ -1,13 +1,16 @@
 #
-# PostgreSQL database interface. Sort of like DBI(3) (Database
-# independent interface for Perl).
+# PostgreSQL database interface.
+# Sort of like DBI(3) (Database independent interface for Perl).
 #
 # Mark Huang <mlhuang@cs.princeton.edu>
 # Copyright (C) 2006 The Trustees of Princeton University
 #
-# $Id$
-# $URL$
-#
+# pylint: disable=c0103, c0111
+
+import subprocess
+import re
+from pprint import pformat
+from datetime import datetime as DateTimeType
 
 import psycopg2
 import psycopg2.extensions
@@ -15,40 +18,10 @@ psycopg2.extensions.register_type(psycopg2.extensions.UNICODE)
 # UNICODEARRAY not exported yet
 psycopg2.extensions.register_type(psycopg2._psycopg.UNICODEARRAY)
 
-import pgdb
-from types import StringTypes, NoneType
-import traceback
-import commands
-import re
-from pprint import pformat
-
-from PLC.Debug import profile, log
+from PLC.Logger import logger
+#from PLC.Debug import profile
 from PLC.Faults import *
 
-if not psycopg2:
-    is8bit = re.compile("[\x80-\xff]").search
-
-    def unicast(typecast):
-        """
-        pgdb returns raw UTF-8 strings. This function casts strings that
-        appear to contain non-ASCII characters to unicode objects.
-        """
-    
-        def wrapper(*args, **kwds):
-            value = typecast(*args, **kwds)
-
-            # pgdb always encodes unicode objects as UTF-8 regardless of
-            # the DB encoding (and gives you no option for overriding
-            # the encoding), so always decode 8-bit objects as UTF-8.
-            if isinstance(value, str) and is8bit(value):
-                value = unicode(value, "utf-8")
-
-            return value
-
-        return wrapper
-
-    pgdb.pgdbTypeCache.typecast = unicast(pgdb.pgdbTypeCache.typecast)
-
 class PostgreSQL:
     def __init__(self, api):
         self.api = api
@@ -59,25 +32,21 @@ class PostgreSQL:
     def cursor(self):
         if self.connection is None:
             # (Re)initialize database connection
-            if psycopg2:
-                try:
-                    # Try UNIX socket first
-                    self.connection = psycopg2.connect(user = self.api.config.PLC_DB_USER,
-                                                       password = self.api.config.PLC_DB_PASSWORD,
-                                                       database = self.api.config.PLC_DB_NAME)
-                except psycopg2.OperationalError:
-                    # Fall back on TCP
-                    self.connection = psycopg2.connect(user = self.api.config.PLC_DB_USER,
-                                                       password = self.api.config.PLC_DB_PASSWORD,
-                                                       database = self.api.config.PLC_DB_NAME,
-                                                       host = self.api.config.PLC_DB_HOST,
-                                                       port = self.api.config.PLC_DB_PORT)
-                self.connection.set_client_encoding("UNICODE")
-            else:
-                self.connection = pgdb.connect(user = self.api.config.PLC_DB_USER,
-                                               password = self.api.config.PLC_DB_PASSWORD,
-                                               host = "%s:%d" % (api.config.PLC_DB_HOST, api.config.PLC_DB_PORT),
-                                               database = self.api.config.PLC_DB_NAME)
+            try:
+                # Try UNIX socket first
+                self.connection = psycopg2.connect(
+                    user=self.api.config.PLC_DB_USER,
+                    password=self.api.config.PLC_DB_PASSWORD,
+                    database=self.api.config.PLC_DB_NAME)
+            except psycopg2.OperationalError:
+                # Fall back on TCP
+                self.connection = psycopg2.connect(
+                    user=self.api.config.PLC_DB_USER,
+                    password=self.api.config.PLC_DB_PASSWORD,
+                    database=self.api.config.PLC_DB_NAME,
+                    host=self.api.config.PLC_DB_HOST,
+                    port=self.api.config.PLC_DB_PORT)
+            self.connection.set_client_encoding("UNICODE")
 
         (self.rowcount, self.description, self.lastrowid) = \
                         (None, None, None)
@@ -89,35 +58,62 @@ class PostgreSQL:
             self.connection.close()
             self.connection = None
 
-    # join insists on getting strings
-    @classmethod
-    def quote_string(self, value):
-        return str(PostgreSQL.quote(value))
+    @staticmethod
+    # From pgdb, and simplify code
+    def _quote(x):
+        if isinstance(x, DateTimeType):
+            x = str(x)
+        elif isinstance(x, str):
+            x = x.encode('utf-8')
+
+        if isinstance(x, bytes):
+            x = "'%s'" % str(x).replace("\\", "\\\\").replace("'", "''")
+        elif isinstance(x, (int, float)):
+            pass
+        elif x is None:
+            x = 'NULL'
+        elif isinstance(x, (list, tuple, set)):
+            x = 'ARRAY[%s]' % ', '.join([str(PostgreSQL._quote(x)) for x in x])
+        elif hasattr(x, '__pg_repr__'):
+            x = x.__pg_repr__()
+        else:
+            raise PLCDBError('Cannot quote type %s' % type(x))
+        return x
+
 
-    @classmethod
     def quote(self, value):
         """
         Returns quoted version of the specified value.
         """
-
-        # The pgdb._quote function is good enough for general SQL
-        # quoting, except for array types.
-        if isinstance(value, (list, tuple, set)):
-            return "ARRAY[%s]" % ", ".join(map (PostgreSQL.quote_string, value))
-        else:
-            return pgdb._quote(value)
+        return PostgreSQL._quote(value)
+
+# following is an unsuccessful attempt to re-use lib code as much as possible
+#    def quote(self, value):
+#        # The pgdb._quote function is good enough for general SQL
+#        # quoting, except for array types.
+#        if isinstance (value, (types.ListType, types.TupleType, set)):
+#            'ARRAY[%s]' % ', '.join( [ str(self.quote(x)) for x in value ] )
+#        else:
+#            try:
+#                # up to PyGreSQL-3.x, function was pgdb._quote
+#                import pgdb
+#                return pgdb._quote(value)
+#            except:
+#                # with PyGreSQL-4.x, use psycopg2's adapt
+#                from psycopg2.extensions import adapt
+#                return adapt (value)
 
     @classmethod
-    def param(self, name, value):
+    def param(cls, name, value):
         # None is converted to the unquoted string NULL
-        if isinstance(value, NoneType):
+        if isinstance(value, type(None)):
             conversion = "s"
         # True and False are also converted to unquoted strings
         elif isinstance(value, bool):
             conversion = "s"
         elif isinstance(value, float):
             conversion = "f"
-        elif not isinstance(value, StringTypes):
+        elif not isinstance(value, str):
             conversion = "d"
         else:
             conversion = "s"
@@ -134,36 +130,35 @@ class PostgreSQL:
     def rollback(self):
         self.connection.rollback()
 
-    def do(self, query, params = None):
+    def do(self, query, params=None):
         cursor = self.execute(query, params)
         cursor.close()
         return self.rowcount
 
     def next_id(self, table_name, primary_key):
-       sequence = "%(table_name)s_%(primary_key)s_seq" % locals()      
-       sql = "SELECT nextval('%(sequence)s')" % locals()
-       rows = self.selectall(sql, hashref = False)
-       if rows: 
-           return rows[0][0]
-               
-       return None 
+        sequence = "{}_{}_seq".format(table_name, primary_key)
+        sql = "SELECT nextval('{}')".format(sequence)
+        rows = self.selectall(sql, hashref=False)
+        if rows:
+            return rows[0][0]
+        return None
 
     def last_insert_id(self, table_name, primary_key):
         if isinstance(self.lastrowid, int):
             sql = "SELECT %s FROM %s WHERE oid = %d" % \
                   (primary_key, table_name, self.lastrowid)
-            rows = self.selectall(sql, hashref = False)
+            rows = self.selectall(sql, hashref=False)
             if rows:
                 return rows[0][0]
 
         return None
 
-    # modified for psycopg2-2.0.7 
+    # modified for psycopg2-2.0.7
     # executemany is undefined for SELECT's
     # see http://www.python.org/dev/peps/pep-0249/
     # accepts either None, a single dict, a tuple of single dict - in which case it execute's
     # or a tuple of several dicts, in which case it executemany's
-    def execute(self, query, params = None):
+    def execute(self, query, params=None):
 
         cursor = self.cursor()
         try:
@@ -176,40 +171,37 @@ class PostgreSQL:
             if psycopg2:
                 query = re.sub(r'(%\([^)]*\)|%)[df]', r'\1s', query)
             # rewrite wildcards set by Filter.py as '***' into '%'
-            query = query.replace ('***','%')
+            query = query.replace('***', '%')
 
             if not params:
                 if self.debug:
-                    print >> log,'execute0',query
+                    logger.debug('execute0: {}'.format(query))
                 cursor.execute(query)
-            elif isinstance(params,dict):
+            elif isinstance(params, dict):
                 if self.debug:
-                    print >> log,'execute-dict: params',params,'query',query%params
-                cursor.execute(query,params)
-            elif isinstance(params,tuple) and len(params)==1:
+                    logger.debug('execute-dict: params {} query {}'
+                                 .format(params, query%params))
+                cursor.execute(query, params)
+            elif isinstance(params, tuple) and len(params) == 1:
                 if self.debug:
-                    print >> log,'execute-tuple',query%params[0]
-                cursor.execute(query,params[0])
+                    logger.debug('execute-tuple {}'.format(query%params[0]))
+                cursor.execute(query, params[0])
             else:
-                param_seq=(params,)
+                param_seq = (params,)
                 if self.debug:
                     for params in param_seq:
-                        print >> log,'executemany',query%params
+                        logger.debug('executemany {}'.format(query%params))
                 cursor.executemany(query, param_seq)
             (self.rowcount, self.description, self.lastrowid) = \
                             (cursor.rowcount, cursor.description, cursor.lastrowid)
-        except Exception, e:
+        except Exception as e:
             try:
                 self.rollback()
             except:
                 pass
-            uuid = commands.getoutput("uuidgen")
-            print >> log, "Database error %s:" % uuid
-            print >> log, e
-            print >> log, "Query:"
-            print >> log, query
-            print >> log, "Params:"
-            print >> log, pformat(params)
+            uuid = subprocess.getoutput("uuidgen")
+            message = "Database error {}: - Query {} - Params {}".format(uuid, query, pformat(params))
+            logger.exception(message)
             raise PLCDBError("Please contact " + \
                              self.api.config.PLC_NAME + " Support " + \
                              "<" + self.api.config.PLC_MAIL_SUPPORT_ADDRESS + ">" + \
@@ -217,7 +209,7 @@ class PostgreSQL:
 
         return cursor
 
-    def selectall(self, query, params = None, hashref = True, key_field = None):
+    def selectall(self, query, params=None, hashref=True, key_field=None):
         """
         Return each row as a dictionary keyed on field name (like DBI
         selectrow_hashref()). If key_field is specified, return rows
@@ -236,22 +228,22 @@ class PostgreSQL:
             # Return each row as a dictionary keyed on field name
             # (like DBI selectrow_hashref()).
             labels = [column[0] for column in self.description]
-            rows = [dict(zip(labels, row)) for row in rows]
+            rows = [dict(list(zip(labels, row))) for row in rows]
 
         if key_field is not None and key_field in labels:
             # Return rows as a dictionary keyed on the specified field
             # (like DBI selectall_hashref()).
-            return dict([(row[key_field], row) for row in rows])
+            return {row[key_field]: row for row in rows}
         else:
             return rows
 
-    def fields(self, table, notnull = None, hasdef = None):
+    def fields(self, table, notnull=None, hasdef=None):
         """
         Return the names of the fields of the specified table.
         """
 
         if hasattr(self, 'fields_cache'):
-            if self.fields_cache.has_key((table, notnull, hasdef)):
+            if (table, notnull, hasdef) in self.fields_cache:
                 return self.fields_cache[(table, notnull, hasdef)]
         else:
             self.fields_cache = {}
@@ -266,7 +258,7 @@ class PostgreSQL:
         if hasdef is not None:
             sql += " AND atthasdef is %(hasdef)s"
 
-        rows = self.selectall(sql, locals(), hashref = False)
+        rows = self.selectall(sql, locals(), hashref=False)
 
         self.fields_cache[(table, notnull, hasdef)] = [row[0] for row in rows]