Differential D6794 Diff 24639 swh/objstorage/backends/winery/database.py

Changeset View

Standalone View

swh/objstorage/backends/winery/database.py

This file was added.

				# Copyright (C) 2021 The Software Heritage developers
				# See the AUTHORS file at the top-level directory of this distribution
				# License: GNU General Public License version 3, or any later version
				# See top-level LICENSE file for more information

				import logging
				import time

				import psycopg2

				LOGGER = logging.getLogger(__name__)
				LOGGER.setLevel(logging.ERROR)


				class Database:
				def __init__(self, dsn):
				self.dsn = dsn

				def create_database(self, database):
				db = psycopg2.connect(f"{self.dsn}/postgres")
				db.autocommit = True
				c = db.cursor()
				c.execute(
				f"SELECT datname FROM pg_catalog.pg_database WHERE datname = '{database}'"
				)
				if c.rowcount == 0:
				c.execute(f"CREATE DATABASE {database}")
				c.close()

				def drop_database(self, database):
				db = psycopg2.connect(f"{self.dsn}/postgres")
				# https://wiki.postgresql.org/wiki/Psycopg2_Tutorial
				# If you want to drop the database you would need to
				# change the isolation level of the database.
				db.set_isolation_level(0)
				db.autocommit = True
				c = db.cursor()
				#
				# Dropping the database may fail because the server takes time
				# to notice a connection was dropped and/or a named cursor is
				# in the process of being deleted. It can happen here or even
				# when deleting all database with the psql cli
				# and there are no bench process active.
				#
				# ERROR: database "i606428a5a6274d1ab09eecc4d019fef7" is being
				# accessed by other users DETAIL: There is 1 other session
				# using the database.
				#
				# See:
				# https://stackoverflow.com/questions/5108876/kill-a-postgresql-session-connection
				#
				# https://www.postgresql.org/docs/current/sql-dropdatabase.html
				#
				# WITH (FORCE) added in postgresql 13 but may also fail because the
				# named cursor may not be handled as a client.
				#
				for i in range(60):
				try:
				c.execute(f"DROP DATABASE IF EXISTS {database}")
				break
				except psycopg2.errors.ObjectInUse:
				LOGGER.warning(f"{database} database drop fails, waiting")
				time.sleep(10)
				continue
				raise Exception("database drop fails {database}")
				c.close()

				def list_databases(self):
				db = psycopg2.connect(f"{self.dsn}/postgres")
				with db.cursor() as c:
				c.execute(
				"SELECT datname FROM pg_database "
				"WHERE datistemplate = false and datname != 'postgres'"
				)
				return [r[0] for r in c.fetchall()]