feat: add extension options, support prql (#167)

* feat: add extension options, support prql

* fix: xfail prql test for windows

* fix: don't xfail md test for windows

CHANGELOG.md

@@ -3,6 +3,16 @@
 All notable changes to this project will be documented in this file.
 
 ## [Unreleased]
+-   Adds a new CLI option, `--extension` or `-e`, which will install and load a named DuckDB extension.
+-   Adds a new CLI option, `--force-install-extensions`, which will re-install the extensions provided
+    with the `-e` option.
+-   Adds a new CLI option, `--custom-extension-repo`, which enables installing extensions other than
+    the official DuckDB extensions.
+-   Taken together, Harlequin can now be loaded with the [PRQL](https://github.com/ywelsch/duckdb-prql) extension. Use PRQL with Harlequin:
+    ```bash
+    harlequin -u -e prql --custom-extension-repo welsch.lu/duckdb/prql/latest
+    ```
+    ([#152](https://github.com/tconbeer/harlequin/issues/152) - thank you [@dljsjr](https://github.com/dljsjr)!)
 
 ## [0.0.23] - 2023-08-03
 

src/harlequin/cli.py

@@ -21,6 +21,27 @@
     is_flag=True,
     help="Allow loading unsigned extensions",
 )
+@click.option(
+    "-e",
+    "--extension",
+    multiple=True,
+    help=(
+        "Install and load the named DuckDB extension when starting "
+        "Harlequin. To install multiple extensions, repeat this option."
+    ),
+)
+@click.option(
+    "--force-install-extensions",
+    is_flag=True,
+    help="Force install all extensions passed with -e.",
+)
+@click.option(
+    "--custom-extension-repo",
+    help=(
+        "A value to pass to DuckDB's custom_extension_repository variable. "
+        "Will be set before loading any extensions that are passing using -e."
+    ),
+)
 @click.option(
     "-t",
     "--theme",
@@ -52,16 +73,22 @@
 def harlequin(
     db_path: List[str],
     read_only: bool,
+    allow_unsigned_extensions: bool,
+    extension: List[str],
+    force_install_extensions: bool,
+    custom_extension_repo: Union[str, None],
     theme: str,
     md_token: Union[str, None],
     md_saas: bool,
-    allow_unsigned_extensions: bool,
 ) -> None:
     if not db_path:
         db_path = [":memory:"]
     tui = Harlequin(
         db_path=db_path,
         read_only=read_only,
+        extensions=extension,
+        force_install_extensions=force_install_extensions,
+        custom_extension_repo=custom_extension_repo,
         theme=theme,
         md_token=md_token,
         md_saas=md_saas,

src/harlequin/duck_ops.py

@@ -2,6 +2,8 @@
 from typing import List, Sequence, Tuple, Union
 
 import duckdb
+from rich import print
+from rich.panel import Panel
 
 from harlequin.exception import HarlequinExit
 
@@ -14,9 +16,12 @@
 def connect(
     db_path: Sequence[Union[str, Path]],
     read_only: bool = False,
+    allow_unsigned_extensions: bool = False,
+    extensions: Union[List[str], None] = None,
+    force_install_extensions: bool = False,
+    custom_extension_repo: Union[str, None] = None,
     md_token: Union[str, None] = None,
     md_saas: bool = False,
-    allow_unsigned_extensions: bool = False,
 ) -> duckdb.DuckDBPyConnection:
     if not db_path:
         db_path = [":memory:"]
@@ -32,9 +37,6 @@ def connect(
         for db in other_dbs:
             connection.execute(f"attach '{db}'{' (READ_ONLY)' if read_only else ''}")
     except (duckdb.CatalogException, duckdb.IOException) as e:
-        from rich import print
-        from rich.panel import Panel
-
         print(
             Panel.fit(
                 str(e),
@@ -45,10 +47,35 @@ def connect(
                 subtitle_align="right",
             )
         )
-
         raise HarlequinExit() from None
-    else:
-        return connection
+
+    if custom_extension_repo:
+        connection.execute(
+            f"SET custom_extension_repository='{custom_extension_repo}';"
+        )
+
+    if extensions:
+        try:
+            for extension in extensions:
+                # todo: support installing from a URL instead.
+                connection.install_extension(
+                    extension=extension, force_install=force_install_extensions
+                )
+                connection.load_extension(extension=extension)
+        except (duckdb.HTTPException, duckdb.IOException) as e:
+            print(
+                Panel.fit(
+                    str(e),
+                    title="DuckDB couldn't install or load your extension.",
+                    title_align="left",
+                    border_style="red",
+                    subtitle="Try again?",
+                    subtitle_align="right",
+                )
+            )
+            raise HarlequinExit() from None
+
+    return connection
 
 
 def get_databases(conn: duckdb.DuckDBPyConnection) -> List[Tuple[str]]:

src/harlequin/tui/app.py

@@ -65,13 +65,16 @@ def __init__(
         self,
         db_path: Sequence[Union[str, Path]],
         read_only: bool = False,
+        allow_unsigned_extensions: bool = False,
+        extensions: Union[List[str], None] = None,
+        force_install_extensions: bool = False,
+        custom_extension_repo: Union[str, None] = None,
         theme: str = "monokai",
         md_token: Union[str, None] = None,
         md_saas: bool = False,
         driver_class: Union[Type[Driver], None] = None,
         css_path: Union[CSSPathType, None] = None,
         watch_css: bool = False,
-        allow_unsigned_extensions: bool = False,
     ):
         super().__init__(driver_class, css_path, watch_css)
         self.theme = theme
@@ -80,9 +83,12 @@ def __init__(
             self.connection = connect(
                 db_path,
                 read_only=read_only,
+                allow_unsigned_extensions=allow_unsigned_extensions,
+                extensions=extensions,
+                force_install_extensions=force_install_extensions,
+                custom_extension_repo=custom_extension_repo,
                 md_token=md_token,
                 md_saas=md_saas,
-                allow_unsigned_extensions=allow_unsigned_extensions,
             )
         except HarlequinExit:
             self.exit()

tests/functional_tests/test_duck_ops.py

@@ -23,12 +23,32 @@ def test_connect(tiny_db: Path, small_db: Path, tmp_path: Path) -> None:
     assert connect([tmp_path / "new.db"])
     assert connect([], allow_unsigned_extensions=True)
     assert connect([tiny_db], allow_unsigned_extensions=True)
-    assert connect([tiny_db, small_db], read_only=True, allow_unsigned_extensions=True)
+    assert connect([tiny_db, small_db], read_only=True)
+
+
+def test_connect_extensions() -> None:
+    assert connect([], extensions=None)
+    assert connect([], extensions=[])
+    assert connect([], extensions=["spatial"])
+    assert connect([], allow_unsigned_extensions=True, extensions=["spatial"])
 
 
 @pytest.mark.xfail(
-    sys.platform == "win32", reason="MotherDuck extension not yet built for Windows."
+    sys.platform == "win32",
+    reason="PRQL extension not yet built for Windows and DuckDB v0.8.1.",
 )
+def test_connect_prql() -> None:
+    # Note: this may fail in the future if the extension doesn't support the latest
+    # duckdb version.
+    assert connect(
+        [],
+        allow_unsigned_extensions=True,
+        extensions=["prql"],
+        custom_extension_repo="welsch.lu/duckdb/prql/latest",
+        force_install_extensions=True,
+    )
+
+
 @pytest.mark.skipif(
     sys.version_info[0:2] != (3, 10), reason="Matrix is hitting MD too many times."
 )