Skip to content

Commit 5bd7bed

Browse files
bakeybakey
authored
feat(dynamodb): add catalog mode for whole-account table discovery (#152)
* feat(dynamodb): add catalog mode for whole-account table discovery Register all accessible DynamoDB tables as a named DataFusion catalog so users can query them with three-part SQL references like myddb.tables.products, matching the catalog experience already available for SQLite/Postgres/MySQL/SeekDB. - List tables via ListTables with pagination. - Infer each table's key schema via DescribeTable. - Register discovered tables under the fixed schema tables. - Update server config to pass hierarchy_level and include DynamoDB in CATALOG_SUPPORTED_SOURCES. - Add unit tests for table/catalog option validation. * Fix implementation issue * fix integration test * add unit test cases * Fix by review
1 parent 4c15923 commit 5bd7bed

7 files changed

Lines changed: 1116 additions & 81 deletions

File tree

‎README.md‎

Lines changed: 17 additions & 17 deletions
Original file line numberDiff line numberDiff line change
@@ -313,23 +313,23 @@ For end-to-end walkthroughs — RAG, recommendations, an agent-native wiki, a si
313313

314314
## Supported Data Sources
315315

316-
| Type | CRUD | Description | Docs |
317-
|------|------|-------------|------|
318-
| PostgreSQL | Full | Table or catalog registration, pgvector KNN | [docs/postgres/](docs/postgres/) |
319-
| MySQL | Full | Table or catalog registration | [docs/mysql/](docs/mysql/) |
320-
| SQLite | Full | Table or catalog registration, sqlite-vec KNN, FTS | [docs/sqlite/](docs/sqlite/) |
321-
| MongoDB | Full | Collections with point lookups | [docs/mongo/](docs/mongo/) |
322-
| Redis | Full | Hashes mapped to SQL rows | [docs/redis/](docs/redis/) |
323-
| DynamoDB | Full | Items mapped to SQL rows, scan + filter pushdown | [docs/dynamodb/](docs/dynamodb/) |
324-
| SeekDB | Full | MySQL-wire CRUD, native FULLTEXT FTS, HNSW VECTOR KNN | [docs/seekdb/](docs/seekdb/) |
325-
| Lance | Read (job-write) | KNN vector search, BM25 FTS; job destination | [docs/lance/](docs/lance/) |
326-
| CSV | Read | Local or remote CSV files | [docs/server.md](docs/server.md) |
327-
| Parquet | Read | Local or remote Parquet files | [docs/server.md](docs/server.md) |
328-
| JSON / NDJSON | Read | Local or remote JSON files | [docs/cli.md](docs/cli.md) |
329-
| S3 / GCS / Azure | Read | CSV, Parquet, Lance from object stores | [docs/S3_USAGE.md](docs/S3_USAGE.md) |
330-
| Apache Iceberg | Read | Schema evolution, partition pruning | [docs/iceberg/](docs/iceberg/) |
331-
| InfluxDB 3 | Read | Time-series measurements over Arrow Flight SQL | [docs/influxdb/](docs/influxdb/) |
332-
| Documents | Read | PDF/Office/ODF/image → per-page markdown, tables, images (local directories; `documents` feature) | [docs/documents.md](docs/documents.md) |
316+
| Type | CRUD | Catalog mode | Description | Docs |
317+
|------|------|--------------|-------------|------|
318+
| PostgreSQL | Full | Yes | Table or catalog registration, pgvector KNN | [docs/postgres/](docs/postgres/) |
319+
| MySQL | Full | Yes | Table or catalog registration | [docs/mysql/](docs/mysql/) |
320+
| SQLite | Full | Yes | Table or catalog registration, sqlite-vec KNN, FTS | [docs/sqlite/](docs/sqlite/) |
321+
| MongoDB | Full | No | Collections with point lookups | [docs/mongo/](docs/mongo/) |
322+
| Redis | Full | No | Hashes mapped to SQL rows | [docs/redis/](docs/redis/) |
323+
| DynamoDB | Full | Yes | Items mapped to SQL rows, table or catalog registration, scan + filter pushdown | [docs/dynamodb/](docs/dynamodb/) |
324+
| SeekDB | Full | Yes | MySQL-wire CRUD, native FULLTEXT FTS, HNSW VECTOR KNN | [docs/seekdb/](docs/seekdb/) |
325+
| Lance | Read (job-write) | No | KNN vector search, BM25 FTS; job destination | [docs/lance/](docs/lance/) |
326+
| CSV | Read | No | Local or remote CSV files | [docs/server.md](docs/server.md) |
327+
| Parquet | Read | No | Local or remote Parquet files | [docs/server.md](docs/server.md) |
328+
| JSON / NDJSON | Read | No | Local or remote JSON files | [docs/cli.md](docs/cli.md) |
329+
| S3 / GCS / Azure | Read | No | CSV, Parquet, Lance from object stores | [docs/S3_USAGE.md](docs/S3_USAGE.md) |
330+
| Apache Iceberg | Read | No | Schema evolution, partition pruning | [docs/iceberg/](docs/iceberg/) |
331+
| InfluxDB 3 | Read | No | Time-series measurements over Arrow Flight SQL | [docs/influxdb/](docs/influxdb/) |
332+
| Documents | Read | No | PDF/Office/ODF/image -> per-page markdown, tables, images (local directories; `documents` feature) | [docs/documents.md](docs/documents.md) |
333333

334334
---
335335

‎crates/cli/src/main.rs‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -928,6 +928,7 @@ async fn register_source(
928928
endpoint,
929929
source.options.as_ref(),
930930
source.is_read_write(),
931+
source.hierarchy_level,
931932
)
932933
.await
933934
.with_context(|| format!("Failed to register DynamoDB '{}'", source.name))?;

‎crates/server/src/config.rs‎

Lines changed: 85 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -240,6 +240,11 @@ pub enum ConfigError {
240240
)]
241241
EmptyAllowedSchemas { name: String },
242242

243+
#[error(
244+
"Data source '{name}' has an empty 'allowed_tables' option. Either omit it to load all DynamoDB tables, or provide a non-empty comma-separated list such as \"products,orders\"."
245+
)]
246+
EmptyAllowedTables { name: String },
247+
243248
#[error("Data source '{name}' has a non-UTF8 path: {path:?}")]
244249
NonUtf8Path { name: String, path: PathBuf },
245250
}
@@ -708,6 +713,7 @@ const CATALOG_SUPPORTED_SOURCES: &[DataSourceType] = &[
708713
DataSourceType::Mysql,
709714
DataSourceType::Sqlite,
710715
DataSourceType::Seekdb,
716+
DataSourceType::Dynamodb,
711717
];
712718

713719
/// Data source types that support read_write access mode
@@ -785,6 +791,22 @@ fn validate_data_sources(data_sources: &[DataSource]) -> Result<()> {
785791
.into());
786792
}
787793
}
794+
795+
if source.source_type == DataSourceType::Dynamodb {
796+
if let Some(value) = source
797+
.options
798+
.as_ref()
799+
.and_then(|o| o.get("allowed_tables"))
800+
{
801+
let has_entry = value.split(',').any(|s| !s.trim().is_empty());
802+
if !has_entry {
803+
return Err(ConfigError::EmptyAllowedTables {
804+
name: source.name.clone(),
805+
}
806+
.into());
807+
}
808+
}
809+
}
788810
}
789811

790812
match (&source.source_type, s3_storage.is_remote_path(&source.path)) {
@@ -1292,6 +1314,7 @@ async fn register_data_source(
12921314
connection_string,
12931315
source.options.as_ref(),
12941316
source.access_mode.is_read_write(),
1317+
source.hierarchy_level,
12951318
)
12961319
.await
12971320
.map_err(|e| {
@@ -2252,6 +2275,68 @@ options:
22522275
);
22532276
}
22542277

2278+
#[test]
2279+
fn validate_rejects_dynamodb_catalog_empty_allowed_tables() {
2280+
let mut options = HashMap::new();
2281+
options.insert("allowed_tables".to_string(), " , , ".to_string());
2282+
let mut source = dynamodb_source(
2283+
"ddb",
2284+
Some("http://localhost:8000"),
2285+
Some(options),
2286+
AccessMode::ReadOnly,
2287+
);
2288+
source.hierarchy_level = HierarchyLevel::Catalog;
2289+
2290+
let err = validate_data_sources(&[source]).unwrap_err();
2291+
let config_err = err.downcast_ref::<ConfigError>().unwrap();
2292+
assert!(
2293+
matches!(
2294+
config_err,
2295+
ConfigError::EmptyAllowedTables { name } if name == "ddb"
2296+
),
2297+
"got {config_err}"
2298+
);
2299+
}
2300+
2301+
#[test]
2302+
fn validate_accepts_dynamodb_catalog_allowed_tables() {
2303+
let mut options = HashMap::new();
2304+
options.insert("allowed_tables".to_string(), "products, orders".to_string());
2305+
let mut source = dynamodb_source(
2306+
"ddb",
2307+
Some("http://localhost:8000"),
2308+
Some(options),
2309+
AccessMode::ReadOnly,
2310+
);
2311+
source.hierarchy_level = HierarchyLevel::Catalog;
2312+
2313+
validate_data_sources(&[source]).expect("valid DynamoDB catalog allow-list");
2314+
}
2315+
2316+
#[test]
2317+
fn validate_rejects_dynamodb_catalog_table_option() {
2318+
let mut options = HashMap::new();
2319+
options.insert("table".to_string(), "products".to_string());
2320+
let mut source = dynamodb_source(
2321+
"ddb",
2322+
Some("http://localhost:8000"),
2323+
Some(options),
2324+
AccessMode::ReadOnly,
2325+
);
2326+
source.hierarchy_level = HierarchyLevel::Catalog;
2327+
2328+
let err = validate_data_sources(&[source]).unwrap_err();
2329+
let config_err = err.downcast_ref::<ConfigError>().unwrap();
2330+
assert!(
2331+
matches!(
2332+
config_err,
2333+
ConfigError::CatalogModeConflictingOptions { name, option }
2334+
if name == "ddb" && option == "table"
2335+
),
2336+
"got {config_err}"
2337+
);
2338+
}
2339+
22552340
#[test]
22562341
fn unsupported_write_mode_error_lists_dynamodb() {
22572342
let err = ConfigError::UnsupportedWriteMode {

‎crates/skardi/src/sources/hierarchy/mod.rs‎

Lines changed: 136 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -63,6 +63,13 @@ pub const CONNECT_TIMEOUT: Duration = Duration::from_secs(5);
6363
/// Maximum attempts (including the first) for [`retry_with_timeout`].
6464
pub const MAX_RETRIES: u32 = 3;
6565

66+
/// Summary returned by catalog assembly helpers.
67+
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
68+
pub struct CatalogBuildReport {
69+
pub registered: usize,
70+
pub skipped: usize,
71+
}
72+
6673
/// How much of an upstream database to expose in DataFusion (single table vs whole catalog).
6774
#[derive(Debug, Clone, Copy, Deserialize, Serialize, Hash, PartialEq, Eq, Default)]
6875
#[serde(rename_all = "lowercase")]
@@ -173,14 +180,35 @@ pub async fn build_catalog<F, Fut>(
173180
session_ctx: &SessionContext,
174181
catalog_name: &str,
175182
schema_tables: Vec<(String, String)>,
176-
mut build_table: F,
183+
build_table: F,
177184
) -> Result<()>
178185
where
179186
F: FnMut(String, String) -> Fut,
180187
Fut: Future<Output = Result<Arc<dyn TableProvider>>>,
181188
{
182-
let catalog_provider = Arc::new(MemoryCatalogProvider::new());
189+
build_catalog_with_required_schemas(
190+
session_ctx,
191+
catalog_name,
192+
schema_tables,
193+
Vec::new(),
194+
build_table,
195+
)
196+
.await
197+
.map(|_| ())
198+
}
183199

200+
/// Fail-fast catalog assembly with schemas that must exist even when no table is registered.
201+
pub async fn build_catalog_with_required_schemas<F, Fut>(
202+
session_ctx: &SessionContext,
203+
catalog_name: &str,
204+
schema_tables: Vec<(String, String)>,
205+
required_schemas: Vec<String>,
206+
mut build_table: F,
207+
) -> Result<CatalogBuildReport>
208+
where
209+
F: FnMut(String, String) -> Fut,
210+
Fut: Future<Output = Result<Arc<dyn TableProvider>>>,
211+
{
184212
// Kick off provider construction concurrently. `build_table` (FnMut) is called eagerly
185213
// and sequentially in the map — the closures it returns are what run in parallel.
186214
let provider_futures: Vec<_> = schema_tables
@@ -200,15 +228,114 @@ where
200228
})
201229
.collect();
202230

203-
let mut prepared: Vec<(String, String, Arc<dyn TableProvider>)> =
231+
let prepared: Vec<(String, String, Arc<dyn TableProvider>)> = stream::iter(provider_futures)
232+
.buffer_unordered(CATALOG_BUILD_CONCURRENCY)
233+
.try_collect()
234+
.await?;
235+
236+
register_catalog_tables(session_ctx, catalog_name, required_schemas, prepared, 0)
237+
}
238+
239+
/// Best-effort catalog assembly. Failed table providers are skipped after a warning.
240+
pub async fn build_catalog_best_effort<F, Fut>(
241+
session_ctx: &SessionContext,
242+
catalog_name: &str,
243+
schema_tables: Vec<(String, String)>,
244+
required_schemas: Vec<String>,
245+
mut build_table: F,
246+
) -> Result<CatalogBuildReport>
247+
where
248+
F: FnMut(String, String) -> Fut,
249+
Fut: Future<Output = Result<Arc<dyn TableProvider>>>,
250+
{
251+
let provider_futures: Vec<_> = schema_tables
252+
.into_iter()
253+
.map(|(schema, table_name)| {
254+
let fut = build_table(schema.clone(), table_name.clone());
255+
let catalog_name = catalog_name.to_string();
256+
async move {
257+
let result = fut.await.with_context(|| {
258+
format!(
259+
"Failed to build table provider for '{}.{}' in catalog '{}'",
260+
schema, table_name, catalog_name
261+
)
262+
});
263+
(schema, table_name, result)
264+
}
265+
})
266+
.collect();
267+
268+
let prepared_results: Vec<(String, String, Result<Arc<dyn TableProvider>>)> =
204269
stream::iter(provider_futures)
205270
.buffer_unordered(CATALOG_BUILD_CONCURRENCY)
206-
.try_collect()
207-
.await?;
271+
.collect()
272+
.await;
273+
274+
let mut skipped = 0usize;
275+
let mut prepared = Vec::new();
276+
for (schema, table_name, provider_result) in prepared_results {
277+
match provider_result {
278+
Ok(provider) => prepared.push((schema, table_name, provider)),
279+
Err(e) => {
280+
skipped += 1;
281+
tracing::warn!(
282+
catalog = %catalog_name,
283+
schema = %schema,
284+
table = %table_name,
285+
error = %e,
286+
"Skipping catalog table after provider build failure"
287+
);
288+
}
289+
}
290+
}
291+
292+
if prepared.is_empty() && skipped > 0 {
293+
tracing::warn!(
294+
catalog = %catalog_name,
295+
skipped,
296+
"Catalog registered with no tables because every table was skipped"
297+
);
298+
}
299+
300+
register_catalog_tables(
301+
session_ctx,
302+
catalog_name,
303+
required_schemas,
304+
prepared,
305+
skipped,
306+
)
307+
}
308+
309+
fn register_catalog_tables(
310+
session_ctx: &SessionContext,
311+
catalog_name: &str,
312+
mut required_schemas: Vec<String>,
313+
mut prepared: Vec<(String, String, Arc<dyn TableProvider>)>,
314+
skipped: usize,
315+
) -> Result<CatalogBuildReport> {
316+
let catalog_provider = Arc::new(MemoryCatalogProvider::new());
317+
318+
required_schemas.sort();
319+
required_schemas.dedup();
320+
for schema in required_schemas {
321+
if catalog_provider.schema(&schema).is_none() {
322+
catalog_provider
323+
.register_schema(&schema, Arc::new(MemorySchemaProvider::new()))
324+
.map_err(|e| {
325+
anyhow::anyhow!(
326+
"Failed to register schema '{}' for catalog '{}': {}",
327+
schema,
328+
catalog_name,
329+
e
330+
)
331+
})?;
332+
}
333+
}
208334

209335
// Deterministic registration order for log output and downstream iteration.
210336
prepared.sort_by(|a, b| (a.0.as_str(), a.1.as_str()).cmp(&(b.0.as_str(), b.1.as_str())));
211337

338+
let registered = prepared.len();
212339
for (schema, table_name, table_provider) in prepared {
213340
if catalog_provider.schema(&schema).is_none() {
214341
catalog_provider
@@ -252,5 +379,8 @@ where
252379
}
253380

254381
session_ctx.register_catalog(catalog_name, catalog_provider);
255-
Ok(())
382+
Ok(CatalogBuildReport {
383+
registered,
384+
skipped,
385+
})
256386
}

0 commit comments

Comments
 (0)