sessionmanager.py 8.5 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242
  1. """A base class session manager."""
  2. # Copyright (c) Jupyter Development Team.
  3. # Distributed under the terms of the Modified BSD License.
  4. import os
  5. import uuid
  6. try:
  7. import sqlite3
  8. except ImportError:
  9. # fallback on pysqlite2 if Python was build without sqlite
  10. from pysqlite2 import dbapi2 as sqlite3
  11. from tornado import gen, web
  12. from traitlets.config.configurable import LoggingConfigurable
  13. from ipython_genutils.py3compat import unicode_type
  14. from traitlets import Instance
  15. class SessionManager(LoggingConfigurable):
  16. kernel_manager = Instance('notebook.services.kernels.kernelmanager.MappingKernelManager')
  17. contents_manager = Instance('notebook.services.contents.manager.ContentsManager')
  18. # Session database initialized below
  19. _cursor = None
  20. _connection = None
  21. _columns = {'session_id', 'path', 'name', 'type', 'kernel_id'}
  22. @property
  23. def cursor(self):
  24. """Start a cursor and create a database called 'session'"""
  25. if self._cursor is None:
  26. self._cursor = self.connection.cursor()
  27. self._cursor.execute("""CREATE TABLE session
  28. (session_id, path, name, type, kernel_id)""")
  29. return self._cursor
  30. @property
  31. def connection(self):
  32. """Start a database connection"""
  33. if self._connection is None:
  34. self._connection = sqlite3.connect(':memory:')
  35. self._connection.row_factory = sqlite3.Row
  36. return self._connection
  37. def close(self):
  38. """Close the sqlite connection"""
  39. if self._cursor is not None:
  40. self._cursor.close()
  41. self._cursor = None
  42. def __del__(self):
  43. """Close connection once SessionManager closes"""
  44. self.close()
  45. def session_exists(self, path):
  46. """Check to see if the session of a given name exists"""
  47. self.cursor.execute("SELECT * FROM session WHERE path=?", (path,))
  48. reply = self.cursor.fetchone()
  49. if reply is None:
  50. return False
  51. else:
  52. return True
  53. def new_session_id(self):
  54. "Create a uuid for a new session"
  55. return unicode_type(uuid.uuid4())
  56. @gen.coroutine
  57. def create_session(self, path=None, name=None, type=None, kernel_name=None, kernel_id=None):
  58. """Creates a session and returns its model"""
  59. session_id = self.new_session_id()
  60. if kernel_id is not None and kernel_id in self.kernel_manager:
  61. pass
  62. else:
  63. kernel_id = yield self.start_kernel_for_session(session_id, path, name, type, kernel_name)
  64. result = yield gen.maybe_future(
  65. self.save_session(session_id, path=path, name=name, type=type, kernel_id=kernel_id)
  66. )
  67. # py2-compat
  68. raise gen.Return(result)
  69. @gen.coroutine
  70. def start_kernel_for_session(self, session_id, path, name, type, kernel_name):
  71. """Start a new kernel for a given session."""
  72. # allow contents manager to specify kernels cwd
  73. kernel_path = self.contents_manager.get_kernel_path(path=path)
  74. kernel_id = yield gen.maybe_future(
  75. self.kernel_manager.start_kernel(path=kernel_path, kernel_name=kernel_name)
  76. )
  77. # py2-compat
  78. raise gen.Return(kernel_id)
  79. def save_session(self, session_id, path=None, name=None, type=None, kernel_id=None):
  80. """Saves the items for the session with the given session_id
  81. Given a session_id (and any other of the arguments), this method
  82. creates a row in the sqlite session database that holds the information
  83. for a session.
  84. Parameters
  85. ----------
  86. session_id : str
  87. uuid for the session; this method must be given a session_id
  88. path : str
  89. the path for the given session
  90. name: str
  91. the name of the session
  92. type: string
  93. the type of the session
  94. kernel_id : str
  95. a uuid for the kernel associated with this session
  96. Returns
  97. -------
  98. model : dict
  99. a dictionary of the session model
  100. """
  101. self.cursor.execute("INSERT INTO session VALUES (?,?,?,?,?)",
  102. (session_id, path, name, type, kernel_id)
  103. )
  104. return self.get_session(session_id=session_id)
  105. def get_session(self, **kwargs):
  106. """Returns the model for a particular session.
  107. Takes a keyword argument and searches for the value in the session
  108. database, then returns the rest of the session's info.
  109. Parameters
  110. ----------
  111. **kwargs : keyword argument
  112. must be given one of the keywords and values from the session database
  113. (i.e. session_id, path, name, type, kernel_id)
  114. Returns
  115. -------
  116. model : dict
  117. returns a dictionary that includes all the information from the
  118. session described by the kwarg.
  119. """
  120. if not kwargs:
  121. raise TypeError("must specify a column to query")
  122. conditions = []
  123. for column in kwargs.keys():
  124. if column not in self._columns:
  125. raise TypeError("No such column: %r", column)
  126. conditions.append("%s=?" % column)
  127. query = "SELECT * FROM session WHERE %s" % (' AND '.join(conditions))
  128. self.cursor.execute(query, list(kwargs.values()))
  129. try:
  130. row = self.cursor.fetchone()
  131. except KeyError:
  132. # The kernel is missing, so the session just got deleted.
  133. row = None
  134. if row is None:
  135. q = []
  136. for key, value in kwargs.items():
  137. q.append("%s=%r" % (key, value))
  138. raise web.HTTPError(404, u'Session not found: %s' % (', '.join(q)))
  139. return self.row_to_model(row)
  140. def update_session(self, session_id, **kwargs):
  141. """Updates the values in the session database.
  142. Changes the values of the session with the given session_id
  143. with the values from the keyword arguments.
  144. Parameters
  145. ----------
  146. session_id : str
  147. a uuid that identifies a session in the sqlite3 database
  148. **kwargs : str
  149. the key must correspond to a column title in session database,
  150. and the value replaces the current value in the session
  151. with session_id.
  152. """
  153. self.get_session(session_id=session_id)
  154. if not kwargs:
  155. # no changes
  156. return
  157. sets = []
  158. for column in kwargs.keys():
  159. if column not in self._columns:
  160. raise TypeError("No such column: %r" % column)
  161. sets.append("%s=?" % column)
  162. query = "UPDATE session SET %s WHERE session_id=?" % (', '.join(sets))
  163. self.cursor.execute(query, list(kwargs.values()) + [session_id])
  164. def row_to_model(self, row):
  165. """Takes sqlite database session row and turns it into a dictionary"""
  166. if row['kernel_id'] not in self.kernel_manager:
  167. # The kernel was killed or died without deleting the session.
  168. # We can't use delete_session here because that tries to find
  169. # and shut down the kernel.
  170. self.cursor.execute("DELETE FROM session WHERE session_id=?",
  171. (row['session_id'],))
  172. raise KeyError
  173. model = {
  174. 'id': row['session_id'],
  175. 'path': row['path'],
  176. 'name': row['name'],
  177. 'type': row['type'],
  178. 'kernel': self.kernel_manager.kernel_model(row['kernel_id'])
  179. }
  180. if row['type'] == 'notebook':
  181. # Provide the deprecated API.
  182. model['notebook'] = {'path': row['path'], 'name': row['name']}
  183. return model
  184. def list_sessions(self):
  185. """Returns a list of dictionaries containing all the information from
  186. the session database"""
  187. c = self.cursor.execute("SELECT * FROM session")
  188. result = []
  189. # We need to use fetchall() here, because row_to_model can delete rows,
  190. # which messes up the cursor if we're iterating over rows.
  191. for row in c.fetchall():
  192. try:
  193. result.append(self.row_to_model(row))
  194. except KeyError:
  195. pass
  196. return result
  197. @gen.coroutine
  198. def delete_session(self, session_id):
  199. """Deletes the row in the session database with given session_id"""
  200. session = self.get_session(session_id=session_id)
  201. yield gen.maybe_future(self.kernel_manager.shutdown_kernel(session['kernel']['id']))
  202. self.cursor.execute("DELETE FROM session WHERE session_id=?", (session_id,))