diff --git a/.github/workflows/static.yml b/.github/workflows/static.yml index efd70221..19a7d437 100644 --- a/.github/workflows/static.yml +++ b/.github/workflows/static.yml @@ -10,18 +10,18 @@ concurrency: jobs: static: - name: Static Analysis - 3.11 + name: Static Analysis - 3.8 runs-on: ubuntu-latest steps: - name: Check out Repo uses: actions/checkout@v4 with: persist-credentials: false - - name: Set up Python 3.11 + - name: Set up Python 3.8 uses: actions/setup-python@v4 id: setup-python with: - python-version: "3.11" + python-version: "3.8" - name: Install Poetry uses: snok/install-poetry@v1 with: diff --git a/.gitignore b/.gitignore index 46b59014..c98bb428 100644 --- a/.gitignore +++ b/.gitignore @@ -12,6 +12,7 @@ !tests/data/**/*.sql Pipfile snapshot_report.html +.harlequin.toml # Byte-compiled / optimized / DLL files __pycache__/ diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index 2f36fb73..bf677561 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -1,15 +1,15 @@ repos: - repo: https://github.com/psf/black - rev: 23.10.0 + rev: 23.11.0 hooks: - id: black - - repo: https://github.com/charliermarsh/ruff-pre-commit - rev: v0.1.1 + - repo: https://github.com/astral-sh/ruff-pre-commit + rev: v0.1.6 hooks: - id: ruff args: [ --fix, --exit-non-zero-on-fix ] - repo: https://github.com/pre-commit/mirrors-mypy - rev: v1.6.1 + rev: v1.7.1 hooks: - id: mypy additional_dependencies: @@ -22,6 +22,8 @@ repos: - pytest - types-pygments - rich-click>=1.7.1 + - questionary + - tomlkit args: - "--disallow-untyped-calls" - "--disallow-untyped-defs" diff --git a/CHANGELOG.md b/CHANGELOG.md index 5fb15a4b..fd58d071 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -4,6 +4,12 @@ All notable changes to this project will be documented in this file. ## [Unreleased] +### Features + +- Harlequin can now be configured using a TOML file. The config file can both specify options for Harlequin (like the theme and row limit) and also for installed adapters (like the host, username, and password for a database connection). The config file can define multiple "profiles" (sets of configuration), and you can select the profile to use when starting Harlequin with the `--profile` option (alias `-P`). By default, Harlequin searches the current directory and home directories for files called either `.harlequin.toml` or `pyproject.toml`, and merges the config it finds in them. You can specify a different path using the `--config-path` option. Values loaded from config files can be overridden by passing CLI options ([#206](https://github.com/tconbeer/harlequin/issues/206)). +- Harlequin now ships with a wizard to make it easy to create or update config files. Simply run Harlequin with the `--config` option. +- Adds a `harlequin` theme. You can use it with `harlequin -t harlequin`. + ## [1.5.0] - 2023-11-28 ### Breaking Changes diff --git a/poetry.lock b/poetry.lock index cc19adfc..20b93bd3 100644 --- a/poetry.lock +++ b/poetry.lock @@ -391,13 +391,13 @@ files = [ [[package]] name = "harlequin-postgres" -version = "0.1.2" +version = "0.1.3" description = "A Harlequin adapter for Postgres." -optional = true +optional = false python-versions = ">=3.8.1,<4.0" files = [ - {file = "harlequin_postgres-0.1.2-py3-none-any.whl", hash = "sha256:76cdf867c792be25ebf281bc4769096826fed6dcf4f9a1cfb3e45ebcc1ca768f"}, - {file = "harlequin_postgres-0.1.2.tar.gz", hash = "sha256:df5014a5d7e729cd84bd1596e6815170ab55c66b7d47d92fb93154624fff39af"}, + {file = "harlequin_postgres-0.1.3-py3-none-any.whl", hash = "sha256:db7ab7f7304a18dce91e41ced87cc8d56a46f9c20aadaefd2306eb27f6e2edd0"}, + {file = "harlequin_postgres-0.1.3.tar.gz", hash = "sha256:a5060b7311a09a12efd5b66f7c3c787ebb8149ee7839625de2dfe931877b1ca1"}, ] [package.dependencies] @@ -948,11 +948,25 @@ nodeenv = ">=0.11.1" pyyaml = ">=5.1" virtualenv = ">=20.10.0" +[[package]] +name = "prompt-toolkit" +version = "3.0.36" +description = "Library for building powerful interactive command lines in Python" +optional = false +python-versions = ">=3.6.2" +files = [ + {file = "prompt_toolkit-3.0.36-py3-none-any.whl", hash = "sha256:aa64ad242a462c5ff0363a7b9cfe696c20d55d9fc60c11fd8e632d064804d305"}, + {file = "prompt_toolkit-3.0.36.tar.gz", hash = "sha256:3e163f254bef5a03b146397d7c1963bd3e2812f0964bb9a24e6ec761fd28db63"}, +] + +[package.dependencies] +wcwidth = "*" + [[package]] name = "psycopg2-binary" version = "2.9.9" description = "psycopg2 - Python-PostgreSQL Database Adapter" -optional = true +optional = false python-versions = ">=3.7" files = [ {file = "psycopg2-binary-2.9.9.tar.gz", hash = "sha256:7f01846810177d829c7692f1f5ada8096762d9172af1b1a28d4ab5b77c923c1c"}, @@ -1223,6 +1237,20 @@ files = [ {file = "PyYAML-6.0.1.tar.gz", hash = "sha256:bfdf460b1736c775f2ba9f6a92bca30bc2095067b8a9d77876d1fad6cc3b4a43"}, ] +[[package]] +name = "questionary" +version = "2.0.1" +description = "Python library to build pretty command line user prompts ⭐️" +optional = false +python-versions = ">=3.8" +files = [ + {file = "questionary-2.0.1-py3-none-any.whl", hash = "sha256:8ab9a01d0b91b68444dff7f6652c1e754105533f083cbe27597c8110ecc230a2"}, + {file = "questionary-2.0.1.tar.gz", hash = "sha256:bcce898bf3dbb446ff62830c86c5c6fb9a22a54146f0f5597d3da43b10d8fc8b"}, +] + +[package.dependencies] +prompt_toolkit = ">=2.0,<=3.0.36" + [[package]] name = "rich" version = "13.7.0" @@ -1418,6 +1446,17 @@ files = [ {file = "tomli-2.0.1.tar.gz", hash = "sha256:de526c12914f0c550d15924c62d72abc48d6fe7364aa87328337a31007fe8a4f"}, ] +[[package]] +name = "tomlkit" +version = "0.12.3" +description = "Style preserving TOML library" +optional = false +python-versions = ">=3.7" +files = [ + {file = "tomlkit-0.12.3-py3-none-any.whl", hash = "sha256:b0a645a9156dc7cb5d3a1f0d4bab66db287fcb8e0430bdd4664a095ea16414ba"}, + {file = "tomlkit-0.12.3.tar.gz", hash = "sha256:75baf5012d06501f07bee5bf8e801b9f343e7aac5a92581f20f80ce632e6b5a4"}, +] + [[package]] name = "tqdm" version = "4.66.1" @@ -1466,13 +1505,13 @@ types-setuptools = "*" [[package]] name = "types-setuptools" -version = "68.2.0.2" +version = "69.0.0.0" description = "Typing stubs for setuptools" optional = false python-versions = ">=3.7" files = [ - {file = "types-setuptools-68.2.0.2.tar.gz", hash = "sha256:09efc380ad5c7f78e30bca1546f706469568cf26084cfab73ecf83dea1d28446"}, - {file = "types_setuptools-68.2.0.2-py3-none-any.whl", hash = "sha256:d5b5ff568ea2474eb573dcb783def7dadfd9b1ff638bb653b3c7051ce5aeb6d1"}, + {file = "types-setuptools-69.0.0.0.tar.gz", hash = "sha256:b0a06219f628c6527b2f8ce770a4f47550e00d3e8c3ad83e2dc31bc6e6eda95d"}, + {file = "types_setuptools-69.0.0.0-py3-none-any.whl", hash = "sha256:8c86195bae2ad81e6dea900a570fe9d64a59dbce2b11cc63c046b03246ea77bf"}, ] [[package]] @@ -1520,6 +1559,17 @@ platformdirs = ">=3.9.1,<5" docs = ["furo (>=2023.7.26)", "proselint (>=0.13)", "sphinx (>=7.1.2)", "sphinx-argparse (>=0.4)", "sphinxcontrib-towncrier (>=0.2.1a0)", "towncrier (>=23.6)"] test = ["covdefaults (>=2.3)", "coverage (>=7.2.7)", "coverage-enable-subprocess (>=1)", "flaky (>=3.7)", "packaging (>=23.1)", "pytest (>=7.4)", "pytest-env (>=0.8.2)", "pytest-freezer (>=0.4.8)", "pytest-mock (>=3.11.1)", "pytest-randomly (>=3.12)", "pytest-timeout (>=2.1)", "setuptools (>=68)", "time-machine (>=2.10)"] +[[package]] +name = "wcwidth" +version = "0.2.12" +description = "Measures the displayed width of unicode strings in a terminal" +optional = false +python-versions = "*" +files = [ + {file = "wcwidth-0.2.12-py2.py3-none-any.whl", hash = "sha256:f26ec43d96c8cbfed76a5075dac87680124fa84e0855195a6184da9c187f133c"}, + {file = "wcwidth-0.2.12.tar.gz", hash = "sha256:f01c104efdf57971bcb756f054dd58ddec5204dd15fa31d6503ea57947d97c02"}, +] + [[package]] name = "yarl" version = "1.9.3" @@ -1644,4 +1694,4 @@ postgres = ["harlequin-postgres"] [metadata] lock-version = "2.0" python-versions = ">=3.8.1,<4.0.0" -content-hash = "0e6d8b017ebf41b2bb628c8ce9c2f07a3eba477c5e56b6c25d67904a10fcdb3e" +content-hash = "d552447fc1f8cf080a29714a3449d2e1b0a4ee4430c41add41b694f62262eda9" diff --git a/pyproject.toml b/pyproject.toml index 8a94e0c8..dae66aed 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -35,13 +35,17 @@ shandy-sqlfmt = ">=0.19.0" platformdirs = "^3.10.0" pyperclip = "^1.8.2" importlib_metadata = { version = ">=4.6.0", python = "<3.10.0" } +tomli = { version = "^2.0.1", python = "<3.11.0" } +tomlkit = "^0.12.3" # database adapters (optional installs for extras) harlequin-postgres = { version = "^0.1", optional = true } +questionary = "^2.0.1" [tool.poetry.group.dev.dependencies] pre-commit = "^3.3.1" textual-dev = "^1.0.1" +harlequin-postgres = "^0.1.3" [tool.poetry.group.static.dependencies] black = "^23.3.0" @@ -66,6 +70,9 @@ postgres = ["harlequin-postgres"] duckdb = "harlequin_duckdb:DuckDbAdapter" sqlite = "harlequin_sqlite:HarlequinSqliteAdapter" +[tool.poetry.plugins."pygments.styles"] +harlequin = "harlequin.colors:HarlequinPygmentsStyle" + [tool.ruff] select = ["A", "B", "E", "F", "I"] target-version = "py38" diff --git a/src/harlequin/adapter.py b/src/harlequin/adapter.py index cb95ac64..e55501cd 100644 --- a/src/harlequin/adapter.py +++ b/src/harlequin/adapter.py @@ -166,7 +166,10 @@ def __init__(self, conn_str: Sequence[str], **options: Any) -> None: - **options (Any): Options received from the command line, config file, or env variables. Adapters should be robust to receiving both subsets and supersets of their declared options. They should disregard any - extra (unexpected) kwargs. + extra (unexpected) kwargs. Adapters should check the types of options, + as they may not be cast to the correct types. + + Raises: HarlequinConfigError if a received option is the wrong value or type. """ pass diff --git a/src/harlequin/app.py b/src/harlequin/app.py index 7da67052..dea4b390 100644 --- a/src/harlequin/app.py +++ b/src/harlequin/app.py @@ -5,8 +5,6 @@ from functools import partial from typing import Dict, List, Optional, Type, Union -from rich import print -from rich.panel import Panel from textual import work from textual.app import App, ComposeResult from textual.binding import Binding @@ -38,9 +36,11 @@ export_callback, ) from harlequin.exception import ( + HarlequinConfigError, HarlequinConnectionError, HarlequinQueryError, HarlequinThemeError, + pretty_print_error, ) @@ -93,7 +93,7 @@ def __init__( self, adapter: HarlequinAdapter, theme: str = "monokai", - max_results: int = 100_000, + max_results: int | str = 100_000, driver_class: Union[Type[Driver], None] = None, css_path: Union[CSSPathType, None] = None, watch_css: bool = False, @@ -101,26 +101,24 @@ def __init__( super().__init__(driver_class, css_path, watch_css) self.adapter = adapter self.theme = theme - self.max_results = max_results - self.limit = min(500, max_results) if max_results > 0 else 500 + try: + self.max_results = int(max_results) + except ValueError: + pretty_print_error( + HarlequinConfigError( + f"limit={max_results!r} was set by config file but is not " + "a valid integer." + ) + ) + self.exit(return_code=2) + else: + self.limit = min(500, self.max_results) if self.max_results > 0 else 500 self.query_timer: Union[float, None] = None try: self.connection = self.adapter.connect() except HarlequinConnectionError as e: - print( - Panel.fit( - str(e), - title=e.title - if e.title - else ( - "Harlequin encountered an error " - "while connecting to the database." - ), - title_align="left", - border_style="red", - ) - ) - self.exit() + pretty_print_error(e) + self.exit(return_code=2) else: if self.connection.init_message: self.notify(self.connection.init_message) @@ -128,20 +126,8 @@ def __init__( try: self.app_colors = HarlequinColors.from_theme(theme) except HarlequinThemeError as e: - print( - Panel.fit( - ( - f"No theme found with the name {e}.\n" - "Theme must be the name of a Pygments Style. " - "You can browse the supported styles here:\n" - "https://pygments.org/styles/" - ), - title="Harlequin couldn't load your theme.", - title_align="left", - border_style="red", - ) - ) - self.exit() + pretty_print_error(e) + self.exit(return_code=2) else: self.design = self.app_colors.design_system self.stylesheet = Stylesheet(variables=self.get_css_variables()) diff --git a/src/harlequin/cli.py b/src/harlequin/cli.py index 8867c994..d771f19d 100644 --- a/src/harlequin/cli.py +++ b/src/harlequin/cli.py @@ -1,24 +1,31 @@ from __future__ import annotations import sys -from typing import Any, Callable +from pathlib import Path +from typing import Any, Callable, Sequence import rich_click as click from harlequin import Harlequin from harlequin.adapter import HarlequinAdapter +from harlequin.colors import GREEN, PINK, PURPLE, YELLOW +from harlequin.config import get_config_for_profile +from harlequin.config_wizard import wizard +from harlequin.exception import HarlequinConfigError, pretty_print_error +from harlequin.plugins import load_plugins if sys.version_info < (3, 10): from importlib_metadata import entry_points, version else: from importlib.metadata import entry_points, version +# configure defaults +DEFAULT_ADAPTER = "duckdb" +DEFAULT_LIMIT = 100_000 +DEFAULT_THEME = "monokai" + # configure the rich click interface (mostly --help options) DOCS_URL = "https://harlequin.sh/docs/getting-started" -GREEN = "#45FFCA" -YELLOW = "#FEFFAC" -PINK = "#FFB6D9" -PURPLE = "#D67BFF" # general click.rich_click.USE_RICH_MARKUP = True @@ -53,7 +60,16 @@ "harlequin": [ { "name": "Harlequin Options", - "options": ["--adapter", "--theme", "--limit", "--version", "--help"], + "options": [ + "--profile", + "--config-path", + "--adapter", + "--theme", + "--limit", + "--config", + "--version", + "--help", + ], }, ] } @@ -82,6 +98,13 @@ def _version_option() -> str: return output +def _config_wizard_callback(ctx: click.Context, param: Any, value: bool) -> None: + if not value or ctx.resilient_parsing: + return + wizard() + ctx.exit(0) + + def build_cli() -> click.Command: """ Loads installed adapters and constructs a click Command that includes options @@ -89,17 +112,7 @@ def build_cli() -> click.Command: Returns: click.Command """ - adapter_eps = entry_points(group="harlequin.adapter") - adapters: dict[str, type[HarlequinAdapter]] = {} - - for ep in adapter_eps: - try: - adapters.update({ep.name: ep.load()}) - except ImportError as e: - print( - f"Harlequin could not load the installed plug-in named {ep.name}." - f"\n\n{e}" - ) + adapters = load_plugins() @click.command() @click.version_option(package_name="harlequin", message=_version_option()) @@ -107,10 +120,37 @@ def build_cli() -> click.Command: "conn_str", nargs=-1, ) + @click.option( + "-P", + "--profile", + help=( + "Select a profile from an available config file to load its values. " + "Other options passed here will take precedence over those loaded " + "from the profile. Use the special profile named None to use Harlequin's " + "defaults, instead of the default profile specified in the config " + "file." + ), + ) + @click.option( + "--config-path", + help=( + "By default, Harlequin finds files named .harlequin.toml in the " + "current directory and the home directory (~) and merges them. " + "Use this option to specify the full path to a config file at " + "a different location." + ), + type=click.Path( + exists=True, + file_okay=True, + dir_okay=False, + resolve_path=True, + path_type=Path, + ), + ) @click.option( "-t", "--theme", - default="monokai", + default=DEFAULT_THEME, show_default=True, help=( "Set the theme (colors) of the Harlequin IDE. " @@ -121,7 +161,7 @@ def build_cli() -> click.Command: @click.option( "--limit", "-l", - default=100_000, + default=DEFAULT_LIMIT, type=click.IntRange(min=0), help=( "Set the maximum number of rows that can be loaded into Harlequin's " @@ -131,7 +171,7 @@ def build_cli() -> click.Command: @click.option( "-a", "--adapter", - default="duckdb", + default=DEFAULT_ADAPTER, show_default=True, type=click.Choice(list(adapters.keys()), case_sensitive=False), help=( @@ -139,13 +179,22 @@ def build_cli() -> click.Command: "to use to connect to the database at CONN_STR." ), ) + @click.option( + "--config", + help=( + "Run the configuration wizard to create or update a Harlequin " + "config file." + ), + is_flag=True, + callback=_config_wizard_callback, + expose_value=True, + is_eager=True, + ) @click.pass_context def inner_cli( ctx: click.Context, - conn_str: tuple[str], - theme: str, - limit: int, - adapter: str, + profile: str | None, + config_path: Path | None, **kwargs: Any, ) -> None: """ @@ -155,6 +204,13 @@ def inner_cli( connection strings (or paths to local db files) for databases to open with Harlequin.[/] """ + # load config from any config files + try: + config = get_config_for_profile(config_path=config_path, profile=profile) + except HarlequinConfigError as e: + pretty_print_error(e) + ctx.exit(2) + # prune the kwargs to only those that don't have their default arguments params = list(kwargs.keys()) for k in params: @@ -163,15 +219,29 @@ def inner_cli( == click.core.ParameterSource.DEFAULT # type: ignore ): kwargs.pop(k) + # conn_str is an arg, not an option, so get_paramter_source is always CLI + elif k == "conn_str" and kwargs[k] == tuple(): + kwargs.pop(k) + + # merge the config and the cli options + config.update(kwargs) # load and instantiate the adapter - adapter_cls: type[HarlequinAdapter] = adapters[adapter] - adapter_instance = adapter_cls(conn_str=conn_str, **kwargs) + adapter = config.pop("adapter", DEFAULT_ADAPTER) + conn_str: Sequence[str] = config.pop("conn_str", tuple()) # type: ignore + if isinstance(conn_str, str): + conn_str = (conn_str,) + adapter_cls: type[HarlequinAdapter] = adapters[adapter] # type: ignore + try: + adapter_instance = adapter_cls(conn_str=conn_str, **config) + except HarlequinConfigError as e: + pretty_print_error(e) + ctx.exit(2) tui = Harlequin( adapter=adapter_instance, - max_results=limit, - theme=theme, + max_results=config.get("limit", DEFAULT_LIMIT), # type: ignore + theme=config.get("theme", DEFAULT_THEME), # type: ignore ) tui.run() diff --git a/src/harlequin/colors.py b/src/harlequin/colors.py index 108f0ec8..dd391446 100644 --- a/src/harlequin/colors.py +++ b/src/harlequin/colors.py @@ -1,12 +1,75 @@ from typing import Dict, Union +from pygments.style import Style as PygmentsStyle from pygments.styles import get_style_by_name -from pygments.token import Token +from pygments.token import ( + Comment, + Error, + Keyword, + Literal, + Name, + Number, + Operator, + Punctuation, + String, + Token, +) from pygments.util import ClassNotFound +from questionary import Style as QuestionaryStyle from textual.design import ColorSystem from harlequin.exception import HarlequinThemeError +GREEN = "#45FFCA" +YELLOW = "#FEFFAC" +PINK = "#FFB6D9" +PURPLE = "#D67BFF" +GRAY = "#777777" +DARK_GRAY = "#333333" +BLACK = "#0C0C0C" +WHITE = "#DDDDDD" + + +class HarlequinPygmentsStyle(PygmentsStyle): + styles = { + Token: WHITE, + Comment: f"{GRAY} italic", + Keyword: f"{YELLOW} bold", + Keyword.Type: f"{YELLOW} nobold", + Keyword.Constant: f"{PINK} bold", + Name: WHITE, + Name.Builtin: YELLOW, + Name.Constant: YELLOW, + Name.Quoted: f"{WHITE} bold", + Name.Variable: f"{WHITE} bold", + String: PINK, + String.Symbol: f"{WHITE} bold", + String.Name: f"{WHITE} bold", + Operator: GREEN, + Punctuation: PURPLE, + Number: f"{PINK} bold", + Literal: PINK, + Error: PINK, + } + background_color = BLACK + highlight_color = DARK_GRAY + + +HARLEQUIN_QUESTIONARY_STYLE = QuestionaryStyle( + [ + ("qmark", f"fg:{GREEN} bold"), + ("question", "bold"), + ("answer", f"fg:{YELLOW} bold"), + ("pointer", f"fg:{YELLOW} bold"), + ("highlighted", f"fg:{YELLOW} bold"), + ("selected", f"fg:{YELLOW} noreverse bold"), + ("separator", f"fg:{PURPLE}"), + ("instruction", "fg:#858585 italic"), + ("text", ""), + ("disabled", "fg:#858585 italic"), + ] +) + def extract_color(s: str) -> str: for part in s.split(" "): @@ -102,7 +165,15 @@ def from_theme(cls, theme: str) -> "HarlequinColors": try: style = get_style_by_name(theme) except ClassNotFound as e: - raise HarlequinThemeError(theme) from e + raise HarlequinThemeError( + ( + f"No theme found with the name {theme}.\n" + "Theme must be the name of a Pygments Style. " + "You can browse the supported styles here:\n" + "https://pygments.org/styles/" + ), + title="Harlequin couldn't load your theme.", + ) from e background = style.background_color highlight = style.highlight_color diff --git a/src/harlequin/components/run_query_bar.py b/src/harlequin/components/run_query_bar.py index 3091112f..6b02acbf 100644 --- a/src/harlequin/components/run_query_bar.py +++ b/src/harlequin/components/run_query_bar.py @@ -25,7 +25,7 @@ def __init__( def compose(self) -> ComposeResult: yield Checkbox("Limit ", id="limit_checkbox") yield Input( - "500", + str(min(500, self.max_results)), id="limit_input", validators=Integer( minimum=0, diff --git a/src/harlequin/config.py b/src/harlequin/config.py index 27b8d89b..c127abcb 100644 --- a/src/harlequin/config.py +++ b/src/harlequin/config.py @@ -1,14 +1,146 @@ +from __future__ import annotations + +import sys from pathlib import Path -from typing import Tuple +from typing import Dict, List, Union + +from harlequin.exception import HarlequinConfigError + +if sys.version_info < (3, 11): + import tomli as tomllib +else: + import tomllib + +CONFIG_FILENAMES = ("pyproject.toml", ".harlequin.toml") # order matters! +SEARCH_DIRS = (Path.home(), Path.cwd()) + +Profile = Dict[str, Union[bool, int, List[str], str, Path]] +Config = Dict[str, Union[str, Dict[str, Profile]]] -def get_init_script(init_path: Path, no_init: bool) -> Tuple[Path, str]: - if no_init: - init_script = "" +def get_config_for_profile(config_path: Path | None, profile: str | None) -> Profile: + config = load_config(config_path) + active_profile = profile or config.get("default_profile", None) + if active_profile is None or active_profile == "None": + return {} + elif active_profile not in config.get("profiles", {}): + raise HarlequinConfigError( + f"Could not load the profile named {active_profile} because it does not " + "exist in any discovered config files.", + title="Harlequin couldn't load your profile.", + ) else: + return config["profiles"][active_profile] # type: ignore + + +def load_config(config_path: Path | None) -> Config: + paths = _find_config_files(config_path) + config = _merge_config_files(paths) + _raise_on_bad_schema(config) + return config + + +def _find_config_files(config_path: Path | None) -> list[Path]: + found_files: list[Path] = [] + if config_path is None: + for filename in CONFIG_FILENAMES: + for p in [p / filename for p in SEARCH_DIRS]: + if p.exists(): + found_files.append(p) + elif config_path.exists(): + found_files.append(config_path) + else: + raise HarlequinConfigError( + f"Config file could not be found at specified path: {config_path}", + title="Harlequin couldn't load your config file.", + ) + return found_files + + +def _merge_config_files(paths: list[Path]) -> Config: + config: Config = {} + for p in paths: try: - with open(init_path.expanduser(), "r") as f: - init_script = f.read() - except OSError: - init_script = "" - return init_path, init_script + with open(p, "rb") as f: + raw_config = tomllib.load(f) + except OSError as e: + raise HarlequinConfigError( + f"Error opening config file at {p}. {e}", + title="Harlequin couldn't load your config file.", + ) from e + except tomllib.TOMLDecodeError as e: + raise HarlequinConfigError( + f"Error decoding config file at {p}. " f"Check for invalid TOML. {e}", + title="Harlequin couldn't load your config file.", + ) from e + relevant_config = ( + raw_config + if p.stem != "pyproject" + else raw_config.get("tool", {}).get("harlequin", {}) + ) + config.update(relevant_config) + return config + + +def _raise_on_bad_schema(config: Config) -> None: + TOP_LEVEL_KEYS = ("default_profile", "profiles") + if not config: + return + + for k in config.keys(): + if k not in TOP_LEVEL_KEYS: + raise HarlequinConfigError( + f"Found unexpected key in config: {k}.\n" + f"Allowed values are {TOP_LEVEL_KEYS}.", + title="Harlequin couldn't load your config file.", + ) + if config.get("profiles", None) is None: + pass + elif not isinstance(config["profiles"], dict): + raise HarlequinConfigError( + "The profiles key must define a table.", + title="Harlequin couldn't load your config file.", + ) + elif not all( + [isinstance(config["profiles"][k], dict) for k in config["profiles"].keys()] + ): + raise HarlequinConfigError( + "The members of the profiles table must be tables.", + title="Harlequin couldn't load your config file.", + ) + elif any(k == "None" for k in config["profiles"].keys()): + raise HarlequinConfigError( + "Config file defines a profile named 'None', which is not allowed.", + title="Harlequin couldn't load your config file.", + ) + else: + for profile_name, opt_dict in config["profiles"].items(): + for option_name in opt_dict.keys(): + if "-" in option_name: + raise HarlequinConfigError( + f"Profile {profile_name} defines an option '{option_name}'," + "which is an invalid name for an option. Did you mean " + f"""'{option_name.strip("-").replace("-", "_")}'?""", + title="Harlequin couldn't load your config file.", + ) + + if (default := config.get("default_profile", None)) is not None and not isinstance( + default, str + ): + raise HarlequinConfigError( + f"Config file sets default_profile to {default}, but that value " + "must be a string.", + title="Harlequin couldn't load your config file.", + ) + elif ( + default is not None + and isinstance(default, str) + and isinstance(config["profiles"], dict) + and default != "None" + and config["profiles"].get(default, None) is None + ): + raise HarlequinConfigError( + f"Config file sets default_profile to {default}, but does not define a " + "profile with that name.", + title="Harlequin couldn't load your config file.", + ) diff --git a/src/harlequin/config_wizard.py b/src/harlequin/config_wizard.py new file mode 100644 index 00000000..51832e5d --- /dev/null +++ b/src/harlequin/config_wizard.py @@ -0,0 +1,261 @@ +from __future__ import annotations + +from pathlib import Path +from typing import Any + +import questionary +import tomlkit +from pygments.styles import get_all_styles +from rich import print as rich_print +from rich.panel import Panel +from rich.syntax import Syntax +from tomlkit.exceptions import TOMLKitError +from tomlkit.toml_document import TOMLDocument +from tomlkit.toml_file import TOMLFile + +from harlequin.adapter import HarlequinAdapter +from harlequin.colors import HARLEQUIN_QUESTIONARY_STYLE, YELLOW +from harlequin.exception import HarlequinWizardError, pretty_print_error +from harlequin.options import ListOption +from harlequin.plugins import load_plugins + + +def wizard() -> None: + try: + _wizard() + except KeyboardInterrupt: + print("Cancelled config updates. No changes were made to any files.") + return + except HarlequinWizardError as e: + pretty_print_error(e) + return + + +def _wizard() -> None: + path, is_pyproject = _prompt_for_path() + config, file = _read_toml(path) + + # extract existing profiles from config file. + if is_pyproject: + full_config = config + config = config.get("tool", {}).get("harlequin", {}) + if "profiles" not in config: + config["profiles"] = {} + profiles = config.get("profiles", {}) + + profile_name = _prompt_for_profile_name(profiles) + selected_profile = profiles.get(profile_name, {}) + + adapters = load_plugins() + adapter = questionary.select( + message="Which adapter should this profile use?", + choices=sorted(adapters.keys()), + default=selected_profile.get("adapter", "duckdb"), + style=HARLEQUIN_QUESTIONARY_STYLE, + ).unsafe_ask() + + conn_str = questionary.text( + message="What connection string(s) should this profile use?", + instruction="Separate items by a space.", + default=" ".join(selected_profile.get("conn_str", [])), + style=HARLEQUIN_QUESTIONARY_STYLE, + ).unsafe_ask() + + theme = questionary.select( + message="What theme should this profile use?", + choices=list(get_all_styles()), + default=selected_profile.get("theme", "monokai"), + style=HARLEQUIN_QUESTIONARY_STYLE, + ).unsafe_ask() + + limit = int( + questionary.text( + message="How many rows should the data table show?", + validate=_validate_int, + default=str(selected_profile.get("limit", 100000)), + style=HARLEQUIN_QUESTIONARY_STYLE, + ).unsafe_ask() + ) + + adapter_cls = adapters[adapter] + adapter_option_choices = ( + [ + questionary.Choice( + title=opt.name, checked=_sluggify_name(opt.name) in selected_profile + ) + for opt in adapter_cls.ADAPTER_OPTIONS + ] + if adapter_cls.ADAPTER_OPTIONS is not None + else [] + ) + which = questionary.checkbox( + message="Which of the following adapter options would you like to set?", + choices=adapter_option_choices, + style=HARLEQUIN_QUESTIONARY_STYLE, + ).unsafe_ask() + + adapter_options = {} + if conn_str: + adapter_options["conn_str"] = conn_str.split(" ") + _prompt_to_set_adapter_options( + adapter_options=adapter_options, + adapter_cls=adapter_cls, + which=which, + selected_profile=selected_profile, + ) + + default_profile = _prompt_to_set_default_profile(profile_name, config, profiles) + + new_profile = { + "adapter": adapter, + "theme": theme, + "limit": limit, + **adapter_options, + } + + _confirm_profile_generation(default_profile, profile_name, new_profile) + + config["profiles"][profile_name] = new_profile # type: ignore + + if is_pyproject: + if "tool" not in full_config: + full_config["tool"] = {} + full_config["tool"]["harlequin"] = config # type: ignore + config = full_config + + file.write(config) + + +def _prompt_for_path() -> tuple[Path, bool]: + raw_path: str = questionary.path( + "What config file do you want to create or update?", + default=".harlequin.toml", + validate=lambda p: True + if p.endswith(".toml") + else "Must have a .toml extension", + style=HARLEQUIN_QUESTIONARY_STYLE, + ).unsafe_ask() + path = Path(raw_path) + is_pyproject = path.stem == "pyproject" + if path.suffix != ".toml": + raise HarlequinWizardError( + msg="Must create a file with a .toml extension.", + title="Harlequin could not create your configuration.", + ) + return path, is_pyproject + + +def _read_toml(path: Path) -> tuple[TOMLDocument, TOMLFile]: + file = TOMLFile(path) + try: + config = file.read() + except OSError: + config = TOMLDocument() + except TOMLKitError as e: + raise HarlequinWizardError( + f"Attempted to load the config file at {path}, but encountered an " + f"error:\n\n{e}", + title="Harlequin could not load the config file.", + ) from e + return config, file + + +def _prompt_for_profile_name(profiles: TOMLDocument) -> str: + NEW_PROFILE_SENTINEL = "[Create a New Profile]" + profile_name = NEW_PROFILE_SENTINEL + if profiles: + profile_name = questionary.select( + message="Which profile would you like to update?", + choices=[NEW_PROFILE_SENTINEL, *profiles.keys()], + style=HARLEQUIN_QUESTIONARY_STYLE, + ).unsafe_ask() + if profile_name == NEW_PROFILE_SENTINEL: + profile_name = questionary.text( + message="What would you like to name your profile?", + style=HARLEQUIN_QUESTIONARY_STYLE, + validate=lambda x: True if x and x != "None" else "Cannot be empty or None", + ).unsafe_ask() + return profile_name + + +def _prompt_to_set_adapter_options( + adapter_options: dict[str, Any], + adapter_cls: type[HarlequinAdapter], + which: list[str], + selected_profile: TOMLDocument, +) -> None: + """ + Mutates passed adapter_options dict. + """ + if which and adapter_cls.ADAPTER_OPTIONS is not None: + for option in adapter_cls.ADAPTER_OPTIONS: + if option.name not in which: + continue + value = option.to_questionary( + selected_profile.get(_sluggify_name(option.name), None) + ).unsafe_ask() + if isinstance(option, ListOption): + value = value.split(" ") + adapter_options.update({_sluggify_name(option.name): value}) + + +def _prompt_to_set_default_profile( + profile_name: str, config: TOMLDocument, profiles: TOMLDocument +) -> str | None: + possible_names = set([profile_name, *profiles.keys()]) + NO_DEFAULT_SENTINEL = "[No default]" + default_profile: str = questionary.select( + message="Would you like to set a default profile?", + choices=[ + NO_DEFAULT_SENTINEL, + *possible_names, + ], + default=config.get("default_profile", None), + style=HARLEQUIN_QUESTIONARY_STYLE, + ).unsafe_ask() + + if default_profile == NO_DEFAULT_SENTINEL: + _ = config.pop("default_profile", None) + return None + else: + config["default_profile"] = default_profile + return default_profile + + +def _confirm_profile_generation( + default_profile: str | None, profile_name: str, new_profile: dict[str, Any] +) -> None: + new_config: dict[str, Any] = ( + {} if default_profile is None else {"default_profile": default_profile} + ) + new_config.update({"profiles": {profile_name: new_profile}}) + new_config_toml = tomlkit.dumps(new_config).rstrip() + + rich_print("[italic] We generated the following profile:[/]") + rich_print( + Panel.fit( + Syntax(code=new_config_toml, lexer="toml", theme="harlequin"), + border_style=YELLOW, + ) + ) + + all_good = questionary.confirm( + "Save this profile?", + style=HARLEQUIN_QUESTIONARY_STYLE, + ).ask() + + if not all_good: + raise KeyboardInterrupt() + + +def _validate_int(raw: str) -> bool: + try: + int(raw) + except ValueError: + return False + else: + return True + + +def _sluggify_name(raw: str) -> str: + return raw.strip("-").replace("-", "_") diff --git a/src/harlequin/exception.py b/src/harlequin/exception.py index 1128104b..2ba9c589 100644 --- a/src/harlequin/exception.py +++ b/src/harlequin/exception.py @@ -3,29 +3,45 @@ class HarlequinExit(Exception): class HarlequinError(Exception): - pass - - -class HarlequinConnectionError(HarlequinError): def __init__(self, msg: str, title: str = "") -> None: super().__init__(msg) self.msg = msg self.title = title +class HarlequinConnectionError(HarlequinError): + pass + + class HarlequinCopyError(HarlequinError): - def __init__(self, msg: str, title: str = "") -> None: - super().__init__(msg) - self.msg = msg - self.title = title + pass class HarlequinQueryError(HarlequinError): - def __init__(self, msg: str, title: str = "") -> None: - super().__init__(msg) - self.msg = msg - self.title = title + pass class HarlequinThemeError(HarlequinError): pass + + +class HarlequinConfigError(HarlequinError): + pass + + +class HarlequinWizardError(HarlequinError): + pass + + +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", + ) + ) diff --git a/src/harlequin/options.py b/src/harlequin/options.py index bb79323d..d2b824df 100644 --- a/src/harlequin/options.py +++ b/src/harlequin/options.py @@ -3,12 +3,14 @@ import re from abc import ABC, abstractmethod from pathlib import Path -from typing import Any, Callable, Generator, Sequence +from typing import Any, Callable, Generator, Iterable, Sequence import click +import questionary from textual.validation import ValidationResult, Validator from textual.widget import Widget +from harlequin.colors import HARLEQUIN_QUESTIONARY_STYLE from harlequin.copy_widgets import ( Input, NoFocusLabel, @@ -91,6 +93,10 @@ def to_click(self) -> Callable[[click.Command], click.Command]: def to_widgets(self) -> Generator[Widget, None, None]: pass + @abstractmethod + def to_questionary(self, existing_value: Any | None = None) -> questionary.Question: + pass + class TextOption(AbstractOption): """ @@ -157,6 +163,31 @@ def to_widgets(self) -> Generator[Widget, None, None]: validators=[_CustomValidator(self.validator)], ) + def to_questionary(self, existing_value: Any | None = None) -> questionary.Question: + def _q_validator(raw: str) -> bool | str | None: + if self.validator is not None: + result = self.validator(raw) + if result[0]: + return True + else: + return result[1] + else: + return True + + try: + safe_existing_value = str(existing_value) + except (ValueError, TypeError): + safe_existing_value = None + + return questionary.text( + message=self.name, + default=safe_existing_value + if safe_existing_value is not None + else self.default or "", + validate=_q_validator, + style=HARLEQUIN_QUESTIONARY_STYLE, + ) + class ListOption(AbstractOption): def __init__( @@ -189,6 +220,21 @@ def to_click(self) -> Callable[[click.Command], click.Command]: def to_widgets(self) -> Generator[Widget, None, None]: raise NotImplementedError("No widget for ListOption.") + def to_questionary(self, existing_value: Any | None = None) -> questionary.Question: + if isinstance(existing_value, str): + safe_existing_value = existing_value + elif isinstance(existing_value, Iterable): + safe_existing_value = " ".join(existing_value) + else: + safe_existing_value = None + + return questionary.text( + message=self.name, + instruction="Separate items by a space.", + default=safe_existing_value if safe_existing_value is not None else "", + style=HARLEQUIN_QUESTIONARY_STYLE, + ) + class PathOption(AbstractOption): """ @@ -263,6 +309,38 @@ def to_widgets(self) -> Generator[Widget, None, None]: tab_advances_focus=True, ) + def to_questionary(self, existing_value: Any | None = None) -> questionary.Question: + def _path_validator(raw_path: str) -> bool | str: + try: + p = Path(raw_path) + except ValueError as e: + return f"Not a valid path! {e}" + if self.exists and not p.exists(): + return f"No file exists at {p}" + + if not self.file_okay and p.is_file(): + return f"{p} is a file, expected a directory." + + if not self.dir_okay and p.is_dir(): + return f"{p} is a directory, expected a file." + + return True + + try: + safe_existing_value = str(existing_value) + except (ValueError, TypeError): + safe_existing_value = None + + return questionary.path( + message=self.name, + default=safe_existing_value + if safe_existing_value is not None + else self.default or "", + only_directories=not self.file_okay, + validate=_path_validator, + style=HARLEQUIN_QUESTIONARY_STYLE, + ) + class SelectOption(AbstractOption): def __init__( @@ -292,17 +370,11 @@ def __init__( """ def to_click(self) -> Callable[[click.Command], click.Command]: - choices: list[str] = [] - for choice in self.choices: - if isinstance(choice, str): - choices.append(choice) - else: - choices.append(choice[0]) return click.option( f"--{self.name}", *self.short_decls, help=self.description, - type=click.Choice(choices=choices, case_sensitive=False), + type=click.Choice(choices=self._flat_choices(), case_sensitive=False), ) def to_widgets(self) -> Generator[Widget, None, None]: @@ -320,6 +392,33 @@ def to_widgets(self) -> Generator[Widget, None, None]: allow_blank=False, ) + def to_questionary(self, existing_value: Any | None = None) -> questionary.Question: + try: + safe_existing_value = str(existing_value) + except (ValueError, TypeError): + safe_existing_value = None + + if safe_existing_value not in self._flat_choices(): + safe_existing_value = None + + return questionary.select( + message=self.name, + choices=self._flat_choices(), + default=safe_existing_value + if safe_existing_value is not None + else self.default, + style=HARLEQUIN_QUESTIONARY_STYLE, + ) + + def _flat_choices(self) -> list[str]: + choices: list[str] = [] + for choice in self.choices: + if isinstance(choice, str): + choices.append(choice) + else: + choices.append(choice[0]) + return choices + class FlagOption(AbstractOption): """ @@ -335,6 +434,18 @@ def to_widgets(self) -> Generator[Widget, None, None]: yield NoFocusLabel(f"{self.label}:", classes="switch_label") yield Switch(id=self.name) + def to_questionary(self, existing_value: Any | None = None) -> questionary.Question: + try: + safe_existing_value = bool(existing_value) + except (ValueError, TypeError): + safe_existing_value = None + + return questionary.confirm( + message=self.name, + default=safe_existing_value if safe_existing_value is not None else False, + style=HARLEQUIN_QUESTIONARY_STYLE, + ) + HarlequinAdapterOption = AbstractOption diff --git a/src/harlequin/plugins.py b/src/harlequin/plugins.py new file mode 100644 index 00000000..c31d70eb --- /dev/null +++ b/src/harlequin/plugins.py @@ -0,0 +1,27 @@ +from __future__ import annotations + +import sys +from typing import Dict + +from harlequin.adapter import HarlequinAdapter + +if sys.version_info < (3, 10): + from importlib_metadata import entry_points +else: + from importlib.metadata import entry_points + + +def load_plugins() -> Dict[str, type[HarlequinAdapter]]: + adapter_eps = entry_points(group="harlequin.adapter") + adapters: Dict[str, type[HarlequinAdapter]] = {} + + for ep in adapter_eps: + try: + adapters.update({ep.name: ep.load()}) + except ImportError as e: + print( + f"Harlequin could not load the installed plug-in named {ep.name}." + f"\n\n{e}" + ) + + return adapters diff --git a/src/harlequin_duckdb/adapter.py b/src/harlequin_duckdb/adapter.py index eb923e8e..7fa7e42e 100644 --- a/src/harlequin_duckdb/adapter.py +++ b/src/harlequin_duckdb/adapter.py @@ -9,6 +9,7 @@ from harlequin.adapter import HarlequinAdapter, HarlequinConnection, HarlequinCursor from harlequin.catalog import Catalog, CatalogItem from harlequin.exception import ( + HarlequinConfigError, HarlequinConnectionError, HarlequinCopyError, HarlequinQueryError, @@ -323,10 +324,10 @@ class DuckDbAdapter(HarlequinAdapter): def __init__( self, conn_str: Sequence[str], - init_path: Path | None = None, - no_init: bool = False, - read_only: bool = False, - allow_unsigned_extensions: bool = False, + init_path: Path | str | None = None, + no_init: bool | str = False, + read_only: bool | str = False, + allow_unsigned_extensions: bool | str = False, extension: list[str] | None = None, force_install_extensions: bool = False, custom_extension_repo: str | None = None, @@ -334,16 +335,26 @@ def __init__( md_saas: bool = False, **_: Any, ) -> None: - self.conn_str = conn_str if conn_str else (":memory:",) - self.init_path = init_path or Path.home() / ".duckdbrc" - self.no_init = no_init - self.read_only = read_only - self.allow_unsigned_extensions = allow_unsigned_extensions - self.extensions = extension if extension is not None else [] - self.force_install_extensions = force_install_extensions - self.custom_extension_repo = custom_extension_repo - self.md_token = md_token - self.md_saas = md_saas + try: + self.conn_str = conn_str if conn_str else (":memory:",) + self.init_path = ( + Path(init_path).resolve() + if init_path is not None + else Path.home() / ".duckdbrc" + ) + self.no_init = bool(no_init) + self.read_only = bool(read_only) + self.allow_unsigned_extensions = bool(allow_unsigned_extensions) + self.extensions = extension if extension is not None else [] + self.force_install_extensions = force_install_extensions + self.custom_extension_repo = custom_extension_repo + self.md_token = md_token + self.md_saas = md_saas + except (ValueError, TypeError) as e: + raise HarlequinConfigError( + msg=f"DuckDB adapter received bad config value: {e}", + title="Harlequin could not initialize the selected adapter.", + ) from e def connect(self) -> DuckDbConnection: primary_db, *other_dbs = self.conn_str diff --git a/src/harlequin_sqlite/adapter.py b/src/harlequin_sqlite/adapter.py index 327e67ce..7b7b5910 100644 --- a/src/harlequin_sqlite/adapter.py +++ b/src/harlequin_sqlite/adapter.py @@ -8,7 +8,11 @@ from harlequin.adapter import HarlequinAdapter, HarlequinConnection, HarlequinCursor from harlequin.catalog import Catalog, CatalogItem -from harlequin.exception import HarlequinConnectionError, HarlequinQueryError +from harlequin.exception import ( + HarlequinConfigError, + HarlequinConnectionError, + HarlequinQueryError, +) from harlequin.options import HarlequinAdapterOption, HarlequinCopyFormat from textual_fastdatatable.backend import AutoBackendType @@ -189,20 +193,26 @@ def __init__( conn_str: Sequence[str], read_only: bool = False, connection_mode: Literal["ro", "rw", "rwc", "memory"] | None = None, - timeout: str = "5.0", - detect_types: str = "0", + timeout: str | float = 5.0, + detect_types: str | int = 0, isolation_level: Literal["DEFERRED", "EXCLUSIVE", "IMMEDIATE"] | None = "DEFERRED", - cached_statements: str = "128", + cached_statements: str | int = 128, **_: Any, ) -> None: - self.conn_str = conn_str if conn_str else (":memory:",) - self.read_only = read_only - self.connection_mode = connection_mode - self.timeout = timeout - self.detect_types = detect_types - self.isolation_level = isolation_level - self.cached_statements = cached_statements + try: + self.conn_str = conn_str if conn_str else (":memory:",) + self.read_only = bool(read_only) + self.connection_mode = connection_mode + self.timeout = float(timeout) + self.detect_types = int(detect_types) + self.isolation_level = isolation_level + self.cached_statements = int(cached_statements) + except (ValueError, TypeError) as e: + raise HarlequinConfigError( + msg=f"SQLite adapter received bad config value: {e}", + title="Harlequin could not initialize the selected adapter.", + ) from e def connect(self) -> HarlequinSqliteConnection: if ( @@ -245,10 +255,10 @@ def connect(self) -> HarlequinSqliteConnection: try: conn = sqlite3.connect( database=primary_db, - timeout=float(self.timeout), - detect_types=int(self.detect_types), + timeout=self.timeout, + detect_types=self.detect_types, isolation_level=self.isolation_level, - cached_statements=int(self.cached_statements), + cached_statements=self.cached_statements, check_same_thread=False, uri=True, ) diff --git a/tests/data/unit_tests/config/bad_option_name.toml b/tests/data/unit_tests/config/bad_option_name.toml new file mode 100644 index 00000000..12e38e47 --- /dev/null +++ b/tests/data/unit_tests/config/bad_option_name.toml @@ -0,0 +1,2 @@ +[profiles.my-duckdb-profile] +read-only = true diff --git a/tests/data/unit_tests/config/default_no_exist.toml b/tests/data/unit_tests/config/default_no_exist.toml new file mode 100644 index 00000000..c78a772e --- /dev/null +++ b/tests/data/unit_tests/config/default_no_exist.toml @@ -0,0 +1,5 @@ +default_profile = "foo" + +[profiles.my-duckdb-profile] +limit = 200_000 +conn_str = ["my-database.db"] \ No newline at end of file diff --git a/tests/data/unit_tests/config/extra_key.toml b/tests/data/unit_tests/config/extra_key.toml new file mode 100644 index 00000000..7f0dcb6b --- /dev/null +++ b/tests/data/unit_tests/config/extra_key.toml @@ -0,0 +1,5 @@ +default_profile = "my-duckdb-profile" +foo = "bar" + +[profiles.my-duckdb-profile] +conn_str = ["my-database.db"] \ No newline at end of file diff --git a/tests/data/unit_tests/config/good_config.toml b/tests/data/unit_tests/config/good_config.toml new file mode 100644 index 00000000..9e2a77f9 --- /dev/null +++ b/tests/data/unit_tests/config/good_config.toml @@ -0,0 +1,24 @@ +default_profile = "my-duckdb-profile" + +[profiles.my-duckdb-profile] +theme = "monokai" +limit = 200_000 +adapter = "duckdb" +conn_str = ["my-database.db"] +read_only = false +allow_unsigned_extensions = false +extension = ["httpfs", "spatial"] +force_install_extensions = false +md_saas = false +init_path = "~/.duckdbrc" + +[profiles.local-postgres] +theme = "fruity" +limit = 10_000 +adapter = "postgres" +host = "localhost" +user = "postgres" +password = "secretadminpassword" +database = "postgres" +port = 5432 +init_path = "~/.psqlrc" diff --git a/tests/data/unit_tests/config/none_profile.toml b/tests/data/unit_tests/config/none_profile.toml new file mode 100644 index 00000000..61ac9f10 --- /dev/null +++ b/tests/data/unit_tests/config/none_profile.toml @@ -0,0 +1,4 @@ +default_profile = "None" + +[profiles.None] +limit = 200_000 \ No newline at end of file diff --git a/tests/data/unit_tests/config/not_toml.toml b/tests/data/unit_tests/config/not_toml.toml new file mode 100644 index 00000000..7e99477e --- /dev/null +++ b/tests/data/unit_tests/config/not_toml.toml @@ -0,0 +1 @@ +theme: fruity \ No newline at end of file diff --git a/tests/data/unit_tests/config/profile_not_table.toml b/tests/data/unit_tests/config/profile_not_table.toml new file mode 100644 index 00000000..e4655de3 --- /dev/null +++ b/tests/data/unit_tests/config/profile_not_table.toml @@ -0,0 +1,2 @@ +[profiles] +foo = "bar" \ No newline at end of file diff --git a/tests/data/unit_tests/config/profiles_not_table.toml b/tests/data/unit_tests/config/profiles_not_table.toml new file mode 100644 index 00000000..34e3840d --- /dev/null +++ b/tests/data/unit_tests/config/profiles_not_table.toml @@ -0,0 +1 @@ +profiles = "bar" \ No newline at end of file diff --git a/tests/data/unit_tests/config/pyproject.toml b/tests/data/unit_tests/config/pyproject.toml new file mode 100644 index 00000000..71ed2c55 --- /dev/null +++ b/tests/data/unit_tests/config/pyproject.toml @@ -0,0 +1,25 @@ +[tool.harlequin] +default_profile = "my-duckdb-profile" + +[tool.harlequin.profiles.my-duckdb-profile] +theme = "monokai" +limit = 200_000 +adapter = "duckdb" +conn_str = ["my-database.db"] +read_only = false +allow_unsigned_extensions = false +extension = ["httpfs", "spatial"] +force_install_extensions = false +md_saas = false +init_path = "~/.duckdbrc" + +[tool.harlequin.profiles.local-postgres] +theme = "fruity" +limit = 10_000 +adapter = "postgres" +host = "localhost" +user = "postgres" +password = "secretadminpassword" +database = "postgres" +port = 5432 +init_path = "~/.psqlrc" diff --git a/tests/functional_tests/conftest.py b/tests/functional_tests/conftest.py index e8926bb1..ca759b54 100644 --- a/tests/functional_tests/conftest.py +++ b/tests/functional_tests/conftest.py @@ -7,3 +7,8 @@ def mock_user_cache_dir(monkeypatch: pytest.MonkeyPatch, tmp_path: Path) -> Path: monkeypatch.setattr("harlequin.cache.user_cache_dir", lambda **_: tmp_path) return tmp_path + + +@pytest.fixture(autouse=True) +def mock_config_loader(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr("harlequin.cli.get_config_for_profile", lambda **_: dict()) diff --git a/tests/unit_tests/test_cli.py b/tests/unit_tests/test_cli.py index 10ce37a3..dd129dfb 100644 --- a/tests/unit_tests/test_cli.py +++ b/tests/unit_tests/test_cli.py @@ -5,6 +5,7 @@ from click.testing import CliRunner from harlequin import Harlequin from harlequin.cli import build_cli +from harlequin.config import Config from harlequin_duckdb import DUCKDB_OPTIONS, DuckDbAdapter @@ -17,7 +18,7 @@ def mock_adapter(monkeypatch: pytest.MonkeyPatch) -> MagicMock: mock_entrypoint.load.return_value = mock_adapter mock_entry_points = MagicMock() mock_entry_points.return_value = [mock_entrypoint] - monkeypatch.setattr("harlequin.cli.entry_points", mock_entry_points) + monkeypatch.setattr("harlequin.plugins.entry_points", mock_entry_points) return mock_adapter @@ -28,9 +29,25 @@ def mock_harlequin(monkeypatch: pytest.MonkeyPatch) -> MagicMock: return mock +@pytest.fixture() +def mock_empty_config(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr("harlequin.cli.get_config_for_profile", lambda **_: dict()) + return None + + +@pytest.fixture() +def mock_load_config(monkeypatch: pytest.MonkeyPatch) -> Config: + config: Config = {"profiles": {"test-profile": {"theme": "fruity"}}} + monkeypatch.setattr("harlequin.config.load_config", lambda *_: config) + return config + + @pytest.mark.parametrize("harlequin_args", ["", ":memory:"]) def test_default( - mock_harlequin: MagicMock, mock_adapter: MagicMock, harlequin_args: str + mock_harlequin: MagicMock, + mock_adapter: MagicMock, + harlequin_args: str, + mock_empty_config: None, ) -> None: runner = CliRunner() res = runner.invoke(build_cli(), args=harlequin_args) @@ -49,7 +66,10 @@ def test_default( "harlequin_args", ["--init-path foo", ":memory: -i foo", "-init foo"] ) def test_custom_init_script( - mock_harlequin: MagicMock, mock_adapter: MagicMock, harlequin_args: str + mock_harlequin: MagicMock, + mock_adapter: MagicMock, + harlequin_args: str, + mock_empty_config: None, ) -> None: runner = CliRunner() res = runner.invoke(build_cli(), args=harlequin_args) @@ -61,7 +81,10 @@ def test_custom_init_script( @pytest.mark.parametrize("harlequin_args", ["--no-init", ":memory: --no-init"]) def test_no_init_script( - mock_harlequin: MagicMock, mock_adapter: MagicMock, harlequin_args: str + mock_harlequin: MagicMock, + mock_adapter: MagicMock, + harlequin_args: str, + mock_empty_config: None, ) -> None: runner = CliRunner() res = runner.invoke(build_cli(), args=harlequin_args) @@ -75,7 +98,10 @@ def test_no_init_script( "harlequin_args", ["--theme one-dark", ":memory: -t one-dark", "foo.db -t one-dark"] ) def test_theme( - mock_harlequin: MagicMock, mock_adapter: MagicMock, harlequin_args: str + mock_harlequin: MagicMock, + mock_adapter: MagicMock, + harlequin_args: str, + mock_empty_config: None, ) -> None: runner = CliRunner() res = runner.invoke(build_cli(), args=harlequin_args) @@ -96,7 +122,10 @@ def test_theme( ], ) def test_limit( - mock_harlequin: MagicMock, mock_adapter: MagicMock, harlequin_args: str + mock_harlequin: MagicMock, + mock_adapter: MagicMock, + harlequin_args: str, + mock_empty_config: None, ) -> None: runner = CliRunner() res = runner.invoke(build_cli(), args=harlequin_args) @@ -115,7 +144,10 @@ def test_limit( ], ) def test_adapter_opt( - mock_harlequin: MagicMock, mock_adapter: MagicMock, harlequin_args: str + mock_harlequin: MagicMock, + mock_adapter: MagicMock, + harlequin_args: str, + mock_empty_config: None, ) -> None: runner = CliRunner() res = runner.invoke(build_cli(), args=harlequin_args) @@ -133,10 +165,109 @@ def test_adapter_opt( ], ) def test_bad_adapter_opt( - mock_harlequin: MagicMock, mock_adapter: MagicMock, harlequin_args: str + mock_harlequin: MagicMock, + mock_adapter: MagicMock, + harlequin_args: str, + mock_empty_config: None, ) -> None: runner = CliRunner() res = runner.invoke(build_cli(), args=harlequin_args) assert res.exit_code == 2 key_words = ["Error", "Invalid", "-a", "-adapter", "duckdb"] assert all([w in res.stdout for w in key_words]) + + +@pytest.mark.parametrize( + "harlequin_args", + [ + "--profile test-profile", + "-P test-profile", + ], +) +def test_profile_opt( + mock_harlequin: MagicMock, + mock_adapter: MagicMock, + harlequin_args: str, + mock_load_config: Config, +) -> None: + runner = CliRunner() + res = runner.invoke(build_cli(), args=harlequin_args) + assert res.exit_code == 0 + mock_harlequin.assert_called_once() + assert mock_harlequin.call_args + assert mock_harlequin.call_args.kwargs["theme"] == "fruity" + + +@pytest.mark.parametrize( + "harlequin_args", + [ + "--profile test-profile -t zenburn", + "-P test-profile --theme zenburn", + ], +) +def test_profile_override( + mock_harlequin: MagicMock, + mock_adapter: MagicMock, + harlequin_args: str, + mock_load_config: Config, +) -> None: + runner = CliRunner() + res = runner.invoke(build_cli(), args=harlequin_args) + assert res.exit_code == 0 + mock_harlequin.assert_called_once() + assert mock_harlequin.call_args + assert mock_harlequin.call_args.kwargs["theme"] == "zenburn" + + +@pytest.mark.parametrize( + "harlequin_args", + [ + "--profile foo", + "-P bar", + ], +) +def test_bad_profile_opt( + mock_harlequin: MagicMock, + mock_adapter: MagicMock, + harlequin_args: str, + mock_load_config: Config, +) -> None: + runner = CliRunner() + res = runner.invoke(build_cli(), args=harlequin_args) + assert res.exit_code == 2 + key_words = ["profile", "config"] + assert all([w in res.stdout for w in key_words]) + + +@pytest.mark.parametrize("filename", ["good_config.toml", "pyproject.toml"]) +def test_config_path( + mock_harlequin: MagicMock, + mock_adapter: MagicMock, + data_dir: Path, + filename: str, +) -> None: + runner = CliRunner() + config_path = data_dir / "unit_tests" / "config" / filename + res = runner.invoke(build_cli(), args=f"--config-path {config_path.as_posix()}") + assert res.exit_code == 0 + mock_harlequin.assert_called_once() + assert mock_harlequin.call_args + # should use default profile of my-duckdb-profile + assert mock_harlequin.call_args.kwargs["max_results"] == 200_000 + mock_adapter.assert_called_once() + assert mock_adapter.call_args.kwargs["conn_str"] == ["my-database.db"] + assert mock_adapter.call_args.kwargs["read_only"] is False + assert mock_adapter.call_args.kwargs["extension"] == ["httpfs", "spatial"] + + +def test_bad_config_exits( + mock_harlequin: MagicMock, + mock_adapter: MagicMock, + data_dir: Path, +) -> None: + runner = CliRunner() + config_path = data_dir / "unit_tests" / "config" / "default_no_exist.toml" + res = runner.invoke(build_cli(), args=f"--config-path {config_path.as_posix()}") + assert res.exit_code == 2 + key_words = ["default_profile", "foo"] + assert all([w in res.stdout for w in key_words]) diff --git a/tests/unit_tests/test_config.py b/tests/unit_tests/test_config.py new file mode 100644 index 00000000..9e2bdfee --- /dev/null +++ b/tests/unit_tests/test_config.py @@ -0,0 +1,67 @@ +from __future__ import annotations + +from pathlib import Path + +import pytest +from harlequin.config import get_config_for_profile, load_config +from harlequin.exception import HarlequinConfigError + + +@pytest.mark.parametrize("filename", ["good_config.toml", "pyproject.toml"]) +def test_load_config(data_dir: Path, filename: str) -> None: + good_config_path = data_dir / "unit_tests" / "config" / filename + good_config = load_config(config_path=good_config_path) + assert isinstance(good_config, dict) + assert "default_profile" in good_config + assert good_config["default_profile"] == "my-duckdb-profile" + assert "profiles" in good_config + expected_profiles = ["my-duckdb-profile", "local-postgres"] + assert all(name in good_config["profiles"] for name in expected_profiles) + assert all( + isinstance(good_config["profiles"][name], dict) for name in expected_profiles # type: ignore + ) + assert good_config["profiles"]["my-duckdb-profile"]["limit"] == 200_000 # type: ignore + + +@pytest.mark.parametrize("filename", ["good_config.toml", "pyproject.toml"]) +def test_load_named_profile(data_dir: Path, filename: str) -> None: + good_config_path = data_dir / "unit_tests" / "config" / filename + config = get_config_for_profile( + config_path=good_config_path, profile="local-postgres" + ) + assert config["port"] == 5432 + assert config["theme"] == "fruity" + + +@pytest.mark.parametrize("filename", ["good_config.toml", "pyproject.toml"]) +def test_load_default_profile(data_dir: Path, filename: str) -> None: + good_config_path = data_dir / "unit_tests" / "config" / filename + config = get_config_for_profile(config_path=good_config_path, profile=None) + assert config["adapter"] == "duckdb" + assert config["theme"] == "monokai" + + +@pytest.mark.parametrize( + "filename,key_words", + [ + ("default_no_exist.toml", ["default_profile", "foo"]), + ("extra_key.toml", ["unexpected key"]), + ("none_profile.toml", ["None", "not allowed"]), + ("not_toml.toml", ["TOML"]), + ("profiles_not_table.toml", ["profiles", "key", "table"]), + ("profile_not_table.toml", ["members", "profiles", "table"]), + ("bad_option_name.toml", ["option", "invalid", "read-only", "read_only"]), + ], +) +def test_bad_config_raises( + data_dir: Path, + filename: str, + key_words: list[str], +) -> None: + config_path = data_dir / "unit_tests" / "config" / filename + with pytest.raises(HarlequinConfigError) as exc_info: + _ = load_config(config_path=config_path) + err = exc_info.value + assert isinstance(err, HarlequinConfigError) + assert "config" in err.title + assert all([w in err.msg for w in key_words])