From ff84517a52c0351c4b68cc4789a5d83c6f24249b Mon Sep 17 00:00:00 2001
From: program-- <justin@justinsingh.me>
Date: Fri, 6 Dec 2024 21:15:00 -0800
Subject: [PATCH] feat: add optional DuckDB query source

---
 DESCRIPTION                |  4 +++-
 NAMESPACE                  |  1 +
 R/query_source_duckdb.R    | 47 ++++++++++++++++++++++++++++++++++++++
 R/zzz.R                    |  3 ++-
 man/hfsubsetR-package.Rd   |  2 ++
 man/query_source_duckdb.Rd | 23 +++++++++++++++++++
 6 files changed, 78 insertions(+), 2 deletions(-)
 create mode 100644 R/query_source_duckdb.R
 create mode 100644 man/query_source_duckdb.Rd

diff --git a/DESCRIPTION b/DESCRIPTION
index a6a6850..40beba2 100644
--- a/DESCRIPTION
+++ b/DESCRIPTION
@@ -27,7 +27,8 @@ Imports:
     sf,
     methods
 Suggests:
-    testthat
+    testthat,
+    duckdb
 License: GPL (>= 3)
 Encoding: UTF-8
 LazyData: true
@@ -44,6 +45,7 @@ Collate:
     'query.R'
     'query_source.R'
     'query_source_arrow.R'
+    'query_source_duckdb.R'
     'query_source_sf.R'
     'query_subset.R'
     'sf_arrow.R'
diff --git a/NAMESPACE b/NAMESPACE
index 70471a8..01e0621 100644
--- a/NAMESPACE
+++ b/NAMESPACE
@@ -17,6 +17,7 @@ export(query_set_layers)
 export(query_set_sink)
 export(query_set_source)
 export(query_source_arrow)
+export(query_source_duckdb)
 export(query_source_sf)
 export(query_subset)
 export(read_sf_dataset)
diff --git a/R/query_source_duckdb.R b/R/query_source_duckdb.R
new file mode 100644
index 0000000..7c68c0b
--- /dev/null
+++ b/R/query_source_duckdb.R
@@ -0,0 +1,47 @@
+#' Create a new DuckDB query source
+#' @param srcname Path or URI to source
+#' @param ... Passed to [DBI::dbConnect].
+#' @param .attach If `TRUE`, then `ATTACH <srcname>` is executed on the DuckDB connection
+#' @param .extensions Optional extensions to install and load on connection
+#' @returns An `hfsubset_query_source_duckdb` object
+#' @export
+query_source_duckdb <- function(srcname, ..., .attach = FALSE, .extensions = c("spatial")) {
+  if (!requireNamespace("duckdb", quietly = TRUE)) {
+    stop("Package `duckdb` is required for query_source_duckdb()")
+  }
+
+  conn <- DBI::dbConnect(duckdb::duckdb(), ...)
+
+  for (ext in .extensions) {
+    DBI::dbExecute(conn, paste("INSTALL", ext))
+    DBI::dbExecute(conn, paste("LOAD", ext))
+  }
+  
+  if (.attach) {
+    DBI::dbExecute(
+      conn,
+      paste0("ATTACH '", srcname, "' AS hydrofabric (TYPE SQLITE)")
+    )
+  }
+
+  structure(
+    list(src = srcname, conn = conn),
+    class = c("hfsubset_query_source_duckdb", "hfsubset_query_source")
+  )
+}
+
+#' @method query_source_layer hfsubset_query_source_duckdb
+#' @keywords internal
+query_source_layer.hfsubset_query_source_duckdb <- function(query_source, layer, ..., query = NULL) {
+  if (missing(layer) && !is.null(query)) {
+    dplyr::tbl(query_source$conn, dbplyr::sql(query), ...)    
+  } else {
+    dplyr::tbl(query_source$conn, paste0("hydrofabric.", layer), ...)
+  }
+}
+
+#' @method query_source_layers hfsubset_query_source_duckdb
+#' @keywords internal
+query_source_layers.hfsubset_query_source_duckdb <- function(query_source, ...) {
+  DBI::dbListTables(query_source$conn)
+}
diff --git a/R/zzz.R b/R/zzz.R
index ae441e9..aa489d7 100644
--- a/R/zzz.R
+++ b/R/zzz.R
@@ -16,7 +16,8 @@
 
   query_source_classes <- c(
     "hfsubset_query_source_arrow",
-    "hfsubset_query_source_sf"
+    "hfsubset_query_source_sf",
+    "hfsubset_query_source_duckdb"
   )
 
   for (query_source_class in query_source_classes) {
diff --git a/man/hfsubsetR-package.Rd b/man/hfsubsetR-package.Rd
index 5cc27fd..97e6cca 100644
--- a/man/hfsubsetR-package.Rd
+++ b/man/hfsubsetR-package.Rd
@@ -6,6 +6,8 @@
 \alias{hfsubsetR-package}
 \title{hfsubsetR: Hydrofabric Subsetter}
 \description{
+\if{html}{\figure{logo.png}{options: style='float: right' alt='logo' width='120'}}
+
 Subset Hydrofabric Data in R.
 }
 \seealso{
diff --git a/man/query_source_duckdb.Rd b/man/query_source_duckdb.Rd
new file mode 100644
index 0000000..bd941f5
--- /dev/null
+++ b/man/query_source_duckdb.Rd
@@ -0,0 +1,23 @@
+% Generated by roxygen2: do not edit by hand
+% Please edit documentation in R/query_source_duckdb.R
+\name{query_source_duckdb}
+\alias{query_source_duckdb}
+\title{Create a new DuckDB query source}
+\usage{
+query_source_duckdb(srcname, ..., .attach = FALSE, .extensions = c("spatial"))
+}
+\arguments{
+\item{srcname}{Path or URI to source}
+
+\item{...}{Passed to [DBI::dbConnect].}
+
+\item{.attach}{If `TRUE`, then `ATTACH <srcname>` is executed on the DuckDB connection}
+
+\item{.extensions}{Optional extensions to install and load on connection}
+}
+\value{
+An `hfsubset_query_source_duckdb` object
+}
+\description{
+Create a new DuckDB query source
+}