mirror of
https://github.com/Dispatcharr/Dispatcharr.git
synced 2026-07-17 16:50:53 +00:00
c1ff5e35aa
- Introduced a new module to determine the process role based on command-line arguments, enhancing PostgreSQL connection labeling for better monitoring. - Updated the database connection parameters to include the application name derived from the process role. - Added unit tests to validate the process role detection logic for different scenarios, including Celery and uWSGI.
47 lines
1.7 KiB
Python
47 lines
1.7 KiB
Python
try:
|
|
import psycopg
|
|
except ImportError as e:
|
|
from django.core.exceptions import ImproperlyConfigured
|
|
|
|
raise ImproperlyConfigured("Error loading psycopg3 module: %s" % e) from e
|
|
|
|
from django.db.backends.postgresql.base import DatabaseWrapper as OriginalDatabaseWrapper
|
|
from django_db_geventpool.backends.base import DatabaseWrapperMixin
|
|
|
|
from .pool import DatabaseConnectionPool
|
|
|
|
|
|
class PostgresConnectionPool(DatabaseConnectionPool):
|
|
DBERROR = psycopg.DatabaseError
|
|
|
|
def __init__(self, *args, **kwargs):
|
|
self.connect = kwargs.pop("connect", psycopg.connect)
|
|
self.connection = None
|
|
maxsize = kwargs.pop("MAX_CONNS", 4)
|
|
reuse = kwargs.pop("REUSE_CONNS", maxsize)
|
|
max_lifetime = kwargs.pop("CONN_MAX_LIFETIME", None)
|
|
self.args = args
|
|
self.kwargs = kwargs
|
|
self.kwargs["client_encoding"] = "UTF8"
|
|
super().__init__(maxsize, reuse, max_lifetime=max_lifetime)
|
|
|
|
def create_connection(self):
|
|
return self.connect(*self.args, **self.kwargs)
|
|
|
|
def check_usable(self, connection):
|
|
connection.cursor().execute("SELECT 1")
|
|
|
|
|
|
class DatabaseWrapper(DatabaseWrapperMixin, OriginalDatabaseWrapper):
|
|
pool_class = PostgresConnectionPool
|
|
INTRANS = psycopg.pq.TransactionStatus.INTRANS
|
|
|
|
def get_connection_params(self) -> dict:
|
|
from dispatcharr.db.process_label import db_application_name
|
|
|
|
conn_params = super().get_connection_params()
|
|
conn_params["application_name"] = db_application_name()
|
|
for attr in ("MAX_CONNS", "REUSE_CONNS", "CONN_MAX_LIFETIME"):
|
|
if attr in self.settings_dict["OPTIONS"]:
|
|
conn_params[attr] = self.settings_dict["OPTIONS"][attr]
|
|
return conn_params
|