Repository: ibis-project/ibis

Issue:
feat: user-facing method for installing duckdb extensions w/o SQL
### Is your feature request related to a problem?

If I want to use DuckDB with an extension such as httpfs, I have to drop into SQL to install and load it. Due to this, I might run into a situation where I set up a connection that attempts to configure behavior provided by a non-installed, non-loaded extension,

```python
c = ibis.duckdb.connect("database.ddb", s3_region="us-east-2")
```

This returns successfully. But if I run

```python
c.raw_sql("SELECT current_setting('s3_region')")
```

I get:

> (duckdb.CatalogException) Catalog Error: Setting with name "s3_region" is not in the catalog, but it exists in the httpfs extension.
> 
> To install and load the extension, run:
> INSTALL httpfs;
> LOAD httpfs;

This isn't hard to work around if you know about DuckDB and `raw_sql` but it might trip users up and could overall be improved.

### Describe the solution you'd like

It would be nice if I could specify extensions to load with the `ibis.duckdb.connect` call, similar to how I can specify configuration options as kwargs. It might also be nice to handle the above error message in a more useful way though that may be out of scope of this feature request.

### What version of ibis are you running?

6.2.0

### What backend(s) are you using, if any?

DuckDB

### Code of Conduct

- [X] I agree to follow this project's Code of Conduct


Discussion:
Thanks for opening this. I'd propose the following interface for this:

- A new optional `extensions` kwarg to `ibis.duckdb.connect` (handed by `ibis.backends.duckdb.Backend.do_connect`) for loading extensions immediately upon connection. This would be used as `ibis.duckdb.connect(extensions=["httpfs"])`.
- A user-facing `load_extension` method on the duckdb backend connection. This would both install and load the extension as necessary (the duckdb library has separate `install_extension` and `load_extension` methods for this). We already have much of the code for doing this in a private method, just need to make it public. This would let the user write `con.load_extension("httpfs")`.
Thanks for the quick reply @jcrist, I think both of those proposed ideas would be useful here.

Resolving patch:
diff --git a/ibis/backends/duckdb/__init__.py b/ibis/backends/duckdb/__init__.py
--- a/ibis/backends/duckdb/__init__.py
+++ b/ibis/backends/duckdb/__init__.py
@@ -6,7 +6,6 @@
 import contextlib
 import os
 import warnings
-from functools import partial
 from pathlib import Path
 from typing import (
     TYPE_CHECKING,
@@ -34,7 +33,7 @@
 from ibis.formats.pandas import PandasData
 
 if TYPE_CHECKING:
-    from collections.abc import Iterable, Iterator, Mapping, MutableMapping
+    from collections.abc import Iterable, Iterator, Mapping, MutableMapping, Sequence
 
     import pandas as pd
     import torch
@@ -173,6 +172,7 @@ def do_connect(
         database: str | Path = ":memory:",
         read_only: bool = False,
         temp_directory: str | Path | None = None,
+        extensions: Sequence[str] | None = None,
         **config: Any,
     ) -> None:
         """Create an Ibis client connected to a DuckDB database.
@@ -186,6 +186,8 @@ def do_connect(
         temp_directory
             Directory to use for spilling to disk. Only set by default for
             in-memory connections.
+        extensions
+            A list of duckdb extensions to install/load upon connection.
         config
             DuckDB configuration parameters. See the [DuckDB configuration
             documentation](https://duckdb.org/docs/sql/configuration) for
@@ -222,6 +224,8 @@ def do_connect(
 
         @sa.event.listens_for(engine, "connect")
         def configure_connection(dbapi_connection, connection_record):
+            if extensions is not None:
+                self._sa_load_extensions(dbapi_connection, extensions)
             dbapi_connection.execute("SET TimeZone = 'UTC'")
             # the progress bar in duckdb <0.8.0 causes kernel crashes in
             # jupyterlab, fixed in https://github.com/duckdb/duckdb/pull/6831
@@ -237,31 +241,36 @@ def configure_connection(dbapi_connection, connection_record):
 
         super().do_connect(engine)
 
-    def _load_extensions(self, extensions):
-        extension_name = sa.column("extension_name")
-        loaded = sa.column("loaded")
-        installed = sa.column("installed")
-        aliases = sa.column("aliases")
-        query = (
-            sa.select(extension_name)
-            .select_from(sa.func.duckdb_extensions())
-            .where(
-                sa.and_(
-                    # extension isn't loaded or isn't installed
-                    sa.not_(loaded & installed),
-                    # extension is one that we're requesting, or an alias of it
-                    sa.or_(
-                        extension_name.in_(extensions),
-                        *map(partial(sa.func.array_has, aliases), extensions),
-                    ),
-                )
-            )
+    @staticmethod
+    def _sa_load_extensions(dbapi_con, extensions):
+        query = """
+        WITH exts AS (
+          SELECT extension_name AS name, aliases FROM duckdb_extensions()
+          WHERE installed AND loaded
         )
+        SELECT name FROM exts
+        UNION (SELECT UNNEST(aliases) AS name FROM exts)
+        """
+        installed = (name for (name,) in dbapi_con.sql(query).fetchall())
+        # Install and load all other extensions
+        todo = set(extensions).difference(installed)
+        for extension in todo:
+            dbapi_con.install_extension(extension)
+            dbapi_con.load_extension(extension)
+
+    def _load_extensions(self, extensions):
         with self.begin() as con:
-            c = con.connection
-            for extension in con.execute(query).scalars():
-                c.install_extension(extension)
-                c.load_extension(extension)
+            self._sa_load_extensions(con.connection, extensions)
+
+    def load_extension(self, extension: str) -> None:
+        """Install and load a duckdb extension by name or path.
+
+        Parameters
+        ----------
+        extension
+            The extension name or path.
+        """
+        self._load_extensions([extension])
 
     def create_schema(
         self, name: str, database: str | None = None, force: bool = False


Test patch:
diff --git a/ibis/backends/duckdb/tests/test_client.py b/ibis/backends/duckdb/tests/test_client.py
new file mode 100644
--- /dev/null
+++ b/ibis/backends/duckdb/tests/test_client.py
@@ -0,0 +1,42 @@
+from __future__ import annotations
+
+import duckdb
+import pytest
+import sqlalchemy as sa
+
+import ibis
+from ibis.conftest import LINUX, SANDBOXED
+
+
+@pytest.mark.xfail(
+    LINUX and SANDBOXED,
+    reason="nix on linux cannot download duckdb extensions or data due to sandboxing",
+    raises=sa.exc.OperationalError,
+)
+def test_connect_extensions():
+    con = ibis.duckdb.connect(extensions=["s3", "sqlite"])
+    results = con.raw_sql(
+        """
+        SELECT loaded FROM duckdb_extensions()
+        WHERE extension_name = 'httpfs' OR extension_name = 'sqlite'
+        """
+    ).fetchall()
+    assert all(loaded for (loaded,) in results)
+
+
+@pytest.mark.xfail(
+    LINUX and SANDBOXED,
+    reason="nix on linux cannot download duckdb extensions or data due to sandboxing",
+    raises=duckdb.IOException,
+)
+def test_load_extension():
+    con = ibis.duckdb.connect()
+    con.load_extension("s3")
+    con.load_extension("sqlite")
+    results = con.raw_sql(
+        """
+        SELECT loaded FROM duckdb_extensions()
+        WHERE extension_name = 'httpfs' OR extension_name = 'sqlite'
+        """
+    ).fetchall()
+    assert all(loaded for (loaded,) in results)