Skip to content

Commit ed37ced

Browse files
hbcarloshbcarlos
committed
Migrate to latest stores
1 parent 874a7ea commit ed37ced

File tree

5 files changed

+54
-56
lines changed

5 files changed

+54
-56
lines changed

‎jupyter_collaboration/app.py

Lines changed: 12 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -3,10 +3,12 @@
33
from __future__ import annotations
44

55
import asyncio
6+
from anyio import create_task_group
7+
from anyio.abc import TaskGroup
68

79
from jupyter_server.extension.application import ExtensionApp
810
from traitlets import Bool, Float, Type
9-
from ypy_websocket.ystore import BaseYStore
11+
from ypy_websocket.stores import BaseYStore
1012

1113
from .handlers import DocSessionHandler, YDocWebSocketHandler
1214
from .loaders import FileLoaderMapping
@@ -22,6 +24,8 @@ class YDocExtension(ExtensionApp):
2224
Enables Real Time Collaboration in JupyterLab
2325
"""
2426

27+
_store: BaseYStore = None
28+
2529
disable_rtc = Bool(False, config=True, help="Whether to disable real time collaboration.")
2630

2731
file_poll_interval = Float(
@@ -79,11 +83,13 @@ def initialize_handlers(self):
7983
# Set configurable parameters to YStore class
8084
for k, v in self.config.get(self.ystore_class.__name__, {}).items():
8185
setattr(self.ystore_class, k, v)
86+
87+
# Instantiate the store
88+
self._store = self.ystore_class(path="jupyter_updates.db", log=self.log)
8289

8390
self.ywebsocket_server = JupyterWebsocketServer(
8491
rooms_ready=False,
8592
auto_clean_rooms=False,
86-
ystore_class=self.ystore_class,
8793
log=self.log,
8894
)
8995

@@ -103,7 +109,7 @@ def initialize_handlers(self):
103109
"document_cleanup_delay": self.document_cleanup_delay,
104110
"document_save_delay": self.document_save_delay,
105111
"file_loaders": self.file_loaders,
106-
"ystore_class": self.ystore_class,
112+
"store": self._store,
107113
"ywebsocket_server": self.ywebsocket_server,
108114
},
109115
),
@@ -120,3 +126,6 @@ async def stop_extension(self):
120126
],
121127
timeout=3,
122128
)
129+
130+
if self._store is not None and self._store.started.is_set():
131+
self._store.stop()

‎jupyter_collaboration/handlers.py

Lines changed: 13 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -15,7 +15,7 @@
1515
from tornado import web
1616
from tornado.websocket import WebSocketHandler
1717
from ypy_websocket.websocket_server import YRoom
18-
from ypy_websocket.ystore import BaseYStore
18+
from ypy_websocket.stores import BaseYStore
1919
from ypy_websocket.yutils import YMessageType, write_var_uint
2020

2121
from .loaders import FileLoaderMapping
@@ -62,6 +62,15 @@ def create_task(self, aw):
6262
task.add_done_callback(self._background_tasks.discard)
6363

6464
async def prepare(self):
65+
# NOTE: Initialize in the ExtensionApp.start_extension once
66+
# https://github.com/jupyter-server/jupyter_server/issues/1329
67+
# is done.
68+
# We are temporarily initializing the store here because `start``
69+
# is an async function
70+
if self._store is not None and not self._store.started.is_set():
71+
await self._store.start()
72+
await self._store.initialize()
73+
6574
if not self._websocket_server.started.is_set():
6675
self.create_task(self._websocket_server.start())
6776
await self._websocket_server.started.wait()
@@ -84,15 +93,13 @@ async def prepare(self):
8493
)
8594

8695
file = self._file_loaders[file_id]
87-
updates_file_path = f".{file_type}:{file_id}.y"
88-
ystore = self._ystore_class(path=updates_file_path, log=self.log)
8996
self.room = DocumentRoom(
9097
self._room_id,
9198
file_format,
9299
file_type,
93100
file,
94101
self.event_logger,
95-
ystore,
102+
self._store,
96103
self.log,
97104
self._document_save_delay,
98105
)
@@ -111,15 +118,15 @@ def initialize(
111118
self,
112119
ywebsocket_server: JupyterWebsocketServer,
113120
file_loaders: FileLoaderMapping,
114-
ystore_class: type[BaseYStore],
121+
store: BaseYStore,
115122
document_cleanup_delay: float | None = 60.0,
116123
document_save_delay: float | None = 1.0,
117124
) -> None:
118125
self._background_tasks = set()
119126
# File ID manager cannot be passed as argument as the extension may load after this one
120127
self._file_id_manager = self.settings["file_id_manager"]
121128
self._file_loaders = file_loaders
122-
self._ystore_class = ystore_class
129+
self._store = store
123130
self._cleanup_delay = document_cleanup_delay
124131
self._document_save_delay = document_save_delay
125132
self._websocket_server = ywebsocket_server

‎jupyter_collaboration/rooms.py

Lines changed: 25 additions & 41 deletions
Original file line numberDiff line numberDiff line change
@@ -11,7 +11,7 @@
1111
from jupyter_events import EventLogger
1212
from jupyter_ydoc import ydocs as YDOCS
1313
from ypy_websocket.websocket_server import YRoom
14-
from ypy_websocket.ystore import BaseYStore, YDocNotFound
14+
from ypy_websocket.stores import BaseYStore, YDocNotFound
1515
from ypy_websocket.yutils import write_var_uint
1616

1717
from .loaders import FileLoader
@@ -104,56 +104,40 @@ async def initialize(self) -> None:
104104
return
105105

106106
self.log.info("Initializing room %s", self._room_id)
107-
108107
model = await self._file.load_content(self._file_format, self._file_type, True)
109108

110109
async with self._update_lock:
111110
# try to apply Y updates from the YStore for this document
112-
read_from_source = True
113-
if self.ystore is not None:
114-
try:
115-
await self.ystore.apply_updates(self.ydoc)
116-
self._emit(
117-
LogLevel.INFO,
118-
"load",
119-
"Content loaded from the store {}".format(
120-
self.ystore.__class__.__qualname__
121-
),
122-
)
123-
self.log.info(
124-
"Content in room %s loaded from the ystore %s",
125-
self._room_id,
126-
self.ystore.__class__.__name__,
127-
)
128-
read_from_source = False
129-
except YDocNotFound:
130-
# YDoc not found in the YStore, create the document from the source file (no change history)
131-
pass
132-
133-
if not read_from_source:
111+
if self.ystore is not None and await self.ystore.exists(self._room_id):
112+
# Load the content from the store
113+
await self.ystore.apply_updates(self._room_id, self.ydoc)
114+
self._emit(LogLevel.INFO, "load", "Content loaded from the store {}".format(self.ystore.__class__.__qualname__))
115+
self.log.info("Content in room %s loaded from the ystore %s", self._room_id, self.ystore.__class__.__name__,)
116+
134117
# if YStore updates and source file are out-of-sync, resync updates with source
135118
if self._document.source != model["content"]:
136-
# TODO: Delete document from the store.
137-
self._emit(
138-
LogLevel.INFO, "initialize", "The file is out-of-sync with the ystore."
139-
)
140-
self.log.info(
141-
"Content in file %s is out-of-sync with the ystore %s",
142-
self._file.path,
143-
self.ystore.__class__.__name__,
144-
)
145-
read_from_source = True
146-
147-
if read_from_source:
119+
self._emit(LogLevel.INFO, "initialize", "The file is out-of-sync with the ystore.")
120+
self.log.info("Content in file %s is out-of-sync with the ystore %s", self._file.path, self.ystore.__class__.__name__,)
121+
122+
doc = await self.ystore.get(self._room_id)
123+
await self.ystore.remove(self._room_id)
124+
version = 0
125+
if "version" in doc:
126+
version = doc["version"] + 1
127+
128+
await self.ystore.create(self._room_id, version)
129+
await self.ystore.encode_state_as_update(self._room_id, self.ydoc)
130+
131+
else:
148132
self._emit(LogLevel.INFO, "load", "Content loaded from disk.")
149-
self.log.info(
150-
"Content in room %s loaded from file %s", self._room_id, self._file.path
151-
)
133+
self.log.info("Content in room %s loaded from file %s", self._room_id, self._file.path)
152134
self._document.source = model["content"]
153135

154-
if self.ystore:
155-
await self.ystore.encode_state_as_update(self.ydoc)
136+
if self.ystore is not None:
137+
await self.ystore.create(self._room_id, 0)
138+
await self.ystore.encode_state_as_update(self._room_id, self.ydoc)
156139

140+
157141
self._last_modified = model["last_modified"]
158142
self._document.dirty = False
159143
self.ready = True

‎jupyter_collaboration/stores.py

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -3,8 +3,8 @@
33

44
from traitlets import Int, Unicode
55
from traitlets.config import LoggingConfigurable
6-
from ypy_websocket.ystore import SQLiteYStore as _SQLiteYStore
7-
from ypy_websocket.ystore import TempFileYStore as _TempFileYStore
6+
from ypy_websocket.stores import SQLiteYStore as _SQLiteYStore
7+
from ypy_websocket.stores import TempFileYStore as _TempFileYStore
88

99

1010
class TempFileYStore(_TempFileYStore):
@@ -17,7 +17,7 @@ class SQLiteYStoreMetaclass(type(LoggingConfigurable), type(_SQLiteYStore)): #
1717

1818
class SQLiteYStore(LoggingConfigurable, _SQLiteYStore, metaclass=SQLiteYStoreMetaclass):
1919
db_path = Unicode(
20-
".jupyter_ystore.db",
20+
"./jupyter_ystore.db",
2121
config=True,
2222
help="""The path to the YStore database. Defaults to '.jupyter_ystore.db' in the current
2323
directory.""",

‎jupyter_collaboration/websocketserver.py

Lines changed: 1 addition & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -9,7 +9,7 @@
99

1010
from tornado.websocket import WebSocketHandler
1111
from ypy_websocket.websocket_server import WebsocketServer, YRoom
12-
from ypy_websocket.ystore import BaseYStore
12+
from ypy_websocket.stores import BaseYStore
1313

1414

1515
class RoomNotFound(LookupError):
@@ -27,13 +27,11 @@ class JupyterWebsocketServer(WebsocketServer):
2727

2828
def __init__(
2929
self,
30-
ystore_class: BaseYStore,
3130
rooms_ready: bool = True,
3231
auto_clean_rooms: bool = True,
3332
log: Logger | None = None,
3433
):
3534
super().__init__(rooms_ready, auto_clean_rooms, log)
36-
self.ystore_class = ystore_class
3735
self.ypatch_nb = 0
3836
self.connected_users: dict[Any, Any] = {}
3937
# Async loop is not yet ready at the object instantiation

0 commit comments

Comments
 (0)