Skip to content

Commit 339a3f4

Browse files
dlqqqdlqqq
committed
fix #111, implement update_session()
1 parent 817ea86 commit 339a3f4

File tree

1 file changed

+118
-38
lines changed

1 file changed

+118
-38
lines changed

‎jupyter_server_documents/session_manager.py‎

Lines changed: 118 additions & 38 deletions
Original file line numberDiff line numberDiff line change
@@ -30,20 +30,47 @@ def yroom_manager(self) -> YRoomManager:
3030
"""The Jupyter Server's YRoom Manager."""
3131
return self.serverapp.web_app.settings["yroom_manager"]
3232

33-
def get_kernel_client(self, kernel_id) -> DocumentAwareKernelClient:
33+
_room_ids: dict[str, str]
34+
"""
35+
Dictionary of room IDs, keyed by session ID.
36+
"""
37+
38+
def __init__(self, *args, **kwargs):
39+
super().__init__(*args, **kwargs)
40+
self._room_ids = {}
41+
42+
def get_kernel_client(self, kernel_id: str) -> DocumentAwareKernelClient:
3443
"""Get the kernel client for a running kernel."""
3544
kernel_manager = self.kernel_manager.get_kernel(kernel_id)
3645
kernel_client = kernel_manager.main_client
3746
return kernel_client
3847

39-
# The `type` argument here is for future proofing this API.
40-
# Today, only notebooks have ydocs with kernels connected.
41-
# In the future, we will include consoles here.
42-
def get_yroom(self, path, type) -> YRoom:
43-
"""Get the yroom for a given path."""
48+
def get_yroom(self, session_id: str) -> YRoom:
49+
"""
50+
Get the `YRoom` for a session given its ID. The session must have
51+
been created first via a call to `create_session()`.
52+
"""
53+
room_id = self._room_ids.get(session_id, None)
54+
yroom = self.yroom_manager.get_room(room_id) if room_id else None
55+
if not yroom:
56+
raise LookupError(f"No room found for session '{session_id}'.")
57+
return yroom
58+
59+
60+
def _init_session_yroom(self, session_id: str, path: str) -> YRoom:
61+
"""
62+
Returns a `YRoom` for a session identified by the given `session_id` and
63+
`path`. This should be called only in `create_session()`.
64+
65+
This method stores the new room ID & session ID in `self._room_ids`. The
66+
`YRoom` for a session can be retrieved via `self.get_yroom()` after this
67+
method is called.
68+
"""
4469
file_id = self.file_id_manager.index(path)
4570
room_id = f"json:notebook:{file_id}"
4671
yroom = self.yroom_manager.get_room(room_id)
72+
self._room_ids[session_id] = room_id
73+
4774
return yroom
4875

4976
async def create_session(
@@ -57,51 +84,104 @@ async def create_session(
5784
"""
5885
After creating a session, connects the yroom to the kernel client.
5986
"""
60-
output = await super().create_session(
87+
session_model = await super().create_session(
6188
path,
6289
name,
6390
type,
6491
kernel_name,
6592
kernel_id
6693
)
94+
session_id = session_model["id"]
6795
if kernel_id is None:
68-
kernel_id = output["kernel"]["id"]
69-
70-
# Connect this session's yroom to the kernel.
71-
if type == "notebook":
72-
# If name or path is None, we cannot map to a yroom,
73-
# so just move on.
74-
if name is None or path is None:
75-
self.log.debug("`name` or `path` was not given, so a yroom was not set up for this session.")
76-
return output
77-
# When JupyterLab creates a session, it uses a fake path
78-
# which is the relative path + UUID, i.e. the notebook
79-
# name is incorrect temporarily. It later makes multiple
80-
# updates to the session to correct the path.
81-
#
82-
# Here, we create the true path to store in the fileID service
83-
# by dropping the UUID and appending the file name.
84-
real_path = os.path.join(os.path.split(path)[0], name)
85-
yroom = self.get_yroom(real_path, type)
86-
# TODO: we likely have a race condition here... need to
87-
# think about it more. Currently, the kernel client gets
88-
# created after the kernel starts fully. We need the
89-
# kernel client instantiated _before_ trying to connect
90-
# the yroom.
91-
kernel_client = self.get_kernel_client(kernel_id)
92-
await kernel_client.add_yroom(yroom)
93-
self.log.info(f"Connected yroom {yroom.room_id} to kernel {kernel_id}. yroom: {yroom}")
94-
else:
95-
self.log.debug(f"Document type {type} is not supported by YRoom.")
96-
return output
96+
kernel_id = session_model["kernel"]["id"]
97+
98+
# If the type is not 'notebook', return the session model immediately
99+
if type != "notebook":
100+
self.log.warning(
101+
f"Document type '{type}' is not recognized by YDocSessionManager."
102+
)
103+
return session_model
104+
105+
# If name or path is None, we cannot map to a yroom,
106+
# so just move on.
107+
if name is None or path is None:
108+
self.log.warning(f"`name` or `path` was not given for new session at '{path}'.")
109+
return session_model
110+
111+
# Otherwise, get a `YRoom` and add it to this session's kernel client.
112+
113+
# When JupyterLab creates a session, it uses a fake path
114+
# which is the relative path + UUID, i.e. the notebook
115+
# name is incorrect temporarily. It later makes multiple
116+
# updates to the session to correct the path.
117+
#
118+
# Here, we create the true path to store in the fileID service
119+
# by dropping the UUID and appending the file name.
120+
real_path = os.path.join(os.path.split(path)[0], name)
121+
122+
# Get YRoom for this session and store its ID in `self._room_ids`
123+
yroom = self._init_session_yroom(session_id, real_path)
124+
125+
# Add YRoom to this session's kernel client
126+
# TODO: we likely have a race condition here... need to
127+
# think about it more. Currently, the kernel client gets
128+
# created after the kernel starts fully. We need the
129+
# kernel client instantiated _before_ trying to connect
130+
# the yroom.
131+
kernel_client = self.get_kernel_client(kernel_id)
132+
await kernel_client.add_yroom(yroom)
133+
self.log.info(f"Connected yroom {yroom.room_id} to kernel {kernel_id}. yroom: {yroom}")
134+
return session_model
135+
136+
137+
async def update_session(self, session_id, **update) -> None:
138+
"""
139+
Updates the session identified by `session_id` using the keyword
140+
arguments passed to this method. Each keyword argument should correspond
141+
to a column in the sessions table.
142+
143+
This class calls the `update_session()` parent method, then updates the
144+
kernel client if `update` contains `kernel_id`.
145+
"""
146+
# Apply update and return early if `kernel_id` was not updated
147+
if "kernel_id" not in update:
148+
return await super().update_session(session_id, **update)
149+
150+
# Otherwise, first remove the YRoom from the old kernel client and add
151+
# the YRoom to the new kernel client, before applying the update.
152+
old_session_info = (await self.get_session(session_id=session_id) or {})
153+
old_kernel_id = old_session_info.get("kernel_id", None)
154+
new_kernel_id = update.get("kernel_id", None)
155+
self.log.info(
156+
f"Updating session '{session_id}' from kernel '{old_kernel_id}' "
157+
f"to kernel '{new_kernel_id}'."
158+
)
159+
yroom = self.get_yroom(session_id)
160+
if old_kernel_id:
161+
old_kernel_client = self.get_kernel_client(old_kernel_id)
162+
await old_kernel_client.remove_yroom(yroom=yroom)
163+
if new_kernel_id:
164+
new_kernel_client = self.get_kernel_client(new_kernel_id)
165+
await new_kernel_client.add_yroom(yroom=yroom)
166+
167+
# Apply update and return
168+
return await super().update_session(session_id, **update)
169+
97170

98171
async def delete_session(self, session_id):
99172
"""
100173
Deletes the session and disconnects the yroom from the kernel client.
101174
"""
102175
session = await self.get_session(session_id=session_id)
103-
kernel_id, path, type = session["kernel"]["id"], session["path"], session["type"]
104-
yroom = self.get_yroom(path, type)
176+
kernel_id = session["kernel"]["id"]
177+
178+
# Remove YRoom from session's kernel client
179+
yroom = self.get_yroom(session_id)
105180
kernel_client = self.get_kernel_client(kernel_id)
106181
await kernel_client.remove_yroom(yroom)
182+
183+
# Remove room ID stored for the session
184+
self._room_ids.pop(session_id, None)
185+
186+
# Delete the session via the parent method
107187
await super().delete_session(session_id)

0 commit comments

Comments
 (0)