Skip to content

Commit

Permalink
refactor(duckdb): source loaded and installed extensions from duckdb
Browse files Browse the repository at this point in the history
  • Loading branch information
cpcloud authored and gforsyth committed Mar 23, 2023
1 parent e9f57eb commit fb06262
Showing 1 changed file with 25 additions and 8 deletions.
33 changes: 25 additions & 8 deletions ibis/backends/duckdb/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -6,6 +6,7 @@
import itertools
import os
import warnings
from functools import partial
from pathlib import Path
from typing import TYPE_CHECKING, Any, Iterable, Iterator, Mapping, MutableMapping

Expand Down Expand Up @@ -143,16 +144,32 @@ def configure_connection(dbapi_connection, connection_record):
super().do_connect(engine)

self._meta = sa.MetaData()
self._extensions = set()

def _load_extensions(self, extensions):
for extension in extensions:
if extension not in self._extensions:
with self.begin() as con:
c = con.connection
c.install_extension(extension)
c.load_extension(extension)
self._extensions.add(extension)
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),
),
)
)
)
with self.begin() as con:
c = con.connection
for extension in con.execute(query).scalars():
c.install_extension(extension)
c.load_extension(extension)

def register(
self,
Expand Down

0 comments on commit fb06262

Please sign in to comment.