datafusion_opensearch/
catalog.rs

1//! A [`SchemaProvider`] that exposes a cluster's indices as tables, so `SHOW TABLES` lists them
2//! and `SELECT … FROM <schema>.<index>` works without registering each index by hand.
3
4use std::collections::HashMap;
5use std::sync::{Arc, RwLock};
6
7use async_trait::async_trait;
8use datafusion::catalog::{SchemaProvider, TableProvider};
9use datafusion::error::{DataFusionError, Result as DataFusionResult};
10
11use crate::client::OpenSearchClient;
12use crate::table::OpenSearchTableFactory;
13use crate::Result;
14
15/// One OpenSearch cluster as a DataFusion schema: every non-system index is a table whose Arrow
16/// schema is derived from its `_mapping` the first time it is referenced (and cached). `SHOW
17/// TABLES` lists them when the session has `information_schema` enabled.
18///
19/// ```no_run
20/// # async fn run() -> Result<(), Box<dyn std::error::Error>> {
21/// use std::sync::Arc;
22/// use datafusion::prelude::SessionContext;
23/// use datafusion_opensearch::{OpenSearchClient, OpenSearchSchemaProvider};
24///
25/// let ctx = SessionContext::new();
26/// let os = OpenSearchSchemaProvider::connect(OpenSearchClient::new("http://localhost:9200")).await?;
27/// ctx.catalog("datafusion").unwrap().register_schema("os", Arc::new(os))?;
28/// ctx.sql("SELECT count(*) FROM os.\"my-index\"").await?.show().await?;
29/// # Ok(()) }
30/// ```
31#[derive(Debug)]
32pub struct OpenSearchSchemaProvider {
33    factory: OpenSearchTableFactory,
34    names: RwLock<Vec<String>>,
35    tables: RwLock<HashMap<String, Arc<dyn TableProvider>>>,
36}
37
38impl OpenSearchSchemaProvider {
39    /// List the cluster's indices now; tables are built lazily on first reference.
40    pub async fn connect(client: OpenSearchClient) -> Result<Self> {
41        let names = client.indices().await?;
42        Ok(Self {
43            factory: OpenSearchTableFactory::new(client),
44            names: RwLock::new(names),
45            tables: RwLock::new(HashMap::new()),
46        })
47    }
48
49    /// Re-list the indices and drop cached tables for indices that no longer exist.
50    pub async fn refresh(&self) -> Result<()> {
51        let names = self.factory.client().indices().await?;
52        self.tables.write().unwrap().retain(|k, _| names.contains(k));
53        *self.names.write().unwrap() = names;
54        Ok(())
55    }
56
57    /// The client the schema talks through.
58    pub fn client(&self) -> &OpenSearchClient {
59        self.factory.client()
60    }
61}
62
63#[async_trait]
64impl SchemaProvider for OpenSearchSchemaProvider {
65    fn table_names(&self) -> Vec<String> {
66        self.names.read().unwrap().clone()
67    }
68
69    async fn table(&self, name: &str) -> DataFusionResult<Option<Arc<dyn TableProvider>>> {
70        if !self.table_exist(name) {
71            return Ok(None);
72        }
73        if let Some(t) = self.tables.read().unwrap().get(name) {
74            return Ok(Some(t.clone()));
75        }
76        let table = self
77            .factory
78            .table_provider(name)
79            .await
80            .map_err(|e| DataFusionError::External(Box::new(e)))?;
81        self.tables.write().unwrap().insert(name.to_string(), table.clone());
82        Ok(Some(table))
83    }
84
85    fn table_exist(&self, name: &str) -> bool {
86        self.names.read().unwrap().iter().any(|n| n == name)
87    }
88}
datafusion_opensearch/catalog.rs

datafusion_opensearch/
catalog.rs
