bug: Unable to load duckdb extensions from a custom restricted environment

[vspinu]

What happened?

Duckdb extensions cannot be loaded on the fly in our restricted environments (our own pypi, firewalls etc). Instead, we can install the relevant packages and load it ourselves.

This is the pure duckdb code which works on our hdfs+delta environment:

import duckdb
import ibis

from fsspec import filesystem
from duckdb_extensions import extension_importer

duckdb.register_filesystem(filesystem("hdfs"))
extension_importer.import_extension("delta")
duckdb.sql("LOAD 'delta';")

duckdb.sql("select count(*) from delta_scan('hdfs://path/to/tbl/')" )

How can I make the above work with ibis?

con = ibis.duckdb.connect()

# these two fail to discover the installed extensions and attempt to auto-install the extension 
# after a long waiting delay they fail with 
# IOException: IO Error: Failed to download extension "delta" at URL "http://extensions.duckdb.org/v1.2.2/linux_amd64_gcc4/delta.duckdb_extension.gz"
 
ibis.read_delta(...)
con.read_delta(...)


# this fails with ParserException: Parser Error: syntax error at or near "'delta'"
con.sql("LOAD delta")

What version of ibis are you using?

10.5.0

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

duckdb

Relevant log output

Code of Conduct

• I agree to follow this project's Code of Conduct

[gforsyth]

Hey @vspinu -- instead of using .sql, which needs to return a table with a schema for us to parse, try raw_sql which just executes directly on the backend.

[vspinu]

Thanks @gforsyth! This solves the parsing error indeed. But is still stuck trying to fetch non-existing resource for me.

The problem is that local extension regestration happens only for the default duckdb connection which is not the one used by ibis in a typical case.

This is what I ended up with:

def import_extension(name: str, con: duckdb.DuckDBPyConnection | None = None):
    import importlib
    import pathlib

    if con is None:
        con = duckdb.default_connection()

    quack_module = importlib.import_module(f"duckdb_extension_{name}")
    module_path = pathlib.Path(str(importlib.resources.files(quack_module)))

    duckdb_version = con.sql("PRAGMA version;").fetchone()[0]

    extension_dir = module_path / "extensions" / duckdb_version
    extension_file = extension_dir / f"{name}.duckdb_extension"

    con.sql(f"INSTALL '{extension_file}'")

duck = ibis.duckdb.connect()
import_extension("delta", duck.con)
duck.con.register_filesystem(filesystem("hdfs"))

duck.con.sql(f"select count(*) from delta_scan('{hdfs_path}')")    # works!
duck.read_delta(hdfs_path).count()   # works!

[vspinu]

Made a PR to duckdb_extensions to support non-default connections.

[cpcloud]

Can this be closed?