fix: connection pool error fix for setup db

frappe/database/postgres/database.py

@@ -126,47 +126,79 @@ def is_interface_error(e):
 class ConnectionPool:
 	_connection_pool = None
 	_lock = threading.Lock()
+	_init_lock = threading.Lock()
 
 	@classmethod
-	def _initialize(cls, conn_settings = {}):
-		if not cls._connection_pool:
-			with cls._lock:  
-				if cls._connection_pool is None and conn_settings:
-					print("Initializing Connection Pool")
-					cls._connection_pool = pool.ThreadedConnectionPool(
-						minconn=5,
-						maxconn=100,
-						**conn_settings
-					)
+	def _initialize(cls, conn_settings={}):
+		with cls._init_lock:  
+			if cls._connection_pool is None and conn_settings:
+				print("Initializing Connection Pool")
+				cls._connection_pool = pool.ThreadedConnectionPool(
+					minconn=5,
+					maxconn=100,
+					**conn_settings
+				)
 		return cls._connection_pool
-	
+
 	@classmethod
-	def get_connection(cls, conn_settings = {}):
-		if not cls._connection_pool:
+	def get_connection(cls, conn_settings={}):
+		if cls._connection_pool is None:
 			cls._initialize(conn_settings)
 		try:
-			conn =  cls._connection_pool.getconn()
+			conn = cls._connection_pool.getconn()
+			if not cls._is_connection_valid(conn):
+				cls._invalidate_pool(conn_settings)
+				conn = cls._connection_pool.getconn()
 		except Exception as e:
-			conn = psycopg2.connect(**conn_settings)
+			print(f"Error getting connection: {e}")
+			conn = cls._create_new_connection(conn_settings)
 		return conn
-	
+
 	@classmethod
-	def put_connection(cls, conn):
+	def _is_connection_valid(cls, conn):
 		try:
-			cls._connection_pool.putconn(conn)
+			with conn.cursor() as cursor:
+				cursor.execute("SELECT 1")
+			return True
+		except Exception:
+			return False
+
+	@classmethod
+	def _invalidate_pool(cls, conn_settings):
+		with cls._lock:
+			print("Invalidating connection pool")
+			cls._connection_pool = None  # Invalidate the pool
+			cls._initialize(conn_settings)  # Reinitialize the pool
+			print("Finished reinitializing the pool")
+
+	@classmethod
+	def _create_new_connection(cls, conn_settings):
+		try:
+			return psycopg2.connect(**conn_settings)
 		except Exception as e:
-			conn.close()
+			print(f"Error creating new connection: {e}")
+			return None
 
 	@classmethod
-	def get_connection_pool(cls, conn_settings = {}):
-		if not cls._connection_pool:
-			cls._initialize(conn_settings)
+	def put_connection(cls, conn):
+		if conn:
+			try:
+				cls._connection_pool.putconn(conn)
+			except Exception as e:
+				print(f"Error returning connection to pool: {e}")
+				conn.close()
+
+	@classmethod
+	def get_connection_pool(cls):
+		if cls._connection_pool is None:
+			cls._initialize()
 		return cls._connection_pool
 
 	@classmethod
 	def close_all_connections(cls):
 		if cls._connection_pool:
 			cls._connection_pool.closeall()
+			cls._connection_pool = None  # Reset pool after closing
 
 class PostgresDatabase(PostgresExceptionUtil, Database):
 	REGEX_CHARACTER = "~"

frappe/database/postgres/setup_db.py

@@ -9,6 +9,19 @@ def setup_database(force, source_sql=None, verbose=False):
 	root_conn = get_root_connection(frappe.flags.root_login, frappe.flags.root_password)
 	root_conn.commit()
 	root_conn.sql("end")
+
+	# Terminate existing connections
+	terminate_query = f"""
+	SELECT pg_terminate_backend(pg_stat_activity.pid)
+	FROM pg_stat_activity
+	WHERE pg_stat_activity.datname = '{frappe.conf.db_name}'
+	AND pid <> pg_backend_pid();
+	"""
+	try:
+		root_conn.sql(terminate_query)
+	except Exception as e:
+		print(f"Error executing terminate_query: {e}")
+
 	root_conn.sql(f'DROP DATABASE IF EXISTS "{frappe.conf.db_name}"')
 
 	# If user exists, just update password