Package backend :: Package server :: Package rhnSQL :: Module driver_postgresql
[hide private]
[frames] | no frames]

Source Code for Module backend.server.rhnSQL.driver_postgresql

  1  # 
  2  # Copyright (c) 2008--2017 Red Hat, Inc. 
  3  # 
  4  # This software is licensed to you under the GNU General Public License, 
  5  # version 2 (GPLv2). There is NO WARRANTY for this software, express or 
  6  # implied, including the implied warranties of MERCHANTABILITY or FITNESS 
  7  # FOR A PARTICULAR PURPOSE. You should have received a copy of GPLv2 
  8  # along with this software; if not, see 
  9  # http://www.gnu.org/licenses/old-licenses/gpl-2.0.txt. 
 10  # 
 11  # Red Hat trademarks are not licensed under GPLv2. No permission is 
 12  # granted to use or replicate Red Hat trademarks that are incorporated 
 13  # in this software or its documentation. 
 14  # 
 15  # 
 16  # 
 17  # Database driver for PostgreSQL 
 18  # 
 19   
 20  import sys 
 21  import string 
 22  import re 
 23  import psycopg2 
 24   
 25  # workaround for python-psycopg2 = 2.0.13 (RHEL6) 
 26  # which does not import extensions by default 
 27  if not hasattr(psycopg2, 'extensions'): 
 28      import psycopg2.extensions 
 29   
 30  import sql_base 
 31  from rhn.UserDictCase import UserDictCase 
 32  from spacewalk.server import rhnSQL 
 33   
 34  from spacewalk.common.usix import BufferType, raise_with_tb 
 35  from spacewalk.common.rhnLog import log_debug, log_error 
 36  from spacewalk.common.rhnException import rhnException 
 37  from const import POSTGRESQL 
 38   
 39   
40 -def convert_named_query_params(query):
41 """ 42 Convert a query with named parameters (i.e. :id, :name, etc) into one 43 that uses %(id)s, %(name)s parameters instead. 44 45 python-psycopg2 requires parameters to be in this form, so to keep our 46 existing queries intact we'll convert them when provided to the 47 postgresql driver. 48 49 RETURNS: the new query with parameters replaced 50 """ 51 log_debug(6, "Converting query for PostgreSQL: %s" % query) 52 new_query = re.sub(r'(\W):(\w+)', r'\1%(\2)s', query.replace('%', '%%')) 53 log_debug(6, "New query: %s" % new_query) 54 return new_query
55 56
57 -class Function(sql_base.Procedure):
58 59 """ 60 Function implementation for PostgreSQL. As there is no support in the Python 61 driver we use direct SQL. 62 """ 63
64 - def __init__(self, name, cursor, ret_type):
65 sql_base.Procedure.__init__(self, name, cursor) 66 self.ret_type = ret_type
67
68 - def __call__(self, *args):
69 log_debug(2, self.name, args) 70 71 # Buildup a string for the positional arguments to the procedure: 72 positional_args = "" 73 i = 1 74 for arg in args: 75 if len(positional_args) == 0: 76 positional_args = "%s" 77 else: 78 positional_args = positional_args + ", %s" 79 i += 1 80 query = "SELECT %s(%s)" % (self.name, positional_args) 81 82 log_debug(2, query, args) 83 try: 84 ret = self.cursor.execute(query, args) 85 except psycopg2.Error: 86 e = sys.exc_info()[1] 87 error_code = 99999 88 m = re.match('ERROR: +-([0-9]+)', e.pgerror) 89 if m: 90 error_code = int(m.group(1)) 91 raise sql_base.SQLSchemaError(error_code, e.pgerror, e) 92 93 if self.ret_type is None: 94 return ret 95 else: 96 return self.cursor.fetchone()[0]
97 98
99 -class Procedure(Function):
100 101 """ 102 PostgreSQL functions are somewhat different than stored procedures in 103 other databases. As a result the python-pgsql does not even implement 104 the Python DBI API callproc method. 105 106 To workaround this and keep rhnSQL database independent, we'll translate 107 any incoming requests to call a procedure into a PostgreSQL query. 108 """ 109
110 - def __init__(self, name, cursor):
111 Function.__init__(self, name, cursor, None) 112 self.ret_type = None
113
114 - def __call__(self, *args):
115 result = Function.__call__(self, *args)
116 # we do not expect any result (this is procedure) 117 # if not (type(result) == 'tuple' and result[0] == ''): 118 #raise rhnSQL.SQLError("Unexpected result returned by procedure %s: %s" % (self.name, str(result))) 119 120
121 -def decimal2intfloat(dec, cursor):
122 "Convert a Decimal to an int or a float with no loss of information." 123 "The dec is passed in as str (not Decimal) so we cannot check its type." 124 if dec is None: 125 return None 126 "If we can convert to int without loss of information, return int, float otherwise." 127 try: 128 if float(dec) == float(int(dec)): 129 return int(dec) 130 return float(dec) 131 except ValueError: 132 return float(dec)
133 134
135 -class Database(sql_base.Database):
136 137 """ Class for PostgreSQL database operations. """ 138
139 - def __init__(self, host=None, port=None, username=None, 140 password=None, database=None, sslmode=None, sslrootcert=None):
141 142 self.username = username 143 self.password = password 144 self.database = database 145 self.sslmode = sslmode 146 self.sslrootcert = sslrootcert 147 148 # Minimum requirements to connect to a PostgreSQL db: 149 if not (self.username and self.database): 150 raise AttributeError("PostgreSQL requires at least a user and database name.") 151 152 if host is None or host == '' or host == 'local': 153 self.host = None 154 self.port = None 155 else: 156 self.host = host 157 self.port = port 158 159 # pgsql module prefers -1 for an unspecified port: 160 if not self.port: 161 self.port = -1 162 163 self.dbh = None 164 165 sql_base.Database.__init__(self)
166
167 - def connect(self, reconnect=1):
168 try: 169 dsndata = { 170 'dbname': self.database, 171 'user': self.username, 172 'password': self.password} 173 if self.host is not None: 174 dsndata['host'] = self.host 175 dsndata['port'] = self.port 176 if self.sslmode is not None and self.sslmode == 'verify-full' and self.sslrootcert is not None: 177 dsndata['sslmode'] = self.sslmode 178 dsndata['sslrootcert'] = self.sslrootcert 179 elif self.sslmode is not None: 180 raise AttributeError("Only sslmode=\"verify-full\" (or None) is supported.") 181 if self.sslmode is not None and self.sslrootcert is None: 182 raise AttributeError("Attribute sslrootcert needs to be set if sslmode is set.") 183 184 self.dbh = psycopg2.connect(" ".join("%s=%s" % (k, re.escape(str(v))) for k, v in dsndata.items())) 185 186 # convert all DECIMAL types to float (let Python to choose one) 187 DEC2INTFLOAT = psycopg2.extensions.new_type(psycopg2._psycopg.DECIMAL.values, 188 'DEC2INTFLOAT', decimal2intfloat) 189 psycopg2.extensions.register_type(DEC2INTFLOAT) 190 except psycopg2.Error: 191 e = sys.exc_info()[1] 192 if reconnect > 0: 193 # Try one more time: 194 return self.connect(reconnect=reconnect - 1) 195 196 # Failed reconnect, time to error out: 197 raise_with_tb(sql_base.SQLConnectError( 198 self.database, e.pgcode, e.pgerror, 199 "All attempts to connect to the database failed"), sys.exc_info()[2])
200
201 - def is_connected_to(self, backend, host, port, username, password, 202 database, sslmode, sslrootcert):
203 if host is None or host == '' or host == 'local': 204 host = None 205 port = None 206 if not port: 207 port = -1 208 return (backend == POSTGRESQL) and (self.host == host) and \ 209 (self.port == port) and (self.username == username) and \ 210 (self.password == password) and (self.database == database) and \ 211 (self.sslmode == sslmode) and (self.sslrootcert == sslrootcert)
212
213 - def check_connection(self):
214 try: 215 c = self.prepare("select 1") 216 c.execute() 217 except: # try to reconnect, that one MUST WORK always 218 log_error("DATABASE CONNECTION TO '%s' LOST" % self.database, 219 "Exception information: %s" % sys.exc_info()[1]) 220 self.connect() # only allow one try
221
222 - def prepare(self, sql, force=0, blob_map=None):
223 return Cursor(dbh=self.dbh, sql=sql, force=force, blob_map=blob_map)
224
225 - def execute(self, sql, *args, **kwargs):
226 cursor = self.prepare(sql) 227 cursor.execute(*args, **kwargs) 228 return cursor
229
230 - def transaction(self, name):
231 if not name: 232 raise rhnException("Can not set a transaction without a name", name) 233 c = self.prepare("savepoint %s" % name) 234 return c.execute()
235
236 - def commit(self):
237 if self.dbh is not None: 238 self.dbh.commit()
239
240 - def rollback(self, name=None):
241 if name: 242 c = self.prepare("rollback to %s" % name) 243 return c.execute() 244 else: 245 return self.dbh.rollback()
246
247 - def procedure(self, name):
248 c = self.dbh.cursor() 249 # Pass the cursor in so we can close it after execute() 250 return Procedure(name, c)
251
252 - def _function(self, name, ret_type):
253 c = self.dbh.cursor() 254 return Function(name, c, ret_type)
255
256 - def cursor(self):
257 return Cursor(dbh=self.dbh)
258
259 - def _read_lob(self, lob):
260 return str(lob)
261 262
263 -class Cursor(sql_base.Cursor):
264 265 """ PostgreSQL specific wrapper over sql_base.Cursor. """ 266
267 - def __init__(self, dbh=None, sql=None, force=None, blob_map=None):
268 269 sql_base.Cursor.__init__(self, dbh, sql, force) 270 self.blob_map = blob_map 271 272 # Accept Oracle style named query params, but convert for python-pgsql 273 # under the hood: 274 temp_sql = "" 275 if self.sql is not None: 276 temp_sql = self.sql 277 self.sql = convert_named_query_params(temp_sql)
278
279 - def _prepare_sql(self):
280 cursor = self.dbh.cursor() 281 return cursor
282
283 - def _execute_wrapper(self, function, *p, **kw):
284 params = ','.join(["%s: %s" % (key, value) for key, value 285 in list(kw.items())]) 286 log_debug(5, "Executing SQL: \"%s\" with bind params: {%s}" 287 % (self.sql, params)) 288 if self.sql is None: 289 raise rhnException("Cannot execute empty cursor") 290 if self.blob_map: 291 for blob_var in list(self.blob_map.keys()): 292 kw[blob_var] = BufferType(kw[blob_var]) 293 294 try: 295 retval = function(*p, **kw) 296 except psycopg2.InternalError: 297 e = sys.exc_info()[1] 298 error_code = 99999 299 m = re.match('ERROR: +-([0-9]+)', e.pgerror) 300 if m: 301 error_code = int(m.group(1)) 302 raise sql_base.SQLSchemaError(error_code, e.pgerror, e) 303 except psycopg2.ProgrammingError: 304 e = sys.exc_info()[1] 305 raise sql_base.SQLStatementPrepareError(self.dbh, e.pgerror, self.sql) 306 except KeyError: 307 e = sys.exc_info()[1] 308 raise sql_base.SQLError("Unable to bound the following variable(s): %s" 309 % (string.join(e.args, " "))) 310 return retval
311
312 - def _execute_(self, args, kwargs):
313 """ 314 PostgreSQL specific execution of the query. 315 """ 316 params = UserDictCase(kwargs) 317 try: 318 self._real_cursor.execute(self.sql, params) 319 except psycopg2.OperationalError: 320 e = sys.exc_info()[1] 321 raise sql_base.SQLError("Cannot execute SQL statement: %s" % str(e)) 322 323 self.description = self._real_cursor.description 324 return self._real_cursor.rowcount
325
326 - def _executemany(self, *args, **kwargs):
327 if not kwargs: 328 return 0 329 330 params = UserDictCase(kwargs) 331 332 # First break all the incoming keyword arg lists into individual 333 # hashes: 334 all_kwargs = [] 335 for key in params.keys(): 336 if len(all_kwargs) < len(params[key]): 337 for i in range(len(params[key])): 338 all_kwargs.append({}) 339 340 i = 0 341 for val in params[key]: 342 all_kwargs[i][key] = val 343 i = i + 1 344 345 self._real_cursor.executemany(self.sql, all_kwargs) 346 self.description = self._real_cursor.description 347 rowcount = self._real_cursor.rowcount 348 return rowcount
349
350 - def update_blob(self, table_name, column_name, where_clause, data, 351 **kwargs):
352 """ 353 PostgreSQL uses bytea columns instead of blobs. Nothing special 354 needs to be done to insert text into one. 355 """ 356 # NOTE: Injecting a :column_name parameter here 357 sql = "UPDATE %s SET %s = :%s %s" % (table_name, column_name, 358 column_name, where_clause) 359 c = rhnSQL.prepare(sql) 360 kwargs[column_name] = data 361 c.execute(**kwargs)
362
363 - def close(self):
364 pass
365