From dbc4a2b7307d0bc2c74a644d69b55ac9c598e822 Mon Sep 17 00:00:00 2001 From: Marc 'risson' Schmitt Date: Mon, 24 Mar 2025 12:51:20 +0100 Subject: [PATCH] lint Signed-off-by: Marc 'risson' Schmitt --- authentik/admin/tests/test_tasks.py | 1 - authentik/tasks/apps.py | 8 +++-- authentik/tasks/broker.py | 7 ++-- authentik/tasks/encoder.py | 32 ------------------- authentik/tasks/management/commands/worker.py | 6 ++-- authentik/tasks/middleware.py | 2 ++ authentik/tasks/models.py | 4 +-- authentik/tasks/results.py | 2 +- authentik/tasks/setup.py | 4 +-- authentik/tasks/tests.py | 2 +- 10 files changed, 19 insertions(+), 49 deletions(-) delete mode 100644 authentik/tasks/encoder.py diff --git a/authentik/admin/tests/test_tasks.py b/authentik/admin/tests/test_tasks.py index 1e8bc43c07..716168eb21 100644 --- a/authentik/admin/tests/test_tasks.py +++ b/authentik/admin/tests/test_tasks.py @@ -1,7 +1,6 @@ """test admin tasks""" from django.core.cache import cache -from django.test import TestCase from requests_mock import Mocker from authentik.admin.tasks import ( diff --git a/authentik/tasks/apps.py b/authentik/tasks/apps.py index 8f9ab2b034..de317d979f 100644 --- a/authentik/tasks/apps.py +++ b/authentik/tasks/apps.py @@ -1,8 +1,10 @@ from datetime import timedelta + import dramatiq -from dramatiq.middleware import AgeLimit, Callbacks, Prometheus, Retries, TimeLimit +from dramatiq.encoder import PickleEncoder +from dramatiq.middleware import AgeLimit, Callbacks, Retries, TimeLimit + from authentik.blueprints.apps import ManagedAppConfig -from authentik.tasks.encoder import JSONPickleEncoder class AuthentikTasksConfig(ManagedAppConfig): @@ -15,7 +17,7 @@ class AuthentikTasksConfig(ManagedAppConfig): from authentik.tasks.broker import PostgresBroker from authentik.tasks.middleware import CurrentTask - dramatiq.set_encoder(JSONPickleEncoder()) + dramatiq.set_encoder(PickleEncoder()) broker = PostgresBroker() # broker.add_middleware(Prometheus()) broker.add_middleware(AgeLimit(max_age=timedelta(days=30).total_seconds() * 1000)) diff --git a/authentik/tasks/broker.py b/authentik/tasks/broker.py index 814295fcb2..21afbc2f8f 100644 --- a/authentik/tasks/broker.py +++ b/authentik/tasks/broker.py @@ -1,5 +1,3 @@ -from dramatiq.middleware import Middleware -from psycopg import sql import functools import logging import time @@ -24,13 +22,14 @@ from dramatiq.broker import Broker, Consumer, MessageProxy from dramatiq.common import compute_backoff, current_millis, dq_name, xq_name from dramatiq.errors import ConnectionError, QueueJoinTimeout from dramatiq.message import Message +from dramatiq.middleware import Middleware from dramatiq.results import Results from pglock.core import _cast_lock_id -from psycopg import Notify +from psycopg import Notify, sql from psycopg.errors import AdminShutdown from structlog.stdlib import get_logger -from authentik.tasks.models import Task, CHANNEL_PREFIX, ChannelIdentifier, TaskState +from authentik.tasks.models import CHANNEL_PREFIX, ChannelIdentifier, Task, TaskState from authentik.tasks.results import PostgresBackend from authentik.tenants.models import Tenant from authentik.tenants.utils import get_current_tenant diff --git a/authentik/tasks/encoder.py b/authentik/tasks/encoder.py deleted file mode 100644 index c6fb399fe6..0000000000 --- a/authentik/tasks/encoder.py +++ /dev/null @@ -1,32 +0,0 @@ -import jsonpickle -import dramatiq.encoder -from typing import Any -from dramatiq.encoder import MessageData -import orjson - - -class OrjsonBackend(jsonpickle.JSONBackend): - def encode(self, obj: Any, indent=None, separators=None) -> str: - return orjson.dumps(obj, option=orjson.OPT_NON_STR_KEYS).decode("utf-8") - - def decode(self, string: str) -> Any: - return orjson.loads(string) - - -class JSONPickleEncoder(dramatiq.encoder.Encoder): - def encode(self, data: MessageData) -> bytes: - return jsonpickle.encode( - data, - backend=OrjsonBackend(), - keys=True, - warn=True, - use_base85=True, - ).encode() - - def decode(self, data: bytes) -> MessageData: - return jsonpickle.decode( - data.decode(), - backend=OrjsonBackend(), - keys=True, - on_missing="warn", - ) diff --git a/authentik/tasks/management/commands/worker.py b/authentik/tasks/management/commands/worker.py index 93b48bf356..ca74d4b416 100644 --- a/authentik/tasks/management/commands/worker.py +++ b/authentik/tasks/management/commands/worker.py @@ -1,10 +1,10 @@ import os - -from django.utils.module_loading import module_has_submodule -from authentik.lib.utils.reflection import get_apps import sys from django.core.management.base import BaseCommand +from django.utils.module_loading import module_has_submodule + +from authentik.lib.utils.reflection import get_apps class Command(BaseCommand): diff --git a/authentik/tasks/middleware.py b/authentik/tasks/middleware.py index 1c1c9de30d..0bcf237030 100644 --- a/authentik/tasks/middleware.py +++ b/authentik/tasks/middleware.py @@ -1,6 +1,8 @@ import contextvars + from dramatiq.message import Message from dramatiq.middleware import Middleware + from authentik.tasks.models import Task diff --git a/authentik/tasks/models.py b/authentik/tasks/models.py index 8d01e3c66e..6c4eccb159 100644 --- a/authentik/tasks/models.py +++ b/authentik/tasks/models.py @@ -1,10 +1,10 @@ -from django.utils.translation import gettext_lazy as _ from enum import StrEnum, auto from uuid import uuid4 -import pgtrigger +import pgtrigger from django.db import models from django.utils import timezone +from django.utils.translation import gettext_lazy as _ from authentik.lib.models import SerializerModel from authentik.tenants.models import Tenant diff --git a/authentik/tasks/results.py b/authentik/tasks/results.py index ff6c7d2bf6..8b5392cbc8 100644 --- a/authentik/tasks/results.py +++ b/authentik/tasks/results.py @@ -4,7 +4,7 @@ from django.utils import timezone from dramatiq.message import Message, get_encoder from dramatiq.results.backend import Missing, MResult, Result, ResultBackend -from authentik.tasks.models import Task, TaskState +from authentik.tasks.models import Task class PostgresBackend(ResultBackend): diff --git a/authentik/tasks/setup.py b/authentik/tasks/setup.py index 92b9fb83bb..5869bf8107 100644 --- a/authentik/tasks/setup.py +++ b/authentik/tasks/setup.py @@ -2,11 +2,11 @@ import os import sys import warnings -from authentik.lib.config import CONFIG from cryptography.hazmat.backends.openssl.backend import backend from defusedxml import defuse_stdlib from django.utils.autoreload import DJANGO_AUTORELOAD_ENV +from authentik.lib.config import CONFIG from lifecycle.migrate import run_migrations from lifecycle.wait_for_db import wait_for_db @@ -40,6 +40,6 @@ if ( ): run_migrations() -import django +import django # noqa: E402 django.setup() diff --git a/authentik/tasks/tests.py b/authentik/tasks/tests.py index 0a1e7b01e3..4a82f623a8 100644 --- a/authentik/tasks/tests.py +++ b/authentik/tasks/tests.py @@ -1,5 +1,5 @@ -from dramatiq import Worker, get_broker from django.test import TransactionTestCase +from dramatiq import Worker, get_broker class TaskTestCase(TransactionTestCase):