Source Code for Module conary.dbstore.postgresql_drv

  1  # 
  2  # Copyright (c) 2005-2008 rPath, Inc. 
  3  # 
  4  # This program is distributed under the terms of the Common Public License, 
  5  # version 1.0. A copy of this license should have been distributed with this 
  6  # source file in a file called LICENSE. If it is not present, the license 
  7  # is always available at http://www.rpath.com/permanent/licenses/CPL-1.0. 
  8  # 
  9  # This program is distributed in the hope that it will be useful, but 
 10  # without any warranty; without even the implied warranty of merchantability 
 11  # or fitness for a particular purpose. See the Common Public License for 
 12  # full details. 
 13  # 
 14   
 15  import re 
 16  import pgsql 
 17   
 18  from base_drv import BaseDatabase, BaseCursor, BaseBinary 
 19  from base_drv import BaseKeywordDict 
 20  import sqlerrors 
 21  import sqllib 
 22   
 23 -class KeywordDict(BaseKeywordDict): 
 24      keys = BaseKeywordDict.keys.copy() 
 25      keys.update( { 
 26          'PRIMARYKEY' : 'SERIAL PRIMARY KEY', 
 27          'BLOB'       : 'BYTEA', 
 28          'MEDIUMBLOB' : 'BYTEA', 
 29          'PATHTYPE'   : 'BYTEA', 
 30          'STRING'     : 'VARCHAR' 
 31          } ) 
 32   
 33 -    def binaryVal(self, len): 
 34          return "BYTEA" 
 35   
 36  # class for encapsulating binary strings for dumb drivers 
 37 -class Binary(BaseBinary): 
 38      __binary__ = True 
 39 -    def __quote__(self): 
 40          return self.s 
 41 -    def __pg_repr__(self): 
 42          return "decode('%s','hex')" % "".join("%02x" % ord(c) for c in self.s) 
 43   
 44  # edit the input query to make it postgres compatible 
 45 -def _mungeSQL(sql): 
 46      keys = [] # needs to be a list because we're dealing with positional args 
 47      def __match(m): 
 48          d = m.groupdict() 
 49          kw = d["kw"][1:] 
 50          if len(kw): # a real keyword 
 51              if kw not in keys: 
 52                  keys.append(kw) 
 53              d["kwIdx"] = keys.index(kw)+1 
 54          else: # if we have just the ? then kw is "" here 
 55              keys.append(None) 
 56              d["kwIdx"] = len(keys) 
 57          return "%(pre)s%(s)s$%(kwIdx)d" % d 
 58   
 59      sql = re.sub("(?i)(?P<pre>[(,<>=]|(LIKE|AND|BETWEEN|LIMIT|OFFSET)\s)(?P<s>\s*)(?P<kw>:\w+|[?])", 
 60                   __match, sql) 
 61      # force dbi compliance here. args or kw or none, no mixes 
 62      if len(keys) and keys[0] is not None: 
 63          return (sql, keys) 
 64      return (sql, []) 
 65   
 66 -class Cursor(BaseCursor): 
 67      binaryClass = Binary 
 68      driver = "postgresql" 
 69   
 70  ##     def binary(self, s): 
 71  ##         return s 
 72   
 73  ##     def frombinary(self, s): 
 74  ##         #return s.decode("string_escape") 
 75  ##         return s 
 76   
 77      # execute with exception translation 
 78 -    def _tryExecute(self, func, *params, **kw): 
 79          try: 
 80              ret = func(*params, **kw) 
 81          except pgsql.DatabaseError, e: 
 82              msg = e.args[0] 
 83              if msg.find("violates foreign key constraint") > 0: 
 84                  raise sqlerrors.ConstraintViolation(msg) 
 85              if re.search('relation \S+ does not exist', msg, re.I): 
 86                  raise sqlerrors.InvalidTable(msg) 
 87              if re.search("duplicate key (value )?violates unique constraint", msg): 
 88                  raise sqlerrors.ColumnNotUnique(msg) 
 89              raise sqlerrors.CursorError(msg, e) 
 90          return ret 
 91   
 92      # we need to "fix" the sql code before calling out 
 93 -    def execute(self, sql, *args, **kw): 
 94          self._executeCheck(sql) 
 95          keys = [] 
 96   
 97          kw.pop("start_transaction", True) 
 98          args, kw  = self._executeArgs(args, kw) 
 99   
100          # don't do unnecessary work 
101          if len(args) or len(kw): 
102              sql, keys = _mungeSQL(sql) 
103   
104          # if we have args, we can not have keywords 
105          if len(args): 
106              if len(kw) or len(keys): 
107                  raise sqlerrors.CursorError( 
108                      "Do not pass both positional and named bind arguments", 
109                      *args, **kw) 
110              ret = self._tryExecute(self._cursor.execute, sql, args) 
111          elif len(keys): # check that all keys used in the query appear in the kw 
112              if False in [kw.has_key(x) for x in keys]: 
113                  raise CursorError( 
114                      "Query keys not defined in named argument dict", 
115                      sorted(keys), sorted(kw.keys())) 
116              # need to transform kw into pozitional args 
117              ret = self._tryExecute(self._cursor.execute, sql, 
118                                     [kw[x] for x in keys]) 
119          else: 
120              ret = self._tryExecute(self._cursor.execute, sql) 
121          if ret == self._cursor: 
122              return self 
123          return ret 
124   
125      # executemany - we have to process the query code 
126 -    def executemany(self, sql, argList, **kw): 
127          self._executeCheck(sql) 
128          kw.pop("start_transaction", True) 
129          sql, keys = _mungeSQL(sql) 
130          if len(keys): 
131              # need to transform the dicts in tuples for the query 
132              return self._tryExecute(self._cursor.executemany, sql, 
133                                      (tuple([row[x] for x in keys]) for row in argList)) 
134          return self._tryExecute(self._cursor.executemany, sql, argList) 
135   
136      # support for prepared statements 
137 -    def compile(self, sql): 
138          self._executeCheck(sql) 
139          sql, keys = _mungeSQL(sql.strip()) 
140          stmt = self.dbh.prepare(sql) 
141          stmt.keys = keys 
142          return stmt 
143 -    def execstmt(self, stmt, *args): 
144          assert(isinstance(stmt, pgsql.PreparedCursor)) 
145          if not len(args): 
146              ret = self._tryExecute(stmt._source.execute) 
147          elif isinstance(args[0], (tuple, list)): 
148              ret = self._tryExecute(stmt._source.execute, *args) 
149          else: 
150              ret = self._tryExecute(stmt._source.execute, args) 
151          if isinstance(ret, int): 
152              return ret 
153          return stmt 
154   
155      # override this with the native version 
156 -    def fields(self): 
157          return self._cursor.fields 
158   
159      # pgsql has its own fetch*_dict methods 
160 -    def fetchone_dict(self): 
161          ret = self._cursor.fetchone_dict() 
162          return sqllib.CaselessDict(ret) 
163 -    def fetchmany_dict(self, size): 
164          return [ sqllib.CaselessDict(x) for x in self._cursor.fetchmany_dict(size) ] 
165 -    def fetchall_dict(self): 
166          return [ sqllib.CaselessDict(x) for x in self._cursor.fetchall_dict() ] 
167   
168      # we have "our own" lastrowid 
169 -    def __getattr__(self, name): 
170          if name == "lastrowid": 
171              return self.lastid() 
172          return BaseCursor.__getattr__(self, name) 
173   
174      # postgresql can not report back the last value from a SERIAL 
175      # PRIMARY KEY column insert, so we have to look it up ourselves 
176 -    def lastid(self): 
177          ret = self.execute("select lastval()").fetchone() 
178          if ret is None: 
179              return 0 
180          return ret[0] 
181   
182  # A cursor class that wraps PostgreSQL's server side cursors 
183 -class IterCursor(Cursor): 
184 -    def _getCursor(self): 
185          assert(self.dbh) 
186          return self.dbh.itercursor() 
187   
188 -class Database(BaseDatabase): 
189      driver = "postgresql" 
190      alive_check = "select version() as version" 
191      cursorClass = Cursor 
192      iterCursorClass = IterCursor 
193