feat: add duckdb client (#15)

Author: gadomski

CHANGELOG.md

@@ -6,6 +6,10 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
 
 ## [Unreleased]
 
+### Added
+
+- DuckDB client ([#15](https://github.com/gadomski/stacrs/pull/15))
+
 ## [0.3.0] - 2024-11-21
 
 ### Removed

data/100-sentinel-2-items.parquet

@@ -0,0 +1,60 @@
+use crate::{Error, Result};
+use pyo3::{exceptions::PyException, prelude::*, types::PyDict};
+use stac_api::python::{StringOrDict, StringOrList};
+use stac_duckdb::Client;
+use std::sync::Mutex;
+
+#[pyclass(frozen)]
+pub struct DuckdbClient(Mutex<Client>);
+
+#[pymethods]
+impl DuckdbClient {
+    #[new]
+    fn new() -> Result<DuckdbClient> {
+        let client = Client::new()?;
+        Ok(DuckdbClient(Mutex::new(client)))
+    }
+
+    #[pyo3(signature = (href, *, intersects=None, ids=None, collections=None, limit=None, bbox=None, datetime=None, include=None, exclude=None, sortby=None, filter=None, query=None, **kwargs))]
+    fn search<'py>(
+        &self,
+        py: Python<'py>,
+        href: String,
+        intersects: Option<StringOrDict>,
+        ids: Option<StringOrList>,
+        collections: Option<StringOrList>,
+        limit: Option<u64>,
+        bbox: Option<Vec<f64>>,
+        datetime: Option<String>,
+        include: Option<StringOrList>,
+        exclude: Option<StringOrList>,
+        sortby: Option<StringOrList>,
+        filter: Option<StringOrDict>,
+        query: Option<Bound<'py, PyDict>>,
+        kwargs: Option<Bound<'py, PyDict>>,
+    ) -> PyResult<Bound<'py, PyDict>> {
+        let search = stac_api::python::search(
+            intersects,
+            ids,
+            collections,
+            limit,
+            bbox,
+            datetime,
+            include,
+            exclude,
+            sortby,
+            filter,
+            query,
+            kwargs,
+        )?;
+        let item_collection = {
+            let client = self
+                .0
+                .lock()
+                .map_err(|err| PyException::new_err(err.to_string()))?;
+            client.search(&href, search).map_err(Error::from)?
+        };
+        let dict = pythonize::pythonize(py, &item_collection)?;
+        dict.extract()
+    }
+}

src/duckdb.rs

@@ -1,13 +1,14 @@
 #![deny(unused_crate_dependencies, warnings)]
 
+mod duckdb;
 mod error;
 mod migrate;
 mod read;
 mod search;
 mod version;
 mod write;
 
-use duckdb as _;
+use ::duckdb as _;
 use error::Error;
 use pyo3::prelude::*;
 
@@ -16,6 +17,7 @@ type Result<T> = std::result::Result<T, Error>;
 /// A collection of functions for working with STAC, using Rust under the hood.
 #[pymodule]
 fn stacrs(m: &Bound<'_, PyModule>) -> PyResult<()> {
+    m.add_class::<duckdb::DuckdbClient>()?;
     m.add_function(wrap_pyfunction!(migrate::migrate, m)?)?;
     m.add_function(wrap_pyfunction!(migrate::migrate_href, m)?)?;
     m.add_function(wrap_pyfunction!(read::read, m)?)?;

src/lib.rs

@@ -1,5 +1,27 @@
 from typing import Any, Optional, Tuple
 
+class DuckdbClient:
+    """A client for querying stac-geoparquet with DuckDB."""
+
+    def search(
+        href: str,
+        *,
+        intersects: Optional[str | dict[str, Any]] = None,
+        ids: Optional[str | list[str]] = None,
+        collections: Optional[str | list[str]] = None,
+        max_items: Optional[int] = None,
+        limit: Optional[int] = None,
+        bbox: Optional[list[float]] = None,
+        datetime: Optional[str] = None,
+        include: Optional[str | list[str]] = None,
+        exclude: Optional[str | list[str]] = None,
+        sortby: Optional[str | list[str]] = None,
+        filter: Optional[str | dict[str, Any]] = None,
+        query: Optional[dict[str, Any]] = None,
+        **kwargs: str,
+    ):
+        """Search a stac-geoparquet file with duckdb"""
+
 def migrate_href(href: str, version: Optional[str] = None) -> dict[str, Any]:
     """
     Migrates a STAC dictionary at the given href to another version.

stacrs.pyi

@@ -0,0 +1,30 @@
+import pytest
+from stacrs import DuckdbClient
+
+
+@pytest.fixture
+def client() -> DuckdbClient:
+    return DuckdbClient()
+
+
+def test_search(client: DuckdbClient) -> None:
+    item_collection = client.search("data/extended-item.parquet")
+    assert len(item_collection["features"]) == 1
+
+
+def test_search_offset(client: DuckdbClient) -> None:
+    item_collection = client.search(
+        "data/100-sentinel-2-items.parquet", offset=0, limit=1
+    )
+    assert (
+        item_collection["features"][0]["id"]
+        == "S2B_MSIL2A_20241203T174629_R098_T13TDE_20241203T211406"
+    )
+
+    item_collection = client.search(
+        "data/100-sentinel-2-items.parquet", offset=1, limit=1
+    )
+    assert (
+        item_collection["features"][0]["id"]
+        == "S2A_MSIL2A_20241201T175721_R141_T13TDE_20241201T213150"
+    )