datasette_connectors/monkey.py

   1 import asyncio
   2 import datasette
   3 from datasette.app import connections
   4 from datasette.inspect import inspect_hash
   5 from datasette.utils import Results
   6 from pathlib import Path
   7 import sqlite3
   8
   9 from . import connectors
  10
  11
  12 def patch_datasette():
  13     """
  14     Monkey patching for original Datasette
  15     """
  16
  17     def inspect(self):
  18         " Inspect the database and return a dictionary of table metadata "
  19         if self._inspect:
  20             return self._inspect
  21
  22         _inspect = {}
  23         files = self.files
  24
  25         for filename in files:
  26             self.files = (filename,)
  27             path = Path(filename)
  28             name = path.stem
  29             if name in _inspect:
  30                 raise Exception("Multiple files with the same stem %s" % name)
  31             try:
  32                 _inspect[name] = self.original_inspect()[name]
  33             except sqlite3.DatabaseError:
  34                 tables, views, dbtype = connectors.inspect(path)
  35                 _inspect[name] = {
  36                     "hash": inspect_hash(path),
  37                     "file": str(path),
  38                     "dbtype": dbtype,
  39                     "tables": tables,
  40                     "views": views,
  41                 }
  42
  43         self.files = files
  44         self._inspect = _inspect
  45         return self._inspect
  46
  47     datasette.app.Datasette.original_inspect = datasette.app.Datasette.inspect
  48     datasette.app.Datasette.inspect = inspect
  49
  50
  51     async def execute(self, db_name, sql, params=None, truncate=False, custom_time_limit=None, page_size=None):
  52         """Executes sql against db_name in a thread"""
  53         page_size = page_size or self.page_size
  54
  55         def is_sqlite3_conn():
  56             conn = getattr(connections, db_name, None)
  57             if not conn:
  58                 info = self.inspect()[db_name]
  59                 return info.get('dbtype', 'sqlite3') == 'sqlite3'
  60             else:
  61                 return isinstance(conn, sqlite3.Connection)
  62
  63         def sql_operation_in_thread():
  64             conn = getattr(connections, db_name, None)
  65             if not conn:
  66                 info = self.inspect()[db_name]
  67                 conn = connectors.connect(info['file'], info['dbtype'])
  68                 setattr(connections, db_name, conn)
  69
  70             rows, truncated, description = conn.execute(
  71                 sql,
  72                 params or {},
  73                 truncate=truncate,
  74                 page_size=page_size,
  75                 max_returned_rows=self.max_returned_rows,
  76             )
  77             return Results(rows, truncated, description)
  78
  79         if is_sqlite3_conn():
  80             return await self.original_execute(db_name, sql, params=params, truncate=truncate, custom_time_limit=custom_time_limit, page_size=page_size)
  81         else:
  82             return await asyncio.get_event_loop().run_in_executor(
  83                 self.executor, sql_operation_in_thread
  84             )
  85
  86     datasette.app.Datasette.original_execute = datasette.app.Datasette.execute
  87     datasette.app.Datasette.execute = execute