common.lib.database

  1"""
  2Database wrapper
  3"""
  4import itertools
  5import psycopg2.extras
  6import psycopg2
  7import logging
  8import time
  9
 10from psycopg2 import sql
 11from psycopg2.extras import execute_values
 12
 13from common.lib.exceptions import DatabaseQueryInterruptedException
 14
 15class Database:
 16	"""
 17	Simple database handler
 18
 19	Offers a number of abstraction methods that limit how much SQL one is
 20	required to write. Also makes the database connection mostly multithreading
 21	proof by instantiating a new cursor for each query (and closing it afterwards)
 22	"""
 23	cursor = None
 24	log = None
 25	appname=""
 26
 27	interrupted = False
 28	interruptable_timeout = 86400  # if a query takes this long, it should be cancelled. see also fetchall_interruptable()
 29	interruptable_job = None
 30
 31	def __init__(self, logger, dbname=None, user=None, password=None, host=None, port=None, appname=None):
 32		"""
 33		Set up database connection
 34
 35		:param logger:  Logger instance
 36		:param dbname:  Database name
 37		:param user:  Database username
 38		:param password:  Database password
 39		:param host:  Database server address
 40		:param port:  Database port
 41		:param appname:  App name, mostly useful to trace connections in pg_stat_activity
 42		"""
 43		self.appname = "4CAT" if not appname else "4CAT-%s" % appname
 44
 45		self.connection = psycopg2.connect(dbname=dbname, user=user, password=password, host=host, port=port, application_name=self.appname)
 46		self.cursor = self.connection.cursor(cursor_factory=psycopg2.extras.RealDictCursor)
 47		self.log = logger
 48
 49		if self.log is None:
 50			self.log = logging
 51
 52	def reconnect(self, tries=3, wait=10):
 53		"""
 54		Reconnect to the database
 55
 56		:param int tries: Number of tries to reconnect
 57        :param int wait: Time to wait between tries (first try is immediate)
 58		"""
 59		for i in range(tries):
 60			try:
 61				self.connection = psycopg2.connect(dbname=self.connection.info.dbname,
 62												   user=self.connection.info.user,
 63												   password=self.connection.info.password,
 64												   host=self.connection.info.host,
 65												   port=self.connection.info.port,
 66												   application_name=self.appname)
 67				self.cursor = self.connection.cursor(cursor_factory=psycopg2.extras.RealDictCursor)
 68				return
 69			except (psycopg2.InterfaceError, psycopg2.OperationalError) as e:
 70				self.log.warning(f"Database connection closed. Reconnecting...\n{e}")
 71				time.sleep(wait)
 72		self.log.error("Failed to reconnect to database after %d tries" % tries)
 73
 74	def _execute_query(self, query, replacements=None, cursor=None):
 75		"""
 76		Execute a query
 77
 78		Simple wrapper to get a cursor to execute a query with. Do not call
 79		from outside this class - use `execute()` instead.
 80
 81		:param string query: Query
 82		:param args: Replacement values
 83		:param cursor: Cursor to use. Default - use common cursor
 84		:return None:
 85		"""
 86		if not cursor:
 87			cursor = self.get_cursor()
 88
 89		self.log.debug("Executing query %s" % cursor.mogrify(query, replacements))
 90		try:
 91			cursor.execute(query, replacements)
 92		except (psycopg2.InterfaceError, psycopg2.OperationalError) as e:
 93			self.log.warning(f"Database Exception: {e}\nReconnecting and retrying query...")
 94			self.reconnect()
 95			cursor = self.get_cursor()
 96			cursor.execute(query, replacements)
 97		return cursor
 98
 99	def execute(self, query, replacements=None, commit=True, cursor=None, close_cursor=True):
100		"""
101		Execute a query, and commit afterwards
102
103		This is required for UPDATE/INSERT/DELETE/etc to stick!
104
105		:param string query:  Query
106		:param cursor: Cursor to use. By default, use the result of
107		`get_cursor()`.
108		:param replacements: Replacement values
109		:param bool commit:  Commit transaction after query?
110		:param bool close_cursor:  Close cursor after query?
111		"""
112		cursor = self._execute_query(query, replacements, cursor)
113
114		if commit:
115			self.commit()
116
117		rowcount = cursor.rowcount
118
119		if close_cursor:
120			cursor.close()
121
122		return rowcount
123
124	def execute_many(self, query, commit=True, replacements=None):
125		"""
126		Execute a query multiple times, each time with different values
127
128		This makes it particularly suitable for INSERT queries, but other types
129		of query using VALUES are possible too.
130
131		:param string query:  Query
132		:param replacements: A list of replacement values
133		:param commit:  Commit transaction after query?
134		"""
135		cursor = self.get_cursor()
136		try:
137			execute_values(cursor, query, replacements)
138		except (psycopg2.InterfaceError, psycopg2.OperationalError) as e:
139			self.log.warning(f"Database Exception: {e}\nReconnecting and retrying query...")
140			self.reconnect()
141			cursor = self.get_cursor()
142			execute_values(cursor, query, replacements)
143
144		cursor.close()
145		if commit:
146			self.commit()
147
148	def update(self, table, data, where=None, commit=True):
149		"""
150		Update a database record
151
152		:param string table:  Table to update
153		:param dict data:  Data to set, Column => Value
154		:param dict where:  Simple conditions, parsed as "column1 = value1 AND column2 = value2" etc
155		:param bool commit:  Whether to commit after executing the query
156
157		:return int: Number of affected rows. Note that this may be unreliable if `commit` is `False`
158		"""
159		if where is None:
160			where = {}
161
162		# build query
163		identifiers = [sql.Identifier(column) for column in data.keys()]
164		identifiers.insert(0, sql.Identifier(table))
165		replacements = list(data.values())
166
167		query = "UPDATE {} SET " + ", ".join(["{} = %s" for column in data])
168		if where:
169			query += " WHERE " + " AND ".join(["{} = %s" for column in where])
170			for column in where.keys():
171				identifiers.append(sql.Identifier(column))
172				replacements.append(where[column])
173
174		query = sql.SQL(query).format(*identifiers)
175
176		rowcount = self.execute(query, replacements=replacements, commit=commit)
177		return rowcount
178
179	def delete(self, table, where, commit=True):
180		"""
181		Delete a database record
182
183		:param string table:  Table to delete from
184		:param dict where:  Simple conditions, parsed as "column1 = value1 AND column2 = value2" etc
185		:param bool commit:  Whether to commit after executing the query
186
187		:return int: Number of affected rows. Note that this may be unreliable if `commit` is `False`
188		"""
189		where_sql = []
190		replacements = []
191		for column in where.keys():
192			if type(where[column]) in (set, tuple, list):
193				where_sql.append("{} IN %s")
194				replacements.append(tuple(where[column]))
195			else:
196				where_sql.append("{} = %s")
197				replacements.append(where[column])
198
199		# build query
200		identifiers = [sql.Identifier(column) for column in where.keys()]
201		identifiers.insert(0, sql.Identifier(table))
202		query = sql.SQL("DELETE FROM {} WHERE " + " AND ".join(where_sql)).format(*identifiers)
203
204		rowcount = self.execute(query, replacements=replacements, commit=commit)
205		return rowcount
206
207	def insert(self, table, data, commit=True, safe=False, constraints=None, return_field=""):
208		"""
209		Create database record
210
211		:param string table:  Table to insert record into
212		:param dict data:   Data to insert
213		:param bool commit: Whether to commit after executing the query
214		:param bool safe: If set to `True`, "ON CONFLICT DO NOTHING" is added to the insert query, so it does not
215						  insert the row and no error is thrown when the insert violates a unique index or other constraint
216		:param tuple constraints: If `safe` is `True`, this tuple may contain the columns that should be used as a
217								  constraint, e.g. ON CONFLICT (name, lastname) DO NOTHING
218		:param str return_field: If not empty or None, this makes the method
219		return this field of the inserted row, instead of the number of
220		affected rows, with `RETURNING`.
221		:return int: Number of affected rows. Note that this may be unreliable if `commit` is `False`
222		"""
223		if constraints is None:
224			constraints = []
225
226		# escape identifiers
227		identifiers = [sql.Identifier(column) for column in data.keys()]
228		identifiers.insert(0, sql.Identifier(table))
229
230		# construct ON NOTHING bit of query
231		if safe:
232			safe_bit = " ON CONFLICT "
233			if constraints:
234				safe_bit += "(" + ", ".join(["{}" for each in constraints]) + ")"
235				identifiers.extend([sql.Identifier(column) for column in constraints])
236			safe_bit += " DO NOTHING"
237		else:
238			safe_bit = ""
239
240		# prepare parameter replacements
241		protoquery = "INSERT INTO {} (%s) VALUES %%s" % ", ".join(["{}" for column in data.keys()]) + safe_bit
242
243		if return_field:
244			protoquery += " RETURNING {}"
245			identifiers.append(sql.Identifier(return_field))
246
247		query = sql.SQL(protoquery).format(*identifiers)
248		replacements = (tuple(data.values()),)
249
250		cursor = self.get_cursor()
251		rowcount = self.execute(query, replacements=replacements, cursor=cursor, commit=commit, close_cursor=False)
252
253		result = rowcount if not return_field else cursor.fetchone()[return_field]
254		cursor.close()
255		return result
256
257	def upsert(self, table, data, commit=True, constraints=None):
258		"""
259		Create or update database record
260
261		If the record could not be inserted because of a constraint, the
262		constraining record is updated instead.
263
264		:param string table:  Table to upsert record into
265		:param dict data:   Data to upsert
266		:param bool commit: Whether to commit after executing the query
267		:param tuple constraints: This tuple may contain the columns that should be used as a
268								  constraint, e.g. ON CONFLICT (name, lastname) DO UPDATE
269		:return int: Number of affected rows. Note that this may be unreliable if `commit` is `False`
270		"""
271		if constraints is None:
272			constraints = []
273
274		# escape identifiers
275		identifiers = [sql.Identifier(column) for column in data.keys()]
276		identifiers.insert(0, sql.Identifier(table))
277
278		# prepare parameter replacements
279		protoquery = "INSERT INTO {} (%s) VALUES %%s" % ", ".join(["{}" for column in data.keys()])
280		protoquery += " ON CONFLICT"
281
282		if constraints:
283			protoquery += " (" + ", ".join(["{}" for each in constraints]) + ")"
284			identifiers.extend([sql.Identifier(column) for column in constraints])
285
286		protoquery += " DO UPDATE SET "
287		protoquery += ", ".join(["%s = EXCLUDED.%s" % (column, column) for column in data.keys()])
288		identifiers.extend(list(itertools.chain.from_iterable([[column, column] for column in data.keys()])))
289
290		query = sql.SQL(protoquery).format(*identifiers)
291		replacements = (tuple(data.values()),)
292
293		rowcount = self.execute(query, replacements=replacements, commit=commit)
294		return rowcount
295
296	def fetchall(self, query, *args):
297		"""
298		Fetch all rows for a query
299
300		:param string query:  Query
301		:param args: Replacement values
302		:param commit:  Commit transaction after query?
303		:return list: The result rows, as a list
304		"""
305		cursor = self._execute_query(query, *args)
306
307		try:
308			result = cursor.fetchall()
309		except AttributeError:
310			result = []
311		except psycopg2.ProgrammingError as e:
312			# there seems to be a bug with psycopg2 where it sometimes raises
313			# this for empty query results even though it shouldn't. this
314			# doesn't seem to indicate an actual problem so we catch the
315			# exception and return an empty list
316			self.rollback()
317			result = []
318			self.log.warning("Caught ProgrammingError: %s" % e)
319
320		cursor.close()
321		self.commit()
322
323		return result
324
325	def fetchone(self, query, *args):
326		"""
327		Fetch one result row
328
329		:param string query: Query
330		:param args: Replacement values
331		:param commit:  Commit transaction after query?
332		:return: The row, as a dictionary, or None if there were no rows
333		"""
334		cursor = self._execute_query(query, *args)
335
336		try:
337			result = cursor.fetchone()
338		except psycopg2.ProgrammingError as e:
339			# no results to fetch
340			self.rollback()
341			result = None
342			self.log.warning("Caught ProgrammingError: %s" % e)
343
344		cursor.close()
345		self.commit()
346
347		return result
348
349	def fetchall_interruptable(self, queue, query, *args):
350		"""
351		Fetch all rows for a query, allowing for interruption
352
353		Before running the query, a job is queued to cancel the query after a
354		set amount of time. The query is expected to complete before this
355		timeout. If the backend is interrupted, however, that job will be
356		executed immediately, to cancel the database query. If this happens, a
357		DatabaseQueryInterruptedException will be raised, but the database
358		object will otherwise remain useable.
359
360		Note that in the event that the cancellation job is run, all queries
361		for this instance of the database object will be cancelled. However,
362		there should never be more than one active query per connection within
363		4CAT.
364
365		:param JobQueue queue:  A job queue object, required to schedule the
366		query cancellation job
367		:param str query:  SQL query
368		:param list args:  Replacement variables
369		:param commit:  Commit transaction after query?
370		:return list:  A list of rows, as dictionaries
371		"""
372		# schedule a job that will cancel the query we're about to make
373		self.interruptable_job = queue.add_job("cancel-pg-query", details={}, remote_id=self.appname, claim_after=time.time() + self.interruptable_timeout)
374
375		# run the query
376		cursor = self.get_cursor()
377		try:
378			cursor = self._execute_query(query, cursor=cursor, *args)
379		except psycopg2.extensions.QueryCanceledError:
380			# interrupted with cancellation worker (or manually)
381			self.log.debug("Query in connection %s was interrupted..." % self.appname)
382			self.rollback()
383			cursor.close()
384			raise DatabaseQueryInterruptedException("Interrupted while querying database")
385
386		# collect results
387		try:
388			result = cursor.fetchall()
389		except AttributeError:
390			result = []
391		except psycopg2.ProgrammingError as e:
392			result = []
393			self.log.warning("Caught ProgrammingError: %s" % e)
394
395		# clean up cancelling job when we have the data
396		self.interruptable_job.finish()
397		self.interruptable_job = None
398
399		cursor.close()
400		self.commit()
401
402		return result
403
404
405	def commit(self):
406		"""
407		Commit the current transaction
408
409		This is required for UPDATE etc to stick.
410		"""
411		self.connection.commit()
412
413	def rollback(self):
414		"""
415		Roll back the current transaction
416		"""
417		self.connection.rollback()
418
419	def close(self):
420		"""
421		Close connection
422
423		Running queries after this is probably a bad idea!
424		"""
425		self.connection.close()
426
427	def get_cursor(self):
428		"""
429		Get a new cursor
430
431		Re-using cursors seems to give issues when using per-thread
432		connections, so simply instantiate a new one each time
433
434		:return: Cursor
435		"""
436		try:
437			return self.connection.cursor(cursor_factory=psycopg2.extras.RealDictCursor)
438		except (psycopg2.InterfaceError, psycopg2.OperationalError) as e:
439			self.log.warning(f"Database Exception: {e}\nReconnecting and retrying query...")
440			self.reconnect()
441			return self.connection.cursor(cursor_factory=psycopg2.extras.RealDictCursor)