Skip to content

Commit

Permalink
feat: add a duckdb handle with the project filesystem registered unde…
Browse files Browse the repository at this point in the history
…r cdf protocol
  • Loading branch information
z3z1ma committed Jun 8, 2024
1 parent 8b0982b commit 3342c1a
Show file tree
Hide file tree
Showing 2 changed files with 22 additions and 1 deletion.
3 changes: 2 additions & 1 deletion src/cdf/core/filesystem.py
Original file line number Diff line number Diff line change
Expand Up @@ -52,7 +52,8 @@ def wrapped(self) -> fsspec.AbstractFileSystem:

options = self.options.copy()
options.setdefault("auto_mkdir", True)
return DirFileSystem(
CdfFs = type("CdfFs", (DirFileSystem,), {"protocol": "cdf"})
return CdfFs(
path=posixpath.join(strip_protocol(self.uri), "x")[:-1],
fs=fsspec.filesystem(self.protocol, **options),
auto_mkdir=True,
Expand Down
20 changes: 20 additions & 0 deletions src/cdf/core/project.py
Original file line number Diff line number Diff line change
Expand Up @@ -48,6 +48,7 @@
from functools import cached_property
from pathlib import Path

import duckdb
import dynaconf
import pydantic
from dynaconf.utils.boxing import DynaBox
Expand Down Expand Up @@ -753,6 +754,25 @@ def feature_flags(self) -> FeatureFlagAdapter:
self.ff_settings, filesystem=self.filesystem
).unwrap()

@cached_property
def duckdb(self) -> duckdb.DuckDBPyConnection:
"""Get a handle to the project's DuckDB connection"""
conn = duckdb.connect(":memory:")
conn.install_extension("httpfs")
conn.install_extension("json")
conn.register_filesystem(self.filesystem.wrapped)
conn.execute("CREATE TABLE workspaces (name TEXT PRIMARY KEY, path TEXT)")
for workspace in self.workspaces:
conn.execute(
"INSERT INTO workspaces (name, path) VALUES (?, ?)",
(workspace.name, workspace.path.as_posix()),
)
return conn

def get_workspace_path(self, name: str) -> M.Result[Path, Exception]:
"""Get the path to a workspace by name"""
return self.get_workspace(name).map(lambda ws: ws.path)

@classmethod
def from_path(cls, root: PathLike):
"""Load configuration data from a project root path using dynaconf.
Expand Down

0 comments on commit 3342c1a

Please sign in to comment.