Differential D3997 Diff 14188 swh/core/db/db_utils.py

Changeset View

Standalone View

swh/core/db/db_utils.py

	Show All 36 Lines
	def jsonize(value):			def jsonize(value):
	"""Convert a value to a psycopg2 JSON object if necessary"""			"""Convert a value to a psycopg2 JSON object if necessary"""
	if isinstance(value, dict):			if isinstance(value, dict):
	return psycopg2.extras.Json(value)			return psycopg2.extras.Json(value)

	return value			return value


	def swh_db_version(			def connect_to_conninfo(
	db_or_conninfo: Union[str, psycopg2.extensions.connection]			db_or_conninfo: Union[str, psycopg2.extensions.connection]
	) -> Optional[int]:			) -> psycopg2.extensions.connection:
	"""Retrieve the swh version if any. In case of the db not initialized,			"""Connect to the database passed in argument
	this returns None. Otherwise, this returns the db's version.

	Args:			Args:
	db_or_conninfo: A database connection, or a database connection info string			db_or_conninfo: A database connection, or a database connection info string

	Returns:			Returns:
	Optional[Int]: Either the db's version or None			a connected database handle

				Raises:
				psycopg2.Error if the database doesn't exist
	"""			"""

	if isinstance(db_or_conninfo, psycopg2.extensions.connection):			if isinstance(db_or_conninfo, psycopg2.extensions.connection):
	db = db_or_conninfo			return db_or_conninfo
	else:
	try:			if "=" not in db_or_conninfo and "//" not in db_or_conninfo:
	if "=" not in db_or_conninfo:
	# Database name			# Database name
	db_or_conninfo = f"dbname={db_or_conninfo}"			db_or_conninfo = f"dbname={db_or_conninfo}"

	db = psycopg2.connect(db_or_conninfo)			db = psycopg2.connect(db_or_conninfo)

				return db


				def swh_db_version(
				db_or_conninfo: Union[str, psycopg2.extensions.connection]
				) -> Optional[int]:
				"""Retrieve the swh version of the database.

				If the database is not initialized, this logs a warning and returns None.

				Args:
				db_or_conninfo: A database connection, or a database connection info string

				Returns:
				Either the version of the database, or None if it couldn't be detected
				"""
				try:
				db = connect_to_conninfo(db_or_conninfo)
	except psycopg2.Error:			except psycopg2.Error:
	logger.exception("Failed to connect to `%s`", db_or_conninfo)			logger.exception("Failed to connect to `%s`", db_or_conninfo)
	# Database not initialized			# Database not initialized
	return None			return None

	try:			try:
	with db.cursor() as c:			with db.cursor() as c:
	query = "select version from dbversion order by dbversion desc limit 1"			query = "select version from dbversion order by dbversion desc limit 1"
	try:			try:
	c.execute(query)			c.execute(query)
	return c.fetchone()[0]			return c.fetchone()[0]
	except psycopg2.errors.UndefinedTable:			except psycopg2.errors.UndefinedTable:
	return None			return None
	except Exception:			except Exception:
	logger.exception("Could not get version from `%s`", db_or_conninfo)			logger.exception("Could not get version from `%s`", db_or_conninfo)
	return None			return None


				def swh_db_flavor(
				db_or_conninfo: Union[str, psycopg2.extensions.connection]
				) -> Optional[str]:
				"""Retrieve the swh flavor of the database.

				If the database is not initialized, or the database doesn't support
				flavors, this returns None.

				Args:
				db_or_conninfo: A database connection, or a database connection info string

				Returns:
				The flavor of the database, or None if it could not be detected.
				"""
				try:
				db = connect_to_conninfo(db_or_conninfo)
				except psycopg2.Error:
				logger.exception("Failed to connect to `%s`", db_or_conninfo)
				# Database not initialized
				return None

				try:
				with db.cursor() as c:
				query = "select swh_get_dbflavor()"
				try:
				c.execute(query)
				return c.fetchone()[0]
				except psycopg2.errors.UndefinedFunction:
				# function not found: no flavor
				return None
				except Exception:
				logger.exception("Could not get flavor from `%s`", db_or_conninfo)
				return None


	# The following code has been imported from psycopg2, version 2.7.4,			# The following code has been imported from psycopg2, version 2.7.4,
	# https://github.com/psycopg/psycopg2/tree/5afb2ce803debea9533e293eef73c92ffce95bcd			# https://github.com/psycopg/psycopg2/tree/5afb2ce803debea9533e293eef73c92ffce95bcd
	# and modified by Software Heritage.			# and modified by Software Heritage.
	#			#
	# Original file: lib/extras.py			# Original file: lib/extras.py
	#			#
	# psycopg2 is free software: you can redistribute it and/or modify it under the			# psycopg2 is free software: you can redistribute it and/or modify it under the
	# terms of the GNU Lesser General Public License as published by the Free			# terms of the GNU Lesser General Public License as published by the Free
	▲ Show 20 Lines • Show All 106 Lines • Show Last 20 Lines