datasette_connectors/monkey.py

   1 import threading
   2 import sqlite3
   3 import datasette.views.base
   4 from datasette.database import Database
   5
   6 from .connectors import ConnectorList
   7
   8 connections = threading.local()
   9
  10
  11 def patch_datasette():
  12     """
  13     Monkey patching for original Datasette
  14     """
  15
  16     async def table_names(self):
  17         try:
  18             return await self.original_table_names()
  19         except sqlite3.DatabaseError:
  20             return ConnectorList.table_names(self.path)
  21
  22     Database.original_table_names = Database.table_names
  23     Database.table_names = table_names
  24
  25
  26     async def hidden_table_names(self):
  27         try:
  28             return await self.original_hidden_table_names()
  29         except sqlite3.DatabaseError:
  30             return ConnectorList.hidden_table_names(self.path)
  31
  32     Database.original_hidden_table_names = Database.hidden_table_names
  33     Database.hidden_table_names = hidden_table_names
  34
  35
  36     async def view_names(self):
  37         try:
  38             return await self.original_view_names()
  39         except sqlite3.DatabaseError:
  40             return ConnectorList.view_names(self.path)
  41
  42     Database.original_view_names = Database.view_names
  43     Database.view_names = view_names
  44
  45
  46     async def table_columns(self, table):
  47         try:
  48             return await self.original_table_columns(table)
  49         except sqlite3.DatabaseError:
  50             return ConnectorList.table_columns(self.path, table)
  51
  52     Database.original_table_columns = Database.table_columns
  53     Database.table_columns = table_columns
  54
  55
  56     async def primary_keys(self, table):
  57         try:
  58             return await self.original_primary_keys(table)
  59         except sqlite3.DatabaseError:
  60             return ConnectorList.primary_keys(self.path, table)
  61
  62     Database.original_primary_keys = Database.primary_keys
  63     Database.primary_keys = primary_keys
  64
  65
  66     async def fts_table(self, table):
  67         try:
  68             return await self.original_fts_table(table)
  69         except sqlite3.DatabaseError:
  70             return ConnectorList.fts_table(self.path, table)
  71
  72     Database.original_fts_table = Database.fts_table
  73     Database.fts_table = fts_table
  74
  75
  76     async def get_all_foreign_keys(self):
  77         try:
  78             return await self.original_get_all_foreign_keys()
  79         except sqlite3.DatabaseError:
  80             return ConnectorList.get_all_foreign_keys(self.path)
  81
  82     Database.original_get_all_foreign_keys = Database.get_all_foreign_keys
  83     Database.get_all_foreign_keys = get_all_foreign_keys
  84
  85
  86     async def table_counts(self, *args, **kwargs):
  87         counts = await self.original_table_counts(**kwargs)
  88         # If all tables has None as counts, an error had ocurred
  89         if len(list(filter(lambda table_count: table_count is not None, counts.values()))) == 0:
  90             return ConnectorList.table_counts(self.path, *args, **kwargs)
  91
  92     Database.original_table_counts = Database.table_counts
  93     Database.table_counts = table_counts