<commit_before>from collections import defaultdict
from typing import TYPE_CHECKING, Dict, List, Tuple, Union

import mariadb
from mariadb.constants import FIELD_TYPE
from pymysql.constants import ER
from pymysql.converters import escape_string

import frappe
from frappe.database.database import Database
from frappe.database.mariadb.schema import MariaDBTable
from frappe.utils import UnicodeWithAttrs, get_datetime, get_table_name

if TYPE_CHECKING:
	from mariadb import ConnectionPool

_SITE_POOLS = defaultdict(frappe._dict)
_MAX_POOL_SIZE = 64
_POOL_SIZE = 1

# _POOL_SIZE is selected "arbitrarily" to avoid overloading the server and being mindful of multitenancy
# init size of connection pool will be _POOL_SIZE for each site. Replica setups will have separate pool.
# This means each site with a replica setup can have 2 active pools of size _POOL_SIZE each. Each pool may
# expand up to _MAX_POOL_SIZE as per requirement. This cannot be a function of @@global.max_connections,
# no. of sites since there may be multiple processes holding connections; and this defines the size for each
# of those processes/workers. Check MariaDBConnectionUtil for connection & pool management.


def is_connection_pooling_enabled() -> bool:
	"""Set `frappe.DISABLE_CONNECTION_POOLING` to enable/disable connection pooling for all on current
	process. This will override config key `disable_database_connection_pooling`. Set key
	`disable_database_connection_pooling` in site config for persistent settings across workers."""

	if frappe.DISABLE_DATABASE_POOLING is not None:
		return not frappe.DISABLE_DATABASE_POOLING
	return frappe.local.conf.disable_database_connection_pooling


class MariaDBExceptionUtil:
	ProgrammingError = mariadb.ProgrammingError
	TableMissingError = mariadb.ProgrammingError
	OperationalError = mariadb.OperationalError
	InternalError = mariadb.InternalError
	SQLError = mariadb.ProgrammingError
	DataError = mariadb.DataError

	@staticmethod
	def is_deadlocked(e: mariadb.Error) -> bool:
		return getattr(e, "errno", None) == ER.LOCK_DEADLOCK

	@staticmethod
	def is_timedout(e: mariadb.Error) -> bool:
		return getattr(e, "errno", None) == ER.LOCK_WAIT_TIMEOUT

	@staticmethod
	def is_table_missing(e: mariadb.Error) -> bool:
		return getattr(e, "errno", None) == ER.NO_SUCH_TABLE

	@staticmethod
	def is_missing_table(e: mariadb.Error) -> bool:
		return MariaDBDatabase.is_table_missing(e)

	@staticmethod
	def is_missing_column(e: mariadb.Error) -> bool:
		return getattr(e, "errno", None) == ER.BAD_FIELD_ERROR

	@staticmethod
	def is_duplicate_fieldname(e: mariadb.Error) -> bool:
		return getattr(e, "errno", None) == ER.DUP_FIELDNAME

	@staticmethod
	def is_duplicate_entry(e: mariadb.Error) -> bool:
		return getattr(e, "errno", None) == ER.DUP_ENTRY

	@staticmethod
	def is_access_denied(e: mariadb.Error) -> bool:
		return getattr(e, "errno", None) == ER.ACCESS_DENIED_ERROR

	@staticmethod
	def cant_drop_field_or_key(e: mariadb.Error) -> bool:
		return getattr(e, "errno", None) == ER.CANT_DROP_FIELD_OR_KEY

	@staticmethod
	def is_syntax_error(e: mariadb.Error) -> bool:
		return getattr(e, "errno", None) == ER.PARSE_ERROR

	@staticmethod
	def is_data_too_long(e: mariadb.Error) -> bool:
		return getattr(e, "errno", None) == ER.DATA_TOO_LONG

	@staticmethod
	def is_primary_key_violation(e: mariadb.Error) -> bool:
		return (
			MariaDBDatabase.is_duplicate_entry(e)
			and "PRIMARY" in e.errmsg
			and isinstance(e, mariadb.IntegrityError)
		)

	@staticmethod
	def is_unique_key_violation(e: mariadb.Error) -> bool:
		return (
			MariaDBDatabase.is_duplicate_entry(e)
			and "Duplicate" in e.errmsg
			and isinstance(e, mariadb.IntegrityError)
		)


class MariaDBConnectionUtil:
	def get_connection(self):
		conn = self._get_connection()
		conn.auto_reconnect = True
		return conn

	def _get_connection(self) -> "mariadb.Connection":
		"""Return MariaDB connection object.

		If frappe.conf.disable_database_connection_pooling is set, return a new connection
		object and close existing pool if exists. Else, return a connection from the pool.
		"""
		global _SITE_POOLS

		# don't pool root connections
		if self.user == "root":
			return self.create_connection()

		if is_connection_pooling_enabled():
			self.close_connection_pools()
			return self.create_connection()

		read_only = frappe.conf.read_from_replica and frappe.conf.replica_host

		if frappe.local.site not in _SITE_POOLS:
			site_pool = self.create_connection_pool(read_only=read_only)
		else:
			site_pool = self.get_connection_pool(read_only=read_only)

		try:
			conn = site_pool.get_connection()
		except mariadb.PoolError:
			# PoolError is raised when the pool is exhausted
			conn = self.create_connection()
			try:
				site_pool.add_connection(conn)
				# log this via frappe.logger & continue - site needs bigger pool...over _POOL_SIZE
			except mariadb.PoolError:
				# PoolError is raised when size limit is reached
				# log this via frappe.logger & continue - site needs a much bigger pool...over _MAX_POOL_SIZE
				pass

		return conn

	def close_connection_pools(self):
		if frappe.local.site in _SITE_POOLS:
			pools = _SITE_POOLS[frappe.local.site]<commit_msg>fix: Correct use of is_connection_pooling_enabled check<commit_after>from collections import defaultdict
from typing import TYPE_CHECKING, Dict, List, Tuple, Union

import mariadb
from mariadb.constants import FIELD_TYPE
from pymysql.constants import ER
from pymysql.converters import escape_string

import frappe
from frappe.database.database import Database
from frappe.database.mariadb.schema import MariaDBTable
from frappe.utils import UnicodeWithAttrs, get_datetime, get_table_name

if TYPE_CHECKING:
	from mariadb import ConnectionPool

_SITE_POOLS = defaultdict(frappe._dict)
_MAX_POOL_SIZE = 64
_POOL_SIZE = 1

# _POOL_SIZE is selected "arbitrarily" to avoid overloading the server and being mindful of multitenancy
# init size of connection pool will be _POOL_SIZE for each site. Replica setups will have separate pool.
# This means each site with a replica setup can have 2 active pools of size _POOL_SIZE each. Each pool may
# expand up to _MAX_POOL_SIZE as per requirement. This cannot be a function of @@global.max_connections,
# no. of sites since there may be multiple processes holding connections; and this defines the size for each
# of those processes/workers. Check MariaDBConnectionUtil for connection & pool management.


def is_connection_pooling_enabled() -> bool:
	"""Set `frappe.DISABLE_CONNECTION_POOLING` to enable/disable connection pooling for all on current
	process. This will override config key `disable_database_connection_pooling`. Set key
	`disable_database_connection_pooling` in site config for persistent settings across workers."""

	if frappe.DISABLE_DATABASE_POOLING is not None:
		return not frappe.DISABLE_DATABASE_POOLING
	return frappe.local.conf.disable_database_connection_pooling


class MariaDBExceptionUtil:
	ProgrammingError = mariadb.ProgrammingError
	TableMissingError = mariadb.ProgrammingError
	OperationalError = mariadb.OperationalError
	InternalError = mariadb.InternalError
	SQLError = mariadb.ProgrammingError
	DataError = mariadb.DataError

	@staticmethod
	def is_deadlocked(e: mariadb.Error) -> bool:
		return getattr(e, "errno", None) == ER.LOCK_DEADLOCK

	@staticmethod
	def is_timedout(e: mariadb.Error) -> bool:
		return getattr(e, "errno", None) == ER.LOCK_WAIT_TIMEOUT

	@staticmethod
	def is_table_missing(e: mariadb.Error) -> bool:
		return getattr(e, "errno", None) == ER.NO_SUCH_TABLE

	@staticmethod
	def is_missing_table(e: mariadb.Error) -> bool:
		return MariaDBDatabase.is_table_missing(e)

	@staticmethod
	def is_missing_column(e: mariadb.Error) -> bool:
		return getattr(e, "errno", None) == ER.BAD_FIELD_ERROR

	@staticmethod
	def is_duplicate_fieldname(e: mariadb.Error) -> bool:
		return getattr(e, "errno", None) == ER.DUP_FIELDNAME

	@staticmethod
	def is_duplicate_entry(e: mariadb.Error) -> bool:
		return getattr(e, "errno", None) == ER.DUP_ENTRY

	@staticmethod
	def is_access_denied(e: mariadb.Error) -> bool:
		return getattr(e, "errno", None) == ER.ACCESS_DENIED_ERROR

	@staticmethod
	def cant_drop_field_or_key(e: mariadb.Error) -> bool:
		return getattr(e, "errno", None) == ER.CANT_DROP_FIELD_OR_KEY

	@staticmethod
	def is_syntax_error(e: mariadb.Error) -> bool:
		return getattr(e, "errno", None) == ER.PARSE_ERROR

	@staticmethod
	def is_data_too_long(e: mariadb.Error) -> bool:
		return getattr(e, "errno", None) == ER.DATA_TOO_LONG

	@staticmethod
	def is_primary_key_violation(e: mariadb.Error) -> bool:
		return (
			MariaDBDatabase.is_duplicate_entry(e)
			and "PRIMARY" in e.errmsg
			and isinstance(e, mariadb.IntegrityError)
		)

	@staticmethod
	def is_unique_key_violation(e: mariadb.Error) -> bool:
		return (
			MariaDBDatabase.is_duplicate_entry(e)
			and "Duplicate" in e.errmsg
			and isinstance(e, mariadb.IntegrityError)
		)


class MariaDBConnectionUtil:
	def get_connection(self):
		conn = self._get_connection()
		conn.auto_reconnect = True
		return conn

	def _get_connection(self) -> "mariadb.Connection":
		"""Return MariaDB connection object.

		If frappe.conf.disable_database_connection_pooling is set, return a new connection
		object and close existing pool if exists. Else, return a connection from the pool.
		"""
		global _SITE_POOLS

		# don't pool root connections
		if self.user == "root":
			return self.create_connection()

		if not is_connection_pooling_enabled():
			self.close_connection_pools()
			return self.create_connection()

		read_only = frappe.conf.read_from_replica and frappe.conf.replica_host

		if frappe.local.site not in _SITE_POOLS:
			site_pool = self.create_connection_pool(read_only=read_only)
		else:
			site_pool = self.get_connection_pool(read_only=read_only)

		try:
			conn = site_pool.get_connection()
		except mariadb.PoolError:
			# PoolError is raised when the pool is exhausted
			conn = self.create_connection()
			try:
				site_pool.add_connection(conn)
				# log this via frappe.logger & continue - site needs bigger pool...over _POOL_SIZE
			except mariadb.PoolError:
				# PoolError is raised when size limit is reached
				# log this via frappe.logger & continue - site needs a much bigger pool...over _MAX_POOL_SIZE
				pass

		return conn

	def close_connection_pools(self):
		if frappe.local.site in _SITE_POOLS:
			pools = _SITE_POOLS[frappe.local.site]
{'fraction_non_alphanumeric': 0.08131720430107527, 'fraction_numerical': 0.0007680491551459293, 'mean_word_length': 5.02486986697513, 'pattern_counts': {'":': 4, '<': 3, '<?xml version=': 0, '>': 33, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '4920244', 'n_tokens_mistral': 3659, 'n_tokens_neox': 3517, 'n_words': 937}