Add DjangoApplicationContext to handle sync_to_async transition for handler function

Author: sdc50

bokeh_django/__init__.py

@@ -3,6 +3,7 @@
 
 # Bokeh imports
 from .apps import DjangoBokehConfig
+from .consumers import AutoloadJsConsumer, WSConsumer
 from .routing import autoload, directory, document
 from .static import static_extensions
 

bokeh_django/consumers.py

@@ -128,7 +128,10 @@ async def _get_session(self) -> ServerSession:
                                    signed=False,
                                    expiration=300,
                                    extra_payload=payload)
-        session = await self.application_context.create_session_if_needed(session_id, self.request, token)
+        try:
+            session = await self.application_context.create_session_if_needed(session_id, self.request, token)
+        except Exception as e:
+            log.exception(e)
         return session
 
 

bokeh_django/routing.py

@@ -21,20 +21,25 @@
 import re
 from pathlib import Path
 from typing import Callable, List, Union
+import weakref
 
 # External imports
 from django.core.asgi import get_asgi_application
 from django.urls import re_path
 from django.urls.resolvers import URLPattern
+from channels.db import database_sync_to_async
+from tornado import gen
 
 # Bokeh imports
 from bokeh.application import Application
 from bokeh.application.handlers.document_lifecycle import DocumentLifecycleHandler
 from bokeh.application.handlers.function import FunctionHandler
 from bokeh.command.util import build_single_handler_application, build_single_handler_applications
-from bokeh.server.contexts import ApplicationContext
+from bokeh.server.contexts import ApplicationContext, BokehSessionContext, _RequestProxy, ServerSession
+from bokeh.document import Document
+from bokeh.util.token import get_token_payload
 
-# Bokeh imports
+# Local imports
 from .consumers import AutoloadJsConsumer, DocConsumer, WSConsumer
 
 # -----------------------------------------------------------------------------
@@ -51,6 +56,70 @@
 # General API
 # -----------------------------------------------------------------------------
 
+class DjangoApplicationContext(ApplicationContext):
+    async def create_session_if_needed(self, session_id: ID, request: HTTPServerRequest | None = None,
+            token: str | None = None) -> ServerSession:
+        # this is because empty session_ids would be "falsey" and
+        # potentially open up a way for clients to confuse us
+        if len(session_id) == 0:
+            raise ProtocolError("Session ID must not be empty")
+
+        if session_id not in self._sessions and \
+           session_id not in self._pending_sessions:
+            future = self._pending_sessions[session_id] = gen.Future()
+
+            doc = Document()
+
+            session_context = BokehSessionContext(session_id,
+                                                  self.server_context,
+                                                  doc,
+                                                  logout_url=self._logout_url)
+            if request is not None:
+                payload = get_token_payload(token) if token else {}
+                if ('cookies' in payload and 'headers' in payload
+                    and not 'Cookie' in payload['headers']):
+                    # Restore Cookie header from cookies dictionary
+                    payload['headers']['Cookie'] = '; '.join([
+                        f'{k}={v}' for k, v in payload['cookies'].items()
+                    ])
+                # using private attr so users only have access to a read-only property
+                session_context._request = _RequestProxy(request,
+                                                         cookies=payload.get('cookies'),
+                                                         headers=payload.get('headers'))
+            session_context._token = token
+
+            # expose the session context to the document
+            # use the _attribute to set the public property .session_context
+            doc._session_context = weakref.ref(session_context)
+
+            try:
+                await self._application.on_session_created(session_context)
+            except Exception as e:
+                log.error("Failed to run session creation hooks %r", e, exc_info=True)
+
+            # This needs to be wrapped in the database_sync_to_async wrapper just in case the handler function accesses
+            # Django ORM.
+
+            await database_sync_to_async(self._application.initialize_document)(doc)
+
+            session = ServerSession(session_id, doc, io_loop=self._loop, token=token)
+            del self._pending_sessions[session_id]
+            self._sessions[session_id] = session
+            session_context._set_session(session)
+            self._session_contexts[session_id] = session_context
+
+            # notify anyone waiting on the pending session
+            future.set_result(session)
+
+        if session_id in self._pending_sessions:
+            # another create_session_if_needed is working on
+            # creating this session
+            session = await self._pending_sessions[session_id]
+        else:
+            session = self._sessions[session_id]
+
+        return session
+
 
 class Routing:
     url: str
@@ -62,7 +131,7 @@ class Routing:
     def __init__(self, url: str, app: ApplicationLike, *, document: bool = False, autoload: bool = False) -> None:
         self.url = url
         self.app = self._fixup(self._normalize(app))
-        self.app_context = ApplicationContext(self.app, url=self.url)
+        self.app_context = DjangoApplicationContext(self.app, url=self.url)
         self.document = document
         self.autoload = autoload
 

setup.py

@@ -4,7 +4,7 @@
 
 setup(
     name='bokeh-django',
-    version='0.0.0',
+    version='0.1.0a2',
     description='',
     long_description='',
     keywords='',