python/lib/django-0.96/django/db/backends/postgresql/base.py

   1 """
   2 PostgreSQL database backend for Django.
   3
   4 Requires psycopg 1: http://initd.org/projects/psycopg1
   5 """
   6
   7 from django.db.backends import util
   8 try:
   9     import psycopg as Database
  10 except ImportError, e:
  11     from django.core.exceptions import ImproperlyConfigured
  12     raise ImproperlyConfigured, "Error loading psycopg module: %s" % e
  13
  14 DatabaseError = Database.DatabaseError
  15
  16 try:
  17     # Only exists in Python 2.4+
  18     from threading import local
  19 except ImportError:
  20     # Import copy of _thread_local.py from Python 2.4
  21     from django.utils._threading_local import local
  22
  23 def smart_basestring(s, charset):
  24     if isinstance(s, unicode):
  25         return s.encode(charset)
  26     return s
  27
  28 class UnicodeCursorWrapper(object):
  29     """
  30     A thin wrapper around psycopg cursors that allows them to accept Unicode
  31     strings as params.
  32
  33     This is necessary because psycopg doesn't apply any DB quoting to
  34     parameters that are Unicode strings. If a param is Unicode, this will
  35     convert it to a bytestring using DEFAULT_CHARSET before passing it to
  36     psycopg.
  37     """
  38     def __init__(self, cursor, charset):
  39         self.cursor = cursor
  40         self.charset = charset
  41
  42     def execute(self, sql, params=()):
  43         return self.cursor.execute(sql, [smart_basestring(p, self.charset) for p in params])
  44
  45     def executemany(self, sql, param_list):
  46         new_param_list = [tuple([smart_basestring(p, self.charset) for p in params]) for params in param_list]
  47         return self.cursor.executemany(sql, new_param_list)
  48
  49     def __getattr__(self, attr):
  50         if self.__dict__.has_key(attr):
  51             return self.__dict__[attr]
  52         else:
  53             return getattr(self.cursor, attr)
  54
  55 postgres_version = None
  56
  57 class DatabaseWrapper(local):
  58     def __init__(self, **kwargs):
  59         self.connection = None
  60         self.queries = []
  61         self.options = kwargs
  62
  63     def cursor(self):
  64         from django.conf import settings
  65         set_tz = False
  66         if self.connection is None:
  67             set_tz = True
  68             if settings.DATABASE_NAME == '':
  69                 from django.core.exceptions import ImproperlyConfigured
  70                 raise ImproperlyConfigured, "You need to specify DATABASE_NAME in your Django settings file."
  71             conn_string = "dbname=%s" % settings.DATABASE_NAME
  72             if settings.DATABASE_USER:
  73                 conn_string = "user=%s %s" % (settings.DATABASE_USER, conn_string)
  74             if settings.DATABASE_PASSWORD:
  75                 conn_string += " password='%s'" % settings.DATABASE_PASSWORD
  76             if settings.DATABASE_HOST:
  77                 conn_string += " host=%s" % settings.DATABASE_HOST
  78             if settings.DATABASE_PORT:
  79                 conn_string += " port=%s" % settings.DATABASE_PORT
  80             self.connection = Database.connect(conn_string, **self.options)
  81             self.connection.set_isolation_level(1) # make transactions transparent to all cursors
  82         cursor = self.connection.cursor()
  83         if set_tz:
  84             cursor.execute("SET TIME ZONE %s", [settings.TIME_ZONE])
  85         cursor = UnicodeCursorWrapper(cursor, settings.DEFAULT_CHARSET)
  86         global postgres_version
  87         if not postgres_version:
  88             cursor.execute("SELECT version()")
  89             postgres_version = [int(val) for val in cursor.fetchone()[0].split()[1].split('.')]
  90         if settings.DEBUG:
  91             return util.CursorDebugWrapper(cursor, self)
  92         return cursor
  93
  94     def _commit(self):
  95         if self.connection is not None:
  96             return self.connection.commit()
  97
  98     def _rollback(self):
  99         if self.connection is not None:
 100             return self.connection.rollback()
 101
 102     def close(self):
 103         if self.connection is not None:
 104             self.connection.close()
 105             self.connection = None
 106
 107 supports_constraints = True
 108
 109 def quote_name(name):
 110     if name.startswith('"') and name.endswith('"'):
 111         return name # Quoting once is enough.
 112     return '"%s"' % name
 113
 114 def dictfetchone(cursor):
 115     "Returns a row from the cursor as a dict"
 116     return cursor.dictfetchone()
 117
 118 def dictfetchmany(cursor, number):
 119     "Returns a certain number of rows from a cursor as a dict"
 120     return cursor.dictfetchmany(number)
 121
 122 def dictfetchall(cursor):
 123     "Returns all rows from a cursor as a dict"
 124     return cursor.dictfetchall()
 125
 126 def get_last_insert_id(cursor, table_name, pk_name):
 127     cursor.execute("SELECT CURRVAL('\"%s_%s_seq\"')" % (table_name, pk_name))
 128     return cursor.fetchone()[0]
 129
 130 def get_date_extract_sql(lookup_type, table_name):
 131     # lookup_type is 'year', 'month', 'day'
 132     # http://www.postgresql.org/docs/8.0/static/functions-datetime.html#FUNCTIONS-DATETIME-EXTRACT
 133     return "EXTRACT('%s' FROM %s)" % (lookup_type, table_name)
 134
 135 def get_date_trunc_sql(lookup_type, field_name):
 136     # lookup_type is 'year', 'month', 'day'
 137     # http://www.postgresql.org/docs/8.0/static/functions-datetime.html#FUNCTIONS-DATETIME-TRUNC
 138     return "DATE_TRUNC('%s', %s)" % (lookup_type, field_name)
 139
 140 def get_limit_offset_sql(limit, offset=None):
 141     sql = "LIMIT %s" % limit
 142     if offset and offset != 0:
 143         sql += " OFFSET %s" % offset
 144     return sql
 145
 146 def get_random_function_sql():
 147     return "RANDOM()"
 148
 149 def get_deferrable_sql():
 150     return " DEFERRABLE INITIALLY DEFERRED"
 151
 152 def get_fulltext_search_sql(field_name):
 153     raise NotImplementedError
 154
 155 def get_drop_foreignkey_sql():
 156     return "DROP CONSTRAINT"
 157
 158 def get_pk_default_value():
 159     return "DEFAULT"
 160
 161 def get_sql_flush(style, tables, sequences):
 162     """Return a list of SQL statements required to remove all data from
 163     all tables in the database (without actually removing the tables
 164     themselves) and put the database in an empty 'initial' state
 165
 166     """
 167     if tables:
 168         if postgres_version[0] >= 8 and postgres_version[1] >= 1:
 169             # Postgres 8.1+ can do 'TRUNCATE x, y, z...;'. In fact, it *has to* in order to be able to
 170             # truncate tables referenced by a foreign key in any other table. The result is a
 171             # single SQL TRUNCATE statement.
 172             sql = ['%s %s;' % \
 173                 (style.SQL_KEYWORD('TRUNCATE'),
 174                  style.SQL_FIELD(', '.join([quote_name(table) for table in tables]))
 175             )]
 176         else:
 177             # Older versions of Postgres can't do TRUNCATE in a single call, so they must use
 178             # a simple delete.
 179             sql = ['%s %s %s;' % \
 180                     (style.SQL_KEYWORD('DELETE'),
 181                      style.SQL_KEYWORD('FROM'),
 182                      style.SQL_FIELD(quote_name(table))
 183                      ) for table in tables]
 184
 185         # 'ALTER SEQUENCE sequence_name RESTART WITH 1;'... style SQL statements
 186         # to reset sequence indices
 187         for sequence_info in sequences:
 188             table_name = sequence_info['table']
 189             column_name = sequence_info['column']
 190             if column_name and len(column_name)>0:
 191                 # sequence name in this case will be <table>_<column>_seq
 192                 sql.append("%s %s %s %s %s %s;" % \
 193                     (style.SQL_KEYWORD('ALTER'),
 194                     style.SQL_KEYWORD('SEQUENCE'),
 195                     style.SQL_FIELD('%s_%s_seq' % (table_name, column_name)),
 196                     style.SQL_KEYWORD('RESTART'),
 197                     style.SQL_KEYWORD('WITH'),
 198                     style.SQL_FIELD('1')
 199                     )
 200                 )
 201             else:
 202                 # sequence name in this case will be <table>_id_seq
 203                 sql.append("%s %s %s %s %s %s;" % \
 204                     (style.SQL_KEYWORD('ALTER'),
 205                      style.SQL_KEYWORD('SEQUENCE'),
 206                      style.SQL_FIELD('%s_id_seq' % table_name),
 207                      style.SQL_KEYWORD('RESTART'),
 208                      style.SQL_KEYWORD('WITH'),
 209                      style.SQL_FIELD('1')
 210                      )
 211                 )
 212         return sql
 213     else:
 214         return []
 215
 216
 217 # Register these custom typecasts, because Django expects dates/times to be
 218 # in Python's native (standard-library) datetime/time format, whereas psycopg
 219 # use mx.DateTime by default.
 220 try:
 221     Database.register_type(Database.new_type((1082,), "DATE", util.typecast_date))
 222 except AttributeError:
 223     raise Exception, "You appear to be using psycopg version 2. Set your DATABASE_ENGINE to 'postgresql_psycopg2' instead of 'postgresql'."
 224 Database.register_type(Database.new_type((1083,1266), "TIME", util.typecast_time))
 225 Database.register_type(Database.new_type((1114,1184), "TIMESTAMP", util.typecast_timestamp))
 226 Database.register_type(Database.new_type((16,), "BOOLEAN", util.typecast_boolean))
 227
 228 OPERATOR_MAPPING = {
 229     'exact': '= %s',
 230     'iexact': 'ILIKE %s',
 231     'contains': 'LIKE %s',
 232     'icontains': 'ILIKE %s',
 233     'gt': '> %s',
 234     'gte': '>= %s',
 235     'lt': '< %s',
 236     'lte': '<= %s',
 237     'startswith': 'LIKE %s',
 238     'endswith': 'LIKE %s',
 239     'istartswith': 'ILIKE %s',
 240     'iendswith': 'ILIKE %s',
 241 }