diff --git a/dispatcharr/db/backends/postgresql_psycopg3/base.py b/dispatcharr/db/backends/postgresql_psycopg3/base.py index f5a49ffd..098a01f4 100644 --- a/dispatcharr/db/backends/postgresql_psycopg3/base.py +++ b/dispatcharr/db/backends/postgresql_psycopg3/base.py @@ -37,7 +37,10 @@ class DatabaseWrapper(DatabaseWrapperMixin, OriginalDatabaseWrapper): 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] diff --git a/dispatcharr/db/process_label.py b/dispatcharr/db/process_label.py new file mode 100644 index 00000000..bd50b6a6 --- /dev/null +++ b/dispatcharr/db/process_label.py @@ -0,0 +1,34 @@ +"""Label Postgres connections by Dispatcharr process role (for pg_stat_activity).""" + +from __future__ import annotations + +import os +import sys + + +def get_process_role(argv: list[str] | None = None) -> str: + argv = argv if argv is not None else sys.argv + argv0 = os.path.basename(argv[0]) if argv else "" + cmdline = " ".join(argv) + + if "celery" in argv0 or any("celery" in arg for arg in argv): + if "beat" in cmdline: + return "celery-beat" + if "-Q" in argv: + try: + if "dvr" in argv[argv.index("-Q") + 1]: + return "celery-dvr" + except (IndexError, ValueError): + pass + return "celery-worker" + if "daphne" in argv0: + return "daphne" + if argv0 == "manage.py" and len(argv) > 1: + return f"manage-{argv[1]}" + if os.environ.get("GEVENT_SUPPORT"): + return "uwsgi" + return "django" + + +def db_application_name() -> str: + return f"Dispatcharr-{get_process_role()}-{os.getpid()}" diff --git a/tests/test_process_label.py b/tests/test_process_label.py new file mode 100644 index 00000000..4b12b0f9 --- /dev/null +++ b/tests/test_process_label.py @@ -0,0 +1,18 @@ +from unittest.mock import patch + +from django.test import SimpleTestCase + +from dispatcharr.db.process_label import get_process_role + + +class ProcessLabelTests(SimpleTestCase): + def test_celery_worker_not_labeled_as_uwsgi(self): + role = get_process_role( + ["/dispatcharrpy/bin/celery", "-A", "dispatcharr", "worker"] + ) + self.assertEqual(role, "celery-worker") + + def test_uwsgi_labeled_when_gevent_support(self): + with patch.dict("os.environ", {"GEVENT_SUPPORT": "True"}): + role = get_process_role(["uwsgi"]) + self.assertEqual(role, "uwsgi")