mirror of https://github.com/dbcli/pgcli
153 lines
5.1 KiB
Python
153 lines
5.1 KiB
Python
import threading
|
|
import os
|
|
from collections import OrderedDict
|
|
|
|
from .pgcompleter import PGCompleter
|
|
|
|
|
|
class CompletionRefresher:
|
|
refreshers = OrderedDict()
|
|
|
|
def __init__(self):
|
|
self._completer_thread = None
|
|
self._restart_refresh = threading.Event()
|
|
|
|
def refresh(self, executor, special, callbacks, history=None, settings=None):
|
|
"""
|
|
Creates a PGCompleter object and populates it with the relevant
|
|
completion suggestions in a background thread.
|
|
|
|
executor - PGExecute object, used to extract the credentials to connect
|
|
to the database.
|
|
special - PGSpecial object used for creating a new completion object.
|
|
settings - dict of settings for completer object
|
|
callbacks - A function or a list of functions to call after the thread
|
|
has completed the refresh. The newly created completion
|
|
object will be passed in as an argument to each callback.
|
|
"""
|
|
if executor.is_virtual_database():
|
|
# do nothing
|
|
return [(None, None, None, "Auto-completion refresh can't be started.")]
|
|
|
|
if self.is_refreshing():
|
|
self._restart_refresh.set()
|
|
return [(None, None, None, "Auto-completion refresh restarted.")]
|
|
else:
|
|
self._completer_thread = threading.Thread(
|
|
target=self._bg_refresh,
|
|
args=(executor, special, callbacks, history, settings),
|
|
name="completion_refresh",
|
|
)
|
|
self._completer_thread.setDaemon(True)
|
|
self._completer_thread.start()
|
|
return [
|
|
(None, None, None, "Auto-completion refresh started in the background.")
|
|
]
|
|
|
|
def is_refreshing(self):
|
|
return self._completer_thread and self._completer_thread.is_alive()
|
|
|
|
def _bg_refresh(self, pgexecute, special, callbacks, history=None, settings=None):
|
|
settings = settings or {}
|
|
completer = PGCompleter(
|
|
smart_completion=True, pgspecial=special, settings=settings
|
|
)
|
|
|
|
if settings.get("single_connection"):
|
|
executor = pgexecute
|
|
else:
|
|
# Create a new pgexecute method to populate the completions.
|
|
executor = pgexecute.copy()
|
|
# If callbacks is a single function then push it into a list.
|
|
if callable(callbacks):
|
|
callbacks = [callbacks]
|
|
|
|
while 1:
|
|
for refresher in self.refreshers.values():
|
|
refresher(completer, executor)
|
|
if self._restart_refresh.is_set():
|
|
self._restart_refresh.clear()
|
|
break
|
|
else:
|
|
# Break out of while loop if the for loop finishes natually
|
|
# without hitting the break statement.
|
|
break
|
|
|
|
# Start over the refresh from the beginning if the for loop hit the
|
|
# break statement.
|
|
continue
|
|
|
|
# Load history into pgcompleter so it can learn user preferences
|
|
n_recent = 100
|
|
if history:
|
|
for recent in history.get_strings()[-n_recent:]:
|
|
completer.extend_query_history(recent, is_init=True)
|
|
|
|
for callback in callbacks:
|
|
callback(completer)
|
|
|
|
if not settings.get("single_connection") and executor.conn:
|
|
# close connection established with pgexecute.copy()
|
|
executor.conn.close()
|
|
|
|
|
|
def refresher(name, refreshers=CompletionRefresher.refreshers):
|
|
"""Decorator to populate the dictionary of refreshers with the current
|
|
function.
|
|
"""
|
|
|
|
def wrapper(wrapped):
|
|
refreshers[name] = wrapped
|
|
return wrapped
|
|
|
|
return wrapper
|
|
|
|
|
|
@refresher("schemata")
|
|
def refresh_schemata(completer, executor):
|
|
completer.set_search_path(executor.search_path())
|
|
completer.extend_schemata(executor.schemata())
|
|
|
|
|
|
@refresher("tables")
|
|
def refresh_tables(completer, executor):
|
|
completer.extend_relations(executor.tables(), kind="tables")
|
|
completer.extend_columns(executor.table_columns(), kind="tables")
|
|
completer.extend_foreignkeys(executor.foreignkeys())
|
|
|
|
|
|
@refresher("views")
|
|
def refresh_views(completer, executor):
|
|
completer.extend_relations(executor.views(), kind="views")
|
|
completer.extend_columns(executor.view_columns(), kind="views")
|
|
|
|
|
|
@refresher("types")
|
|
def refresh_types(completer, executor):
|
|
completer.extend_datatypes(executor.datatypes())
|
|
|
|
|
|
@refresher("databases")
|
|
def refresh_databases(completer, executor):
|
|
completer.extend_database_names(executor.databases())
|
|
|
|
|
|
@refresher("casing")
|
|
def refresh_casing(completer, executor):
|
|
casing_file = completer.casing_file
|
|
if not casing_file:
|
|
return
|
|
generate_casing_file = completer.generate_casing_file
|
|
if generate_casing_file and not os.path.isfile(casing_file):
|
|
casing_prefs = "\n".join(executor.casing())
|
|
with open(casing_file, "w") as f:
|
|
f.write(casing_prefs)
|
|
if os.path.isfile(casing_file):
|
|
with open(casing_file) as f:
|
|
completer.extend_casing([line.strip() for line in f])
|
|
|
|
|
|
@refresher("functions")
|
|
def refresh_functions(completer, executor):
|
|
completer.extend_functions(executor.functions())
|