Migrate to latest stores

Author: hbcarlos

jupyter_collaboration/app.py

@@ -3,10 +3,12 @@
 from __future__ import annotations
 
 import asyncio
+from anyio import create_task_group
+from anyio.abc import TaskGroup
 
 from jupyter_server.extension.application import ExtensionApp
 from traitlets import Bool, Float, Type
-from ypy_websocket.ystore import BaseYStore
+from ypy_websocket.stores import BaseYStore
 
 from .handlers import DocSessionHandler, YDocWebSocketHandler
 from .loaders import FileLoaderMapping
@@ -22,6 +24,8 @@ class YDocExtension(ExtensionApp):
     Enables Real Time Collaboration in JupyterLab
     """
 
+    _store: BaseYStore = None
+
     disable_rtc = Bool(False, config=True, help="Whether to disable real time collaboration.")
 
     file_poll_interval = Float(
@@ -79,11 +83,13 @@ def initialize_handlers(self):
         # Set configurable parameters to YStore class
         for k, v in self.config.get(self.ystore_class.__name__, {}).items():
             setattr(self.ystore_class, k, v)
+        
+        # Instantiate the store
+        self._store = self.ystore_class(path="jupyter_updates.db", log=self.log)
 
         self.ywebsocket_server = JupyterWebsocketServer(
             rooms_ready=False,
             auto_clean_rooms=False,
-            ystore_class=self.ystore_class,
             log=self.log,
         )
 
@@ -103,7 +109,7 @@ def initialize_handlers(self):
                         "document_cleanup_delay": self.document_cleanup_delay,
                         "document_save_delay": self.document_save_delay,
                         "file_loaders": self.file_loaders,
-                        "ystore_class": self.ystore_class,
+                        "store": self._store,
                         "ywebsocket_server": self.ywebsocket_server,
                     },
                 ),
@@ -120,3 +126,6 @@ async def stop_extension(self):
             ],
             timeout=3,
         )
+
+        if self._store is not None and self._store.started.is_set():
+            self._store.stop()

jupyter_collaboration/handlers.py

@@ -15,7 +15,7 @@
 from tornado import web
 from tornado.websocket import WebSocketHandler
 from ypy_websocket.websocket_server import YRoom
-from ypy_websocket.ystore import BaseYStore
+from ypy_websocket.stores import BaseYStore
 from ypy_websocket.yutils import YMessageType, write_var_uint
 
 from .loaders import FileLoaderMapping
@@ -62,6 +62,15 @@ def create_task(self, aw):
         task.add_done_callback(self._background_tasks.discard)
 
     async def prepare(self):
+        # NOTE: Initialize in the ExtensionApp.start_extension once 
+        # https://github.com/jupyter-server/jupyter_server/issues/1329
+        # is done.
+        # We are temporarily initializing the store here because `start``
+        # is an async function
+        if self._store is not None and not self._store.started.is_set():
+            await self._store.start()
+            await self._store.initialize()
+
         if not self._websocket_server.started.is_set():
             self.create_task(self._websocket_server.start())
             await self._websocket_server.started.wait()
@@ -84,15 +93,13 @@ async def prepare(self):
                     )
 
                 file = self._file_loaders[file_id]
-                updates_file_path = f".{file_type}:{file_id}.y"
-                ystore = self._ystore_class(path=updates_file_path, log=self.log)
                 self.room = DocumentRoom(
                     self._room_id,
                     file_format,
                     file_type,
                     file,
                     self.event_logger,
-                    ystore,
+                    self._store,
                     self.log,
                     self._document_save_delay,
                 )
@@ -111,15 +118,15 @@ def initialize(
         self,
         ywebsocket_server: JupyterWebsocketServer,
         file_loaders: FileLoaderMapping,
-        ystore_class: type[BaseYStore],
+        store: BaseYStore,
         document_cleanup_delay: float | None = 60.0,
         document_save_delay: float | None = 1.0,
     ) -> None:
         self._background_tasks = set()
         # File ID manager cannot be passed as argument as the extension may load after this one
         self._file_id_manager = self.settings["file_id_manager"]
         self._file_loaders = file_loaders
-        self._ystore_class = ystore_class
+        self._store = store
         self._cleanup_delay = document_cleanup_delay
         self._document_save_delay = document_save_delay
         self._websocket_server = ywebsocket_server

jupyter_collaboration/rooms.py

@@ -11,7 +11,7 @@
 from jupyter_events import EventLogger
 from jupyter_ydoc import ydocs as YDOCS
 from ypy_websocket.websocket_server import YRoom
-from ypy_websocket.ystore import BaseYStore, YDocNotFound
+from ypy_websocket.stores import BaseYStore, YDocNotFound
 from ypy_websocket.yutils import write_var_uint
 
 from .loaders import FileLoader
@@ -104,56 +104,40 @@ async def initialize(self) -> None:
                 return
 
             self.log.info("Initializing room %s", self._room_id)
-
             model = await self._file.load_content(self._file_format, self._file_type, True)
 
             async with self._update_lock:
                 # try to apply Y updates from the YStore for this document
-                read_from_source = True
-                if self.ystore is not None:
-                    try:
-                        await self.ystore.apply_updates(self.ydoc)
-                        self._emit(
-                            LogLevel.INFO,
-                            "load",
-                            "Content loaded from the store {}".format(
-                                self.ystore.__class__.__qualname__
-                            ),
-                        )
-                        self.log.info(
-                            "Content in room %s loaded from the ystore %s",
-                            self._room_id,
-                            self.ystore.__class__.__name__,
-                        )
-                        read_from_source = False
-                    except YDocNotFound:
-                        # YDoc not found in the YStore, create the document from the source file (no change history)
-                        pass
-
-                if not read_from_source:
+                if self.ystore is not None and await self.ystore.exists(self._room_id):
+                    # Load the content from the store
+                    await self.ystore.apply_updates(self._room_id, self.ydoc)
+                    self._emit(LogLevel.INFO, "load", "Content loaded from the store {}".format(self.ystore.__class__.__qualname__))
+                    self.log.info("Content in room %s loaded from the ystore %s", self._room_id, self.ystore.__class__.__name__,)
+
                     # if YStore updates and source file are out-of-sync, resync updates with source
                     if self._document.source != model["content"]:
-                        # TODO: Delete document from the store.
-                        self._emit(
-                            LogLevel.INFO, "initialize", "The file is out-of-sync with the ystore."
-                        )
-                        self.log.info(
-                            "Content in file %s is out-of-sync with the ystore %s",
-                            self._file.path,
-                            self.ystore.__class__.__name__,
-                        )
-                        read_from_source = True
-
-                if read_from_source:
+                        self._emit(LogLevel.INFO, "initialize", "The file is out-of-sync with the ystore.")
+                        self.log.info("Content in file %s is out-of-sync with the ystore %s", self._file.path, self.ystore.__class__.__name__,)
+                        
+                        doc = await self.ystore.get(self._room_id)
+                        await self.ystore.remove(self._room_id)
+                        version = 0
+                        if "version" in doc:
+                            version = doc["version"] + 1
+
+                        await self.ystore.create(self._room_id, version)
+                        await self.ystore.encode_state_as_update(self._room_id, self.ydoc)
+
+                else:
                     self._emit(LogLevel.INFO, "load", "Content loaded from disk.")
-                    self.log.info(
-                        "Content in room %s loaded from file %s", self._room_id, self._file.path
-                    )
+                    self.log.info("Content in room %s loaded from file %s", self._room_id, self._file.path)
                     self._document.source = model["content"]
 
-                    if self.ystore:
-                        await self.ystore.encode_state_as_update(self.ydoc)
+                    if self.ystore is not None:
+                        await self.ystore.create(self._room_id, 0)
+                        await self.ystore.encode_state_as_update(self._room_id, self.ydoc)
 
+                
                 self._last_modified = model["last_modified"]
                 self._document.dirty = False
                 self.ready = True

jupyter_collaboration/stores.py

@@ -3,8 +3,8 @@
 
 from traitlets import Int, Unicode
 from traitlets.config import LoggingConfigurable
-from ypy_websocket.ystore import SQLiteYStore as _SQLiteYStore
-from ypy_websocket.ystore import TempFileYStore as _TempFileYStore
+from ypy_websocket.stores import SQLiteYStore as _SQLiteYStore
+from ypy_websocket.stores import TempFileYStore as _TempFileYStore
 
 
 class TempFileYStore(_TempFileYStore):
@@ -17,7 +17,7 @@ class SQLiteYStoreMetaclass(type(LoggingConfigurable), type(_SQLiteYStore)):  #
 
 class SQLiteYStore(LoggingConfigurable, _SQLiteYStore, metaclass=SQLiteYStoreMetaclass):
     db_path = Unicode(
-        ".jupyter_ystore.db",
+        "./jupyter_ystore.db",
         config=True,
         help="""The path to the YStore database. Defaults to '.jupyter_ystore.db' in the current
         directory.""",

jupyter_collaboration/websocketserver.py

@@ -9,7 +9,7 @@
 
 from tornado.websocket import WebSocketHandler
 from ypy_websocket.websocket_server import WebsocketServer, YRoom
-from ypy_websocket.ystore import BaseYStore
+from ypy_websocket.stores import BaseYStore
 
 
 class RoomNotFound(LookupError):
@@ -27,13 +27,11 @@ class JupyterWebsocketServer(WebsocketServer):
 
     def __init__(
         self,
-        ystore_class: BaseYStore,
         rooms_ready: bool = True,
         auto_clean_rooms: bool = True,
         log: Logger | None = None,
     ):
         super().__init__(rooms_ready, auto_clean_rooms, log)
-        self.ystore_class = ystore_class
         self.ypatch_nb = 0
         self.connected_users: dict[Any, Any] = {}
         # Async loop is not yet ready at the object instantiation