lint

Signed-off-by: Marc 'risson' Schmitt <marc.schmitt@risson.space>

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 (

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))

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

authentik/tasks/encoder.py

@@ -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",
-        )

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):

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
 
 

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

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):

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()

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):