datasette_connectors/monkey.py

   1 import asyncio
   2 import threading
   3 import sqlite3
   4
   5 import datasette.views.base
   6 from datasette.tracer import trace
   7 from datasette.database import Database
   8 from datasette.database import Results
   9
  10 from .connectors import ConnectorList
  11
  12 connections = threading.local()
  13
  14
  15 def patch_datasette():
  16     """
  17     Monkey patching for original Datasette
  18     """
  19
  20     def connect(self, write=False):
  21         try:
  22             # Check if it's a sqlite database
  23             conn = self.original_connect(write=write)
  24             conn.execute("select name from sqlite_master where type='table'")
  25             return conn
  26         except sqlite3.DatabaseError:
  27             conn = ConnectorList.connect(self.path)
  28             return conn
  29
  30     Database.original_connect = Database.connect
  31     Database.connect = connect
  32
  33
  34     async def execute_fn(self, fn):
  35         def in_thread():
  36             conn = getattr(connections, self.name, None)
  37             if not conn:
  38                 conn = self.connect()
  39                 if isinstance(conn, sqlite3.Connection):
  40                     self.ds._prepare_connection(conn, self.name)
  41                 setattr(connections, self.name, conn)
  42             return fn(conn)
  43
  44         return await asyncio.get_event_loop().run_in_executor(
  45             self.ds.executor, in_thread
  46         )
  47
  48     Database.original_execute_fn = Database.execute_fn
  49     Database.execute_fn = execute_fn