Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Refactor db connect to use worker #399

Merged
merged 1 commit into from
Jan 10, 2024
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
5 changes: 5 additions & 0 deletions CHANGELOG.md
Original file line number Diff line number Diff line change
Expand Up @@ -4,6 +4,11 @@ All notable changes to this project will be documented in this file.

## [Unreleased]

### Features

- Harlequin now loads immediately and connects to your database in the background ([#393](https://github.com/tconbeer/harlequin/issues/393)).
- Harlequin shows a loading indicator before the Data Catalog is hydrated for the first time ([#396](https://github.com/tconbeer/harlequin/issues/396)).

## [1.9.2] - 2024-01-10

### Features
Expand Down
8 changes: 4 additions & 4 deletions poetry.lock

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

1 change: 1 addition & 0 deletions pyproject.toml
Original file line number Diff line number Diff line change
Expand Up @@ -50,6 +50,7 @@ harlequin-trino = { version = "^0.1", optional = true }
pre-commit = "^3.3.1"
textual-dev = "^1.0.1"
harlequin-postgres = "^0.2"
harlequin-mysql = "^0.1.1"

[tool.poetry.group.static.dependencies]
black = "^23.3.0"
Expand Down
87 changes: 71 additions & 16 deletions src/harlequin/app.py
Original file line number Diff line number Diff line change
Expand Up @@ -23,6 +23,7 @@
from textual_fastdatatable import DataTable
from textual_fastdatatable.backend import AutoBackendType

from harlequin import HarlequinConnection
from harlequin.adapter import HarlequinAdapter, HarlequinCursor
from harlequin.autocomplete import completer_factory
from harlequin.cache import BufferState, Cache, write_cache
Expand All @@ -42,12 +43,20 @@
from harlequin.exception import (
HarlequinConfigError,
HarlequinConnectionError,
HarlequinError,
HarlequinQueryError,
HarlequinThemeError,
pretty_error_message,
pretty_print_error,
)


class DatabaseConnected(Message):
def __init__(self, connection: HarlequinConnection) -> None:
super().__init__()
self.connection = connection


class QuerySubmitted(Message):
def __init__(self, query_text: str, limit: int | None) -> None:
super().__init__()
Expand Down Expand Up @@ -126,14 +135,7 @@ def __init__(
)
self.exit(return_code=2)
self.query_timer: Union[float, None] = None
try:
self.connection = self.adapter.connect()
except HarlequinConnectionError as e:
pretty_print_error(e)
self.exit(return_code=2)
else:
if self.connection.init_message:
self.notify(self.connection.init_message)
self.connection: HarlequinConnection | None = None

try:
self.app_colors = HarlequinColors.from_theme(theme)
Expand All @@ -150,9 +152,12 @@ def compose(self) -> ComposeResult:
yield DataCatalog(type_color=self.app_colors.gray)
with Vertical(id="main_panel"):
yield EditorCollection(language="sql", theme=self.theme)
yield RunQueryBar(max_results=self.max_results)
yield RunQueryBar(
max_results=self.max_results, classes="non-responsive"
)
yield ResultsViewer(
max_results=self.max_results, type_color=self.app_colors.gray
max_results=self.max_results,
type_color=self.app_colors.gray,
)
yield Footer()

Expand Down Expand Up @@ -186,8 +191,9 @@ async def on_mount(self) -> None:

self.editor.focus()
self.run_query_bar.checkbox.value = False
self.data_catalog.loading = True

self.update_schema_data()
self._connect()

@on(Button.Pressed, "#run_query")
def submit_query_from_run_query_bar(self, message: Button.Pressed) -> None:
Expand All @@ -208,6 +214,17 @@ def submit_query_from_editor(self, message: CodeEditor.Submitted) -> None:
)
)

@on(DatabaseConnected)
def initialize_app(self, message: DatabaseConnected) -> None:
self.connection = message.connection
self.run_query_bar.set_responsive()
self.results_viewer.show_table(did_run=False)
if message.connection.init_message:
self.notify(message.connection.init_message, title="Database Connected.")
else:
self.notify("Database Connected.")
self.update_schema_data()

@on(DataCatalog.NodeSubmitted)
def insert_node_into_editor(
self, message: DataCatalog.NodeSubmitted[CatalogItem]
Expand Down Expand Up @@ -277,11 +294,11 @@ def copy_data_to_clipboard(self, message: DataTable.SelectionCopied) -> None:
self.notify("Selected data copied to clipboard.")

@on(Worker.StateChanged)
def handle_worker_error(self, message: Worker.StateChanged) -> None:
async def handle_worker_error(self, message: Worker.StateChanged) -> None:
if message.state == WorkerState.ERROR:
self._handle_worker_error(message)
await self._handle_worker_error(message)

def _handle_worker_error(self, message: Worker.StateChanged) -> None:
async def _handle_worker_error(self, message: Worker.StateChanged) -> None:
if (
message.worker.name == "update_schema_data"
and message.worker.error is not None
Expand All @@ -296,16 +313,33 @@ def _handle_worker_error(self, message: Worker.StateChanged) -> None:
):
self.run_query_bar.set_responsive()
self.results_viewer.show_table()
header = getattr(message.worker.error, "title", "Query Error")
header = getattr(
message.worker.error, "title", message.worker.error.__class__.__name__
)
self._push_error_modal(
title="Query Error",
header=header,
error=message.worker.error,
)
elif message.worker.name == "_connect" and message.worker.error is not None:
title = getattr(
message.worker.error,
"title",
"Harlequin could not connect to your database.",
)
error = (
message.worker.error
if isinstance(message.worker.error, HarlequinError)
else HarlequinConnectionError(
msg=str(message.worker.error), title=title
)
)
self.exit(return_code=2, message=pretty_error_message(error))

@on(NewCatalog)
def update_tree_and_completers(self, message: NewCatalog) -> None:
self.data_catalog.update_tree(message.catalog)
self.data_catalog.loading = False
self.update_completers(message.catalog)

@on(QueriesExecuted)
Expand Down Expand Up @@ -383,6 +417,8 @@ def watch_full_screen(self, full_screen: bool) -> None:

@on(QuerySubmitted)
def execute_query(self, message: QuerySubmitted) -> None:
if self.connection is None:
return
if message.query_text:
self.full_screen = False
self.run_query_bar.set_not_responsive()
Expand Down Expand Up @@ -472,6 +508,17 @@ def action_toggle_sidebar(self) -> None:
else:
self.sidebar_hidden = not self.sidebar_hidden

@work(
thread=True,
exclusive=True,
exit_on_error=False,
group="connect",
description="Connecting to DB",
)
def _connect(self) -> None:
connection = self.adapter.connect()
self.post_message(DatabaseConnected(connection=connection))

@work(
thread=True,
exclusive=True,
Expand All @@ -480,6 +527,8 @@ def action_toggle_sidebar(self) -> None:
description="Executing queries.",
)
def _execute_query(self, message: QuerySubmitted) -> None:
if self.connection is None:
return
cursors: Dict[str, HarlequinCursor] = {}
queries = self._split_query_text(message.query_text)
for q in queries:
Expand Down Expand Up @@ -514,7 +563,7 @@ def _push_error_modal(self, title: str, header: str, error: BaseException) -> No
title=title,
header=header,
error=error,
)
),
)

@work(
Expand Down Expand Up @@ -543,6 +592,8 @@ def _fetch_data(

@work(thread=True, exclusive=True, exit_on_error=True, group="completer_builders")
def update_completers(self, catalog: Catalog) -> None:
if self.connection is None:
return
if (
self.editor_collection.word_completer is not None
and self.editor_collection.member_completer is not None
Expand All @@ -561,6 +612,8 @@ def update_completers(self, catalog: Catalog) -> None:

@work(thread=True, exclusive=True, exit_on_error=False, group="schema_updaters")
def update_schema_data(self) -> None:
if self.connection is None:
return
catalog = self.connection.get_catalog()
self.post_message(NewCatalog(catalog=catalog))

Expand All @@ -570,6 +623,8 @@ def _validate_selection(self) -> str:
return the empty string.
"""
selection = self.editor.selected_text
if self.connection is None:
return selection
if selection:
try:
return self.connection.validate_sql(selection)
Expand Down
22 changes: 13 additions & 9 deletions src/harlequin/exception.py
Original file line number Diff line number Diff line change
@@ -1,3 +1,6 @@
from rich.panel import Panel


class HarlequinExit(Exception):
pass

Expand Down Expand Up @@ -39,13 +42,14 @@ class HarlequinTzDataError(HarlequinError):

def pretty_print_error(error: HarlequinError) -> None:
from rich import print
from rich.panel import Panel

print(
Panel.fit(
str(error),
title=error.title if error.title else ("Harlequin encountered an error."),
title_align="left",
border_style="red",
)

print(pretty_error_message(error))


def pretty_error_message(error: HarlequinError) -> Panel:
return Panel.fit(
str(error),
title=error.title if error.title else ("Harlequin encountered an error."),
title_align="left",
border_style="red",
)
4 changes: 4 additions & 0 deletions src/harlequin/global.tcss
Original file line number Diff line number Diff line change
Expand Up @@ -77,6 +77,10 @@ Toast.-error {
border-left: wide $error;
}

Toast .toast--title {
color: $primary;
}

/* DATA CATALOG */
DataCatalog {
height: 1fr;
Expand Down