# Source code for django_celery_results.backends.database

    import binascii
    import json

    from celery import maybe_signature, states
    from celery.backends.base import BaseDictBackend, get_current_task
    from celery.exceptions import ChordError
    from celery.result import GroupResult, allow_join_result, result_from_tuple
    from celery.utils.log import get_logger
    from celery.utils.serialization import b64decode, b64encode
    from django.db import connection, router, transaction
    from django.db.models.functions import Now
    from django.db.utils import InterfaceError
    from kombu.exceptions import DecodeError

    from ..models import ChordCounter
    from ..models import GroupResult as GroupResultModel
    from ..models import TaskResult

    EXCEPTIONS_TO_CATCH = (InterfaceError,)

    try:
        from psycopg2 import InterfaceError as Psycopg2InterfaceError
        EXCEPTIONS_TO_CATCH += (Psycopg2InterfaceError,)
    except ImportError:
        pass

    logger = get_logger(__name__)

    class DatabaseBackend(BaseDictBackend):
        """The Django database backend, using models to store task state."""

        TaskModel = TaskResult
        GroupModel = GroupResultModel
        subpolling_interval = 0.5

        def exception_safe_to_retry(self, exc):
            """Check if an exception is safe to retry.

            Backends have to overload this method with correct predicates
            dealing with their exceptions.

            By default no exception is safe to retry, it's up to
            backend implementation to define which exceptions are safe.

            For Celery / django-celery-results, retry Django / Psycopg2
            InterfaceErrors, like "Connection already closed", with new connection.

            Set result_backend_always_retry to True in order to enable retries.
            """
            for exc_type in EXCEPTIONS_TO_CATCH:
                if isinstance(exc, exc_type):
                    # Only called if InterfaceError occurs and always_retry is True
                    connection.close()
                    return True
            return False

        def _get_extended_properties(self, request, traceback):
            extended_props = {
                'periodic_task_name': None,
                'task_args': None,
                'task_kwargs': None,
                'task_name': None,
                'traceback': None,
                'worker': None,
            }
            if request and self.app.conf.find_value_for_key('extended', 'result'):

                if getattr(request, 'argsrepr', None) is not None:
                    # task protocol 2
                    task_args = request.argsrepr
                else:
                    # task protocol 1
                    task_args = getattr(request, 'args', None)

                if getattr(request, 'kwargsrepr', None) is not None:
                    # task protocol 2
                    task_kwargs = request.kwargsrepr
                else:
                    # task protocol 1
                    task_kwargs = getattr(request, 'kwargs', None)

                # Encode input arguments
                if task_args is not None:
                    _, _, task_args = self.encode_content(task_args)

                if task_kwargs is not None:
                    _, _, task_kwargs = self.encode_content(task_kwargs)

                periodic_task_name = getattr(request, 'periodic_task_name', None)

                extended_props.update({
                    'periodic_task_name': periodic_task_name,
                    'task_args': task_args,
                    'task_kwargs': task_kwargs,
                    'task_name': getattr(request, 'task', None),
                    'traceback': traceback,
                    'worker': getattr(request, 'hostname', None),
                })

            return extended_props

        def _get_meta_from_request(self, request=None):
            """
            Use the request or get_current_task to evaluate the `meta` attribute.

            With this, is possible to assign arbitrary data in request.meta to be
            retrieve and stored on the TaskResult.
            """
            request = request or getattr(get_current_task(), "request", None)
            return getattr(request, "meta", {})

        def _store_result(
                self,
                task_id,
                result,
                status,
                traceback=None,
                request=None,
                using=None
        ):
            """Store return value and status of an executed task."""
            content_type, content_encoding, result = self.encode_content(result)

            meta = {
                **self._get_meta_from_request(request),
                "children": self.current_task_children(request),
            }
            _, _, encoded_meta = self.encode_content(
                meta,
            )

            task_props = {
                'content_encoding': content_encoding,
                'content_type': content_type,
                'meta': encoded_meta,
                'result': result,
                'status': status,
                'task_id': task_id,
                'traceback': traceback,
                'using': using,
            }

            task_props.update(
                self._get_extended_properties(request, traceback)
            )

            if status == states.STARTED:
                task_props['date_started'] = Now()

            self.TaskModel._default_manager.store_result(**task_props)
            return result

        def _get_task_meta_for(self, task_id):
            """Get task metadata for a task by id."""
            obj = self.TaskModel._default_manager.get_task(task_id)
            res = obj.as_dict()
            meta = self.decode_content(obj, res.pop('meta', None)) or {}
            result = self.decode_content(obj, res.get('result'))

            task_args = res.get('task_args')
            task_kwargs = res.get('task_kwargs')
            try:
                task_args = self.decode_content(obj, task_args)
                task_kwargs = self.decode_content(obj, task_kwargs)
            except (DecodeError, binascii.Error):
                pass

            # the right names are args/kwargs, not task_args/task_kwargs,
            # keep both for backward compatibility
            res.update(
                meta,
                result=result,
                task_args=task_args,
                task_kwargs=task_kwargs,
                args=task_args,
                kwargs=task_kwargs,
            )
            return self.meta_from_decoded(res)

        def encode_content(self, data):
            content_type, content_encoding, content = self._encode(data)
            if content_encoding == 'binary':
                content = b64encode(content)
            return content_type, content_encoding, content

        def decode_content(self, obj, content):
            if content:
                if obj.content_encoding == 'binary':
                    content = b64decode(content)
                return self.decode(content)

        def _forget(self, task_id):
            try:
                self.TaskModel._default_manager.get(task_id=task_id).delete()
            except self.TaskModel.DoesNotExist:
                pass

        def cleanup(self):
            """Delete expired metadata."""
            if not self.expires:
                return
            self.TaskModel._default_manager.delete_expired(self.expires)
            self.GroupModel._default_manager.delete_expired(self.expires)

        def _restore_group(self, group_id):
            """return result value for a group by id."""
            group_result = self.GroupModel._default_manager.get_group(group_id)

            if group_result:
                res = group_result.as_dict()
                decoded_result = self.decode_content(group_result, res["result"])
                res["result"] = None
                if decoded_result:
                    res["result"] = result_from_tuple(decoded_result, app=self.app)
                return res

        def _save_group(self, group_id, group_result):
            """Store return value of group"""
            content_type, content_encoding, result = self.encode_content(
                group_result.as_tuple()
            )
            self.GroupModel._default_manager.store_group_result(
                content_type, content_encoding, group_id, result
            )
            return group_result

        def _delete_group(self, group_id):
            try:
                self.GroupModel._default_manager.get_group(group_id).delete()
            except self.TaskModel.DoesNotExist:
                pass

        def apply_chord(self, header_result_args, body, **kwargs):
            """Add a ChordCounter with the expected number of results"""
            if not isinstance(header_result_args, GroupResult):
                # Celery 5.1 provides the GroupResult args
                header_result = self.app.GroupResult(*header_result_args)
            else:
                # celery <5.1 will pass a GroupResult object
                header_result = header_result_args
            results = [r.as_tuple() for r in header_result]
            chord_size = body.get("chord_size", None) or len(results)
            data = json.dumps(results)
            ChordCounter.objects.create(
                group_id=header_result.id, sub_tasks=data, count=chord_size
            )

        def on_chord_part_return(self, request, state, result, **kwargs):
            """Called on finishing each part of a Chord header"""
            tid, gid = request.id, request.group
            if not gid or not tid:
                return
            call_callback = False
            with transaction.atomic(using=router.db_for_write(ChordCounter)):
                # We need to know if `count` hits 0.
                # wrap the update in a transaction
                # with a `select_for_update` lock to prevent race conditions.
                # SELECT FOR UPDATE is not supported on all databases
                try:
                    chord_counter = (
                        ChordCounter.objects.select_for_update()
                        .get(group_id=gid)
                    )
                except ChordCounter.DoesNotExist:
                    logger.warning("Can't find ChordCounter for Group %s", gid)
                    return
                chord_counter.count -= 1
                if chord_counter.count != 0:
                    chord_counter.save(update_fields=["count"])
                else:
                    # Last task in the chord header has finished
                    call_callback = True
                    chord_counter.delete()

            if call_callback:
                deps = chord_counter.group_result(app=self.app)
                if deps.ready():
                    callback = maybe_signature(request.chord, app=self.app)
                    trigger_callback(
                        app=self.app,
                        callback=callback,
                        group_result=deps
                    )

    def trigger_callback(app, callback, group_result):
        """Add the callback to the queue or mark the callback as failed
        Implementation borrowed from `celery.app.builtins.unlock_chord`
        """
        if group_result.supports_native_join:
            j = group_result.join_native
        else:
            j = group_result.join

        try:
            with allow_join_result():
                ret = j(timeout=app.conf.result_chord_join_timeout, propagate=True)
        except Exception as exc:  # pylint: disable=broad-except
            try:
                culprit = next(group_result._failed_join_report())
                reason = f"Dependency {culprit.id} raised {exc!r}"
            except StopIteration:
                reason = repr(exc)
            logger.exception("Chord %r raised: %r", group_result.id, exc)
            app.backend.chord_error_from_stack(callback, ChordError(reason))
        else:
            try:
                callback.delay(ret)
            except Exception as exc:  # pylint: disable=broad-except
                logger.exception("Chord %r raised: %r", group_result.id, exc)
                app.backend.chord_error_from_stack(
                    callback, exc=ChordError(f"Callback error: {exc!r}")
                )