mirror of https://github.com/buster-so/buster.git
84 lines
2.4 KiB
Rust
84 lines
2.4 KiB
Rust
|
use anyhow::{anyhow, Result};
|
||
|
|
use database::{
|
||
|
|
pool::get_pg_pool,
|
||
|
|
schema::datasets,
|
||
|
|
};
|
||
|
|
use diesel::{ExpressionMethods, QueryDsl};
|
||
|
|
use diesel_async::RunQueryDsl;
|
||
|
|
use serde::{Deserialize, Serialize};
|
||
|
|
use uuid::Uuid;
|
||
|
|
use std::collections::HashMap;
|
||
|
|
|
||
|
|
/// Response structure that maps dataset IDs to their data source IDs
|
||
|
|
#[derive(Debug, Serialize, Deserialize)]
|
||
|
|
pub struct DatasetWithDataSource {
|
||
|
|
pub dataset_id: Uuid,
|
||
|
|
pub data_source_id: Uuid,
|
||
|
|
pub name: String,
|
||
|
|
}
|
||
|
|
|
||
|
|
/// Helper function to get data source IDs for an array of dataset IDs
|
||
|
|
///
|
||
|
|
/// # Arguments
|
||
|
|
///
|
||
|
|
/// * `dataset_ids` - An array of dataset UUIDs to lookup
|
||
|
|
///
|
||
|
|
/// # Returns
|
||
|
|
///
|
||
|
|
/// A Result containing a Vec of DatasetWithDataSource structs, each mapping a dataset ID to its data source ID
|
||
|
|
pub async fn get_data_sources_for_datasets(dataset_ids: &[Uuid]) -> Result<Vec<DatasetWithDataSource>> {
|
||
|
|
if dataset_ids.is_empty() {
|
||
|
|
return Ok(vec![]);
|
||
|
|
}
|
||
|
|
|
||
|
|
let mut conn = match get_pg_pool().get().await {
|
||
|
|
Ok(conn) => conn,
|
||
|
|
Err(e) => return Err(anyhow!("Failed to get database connection: {}", e)),
|
||
|
|
};
|
||
|
|
|
||
|
|
let datasets_result = datasets::table
|
||
|
|
.filter(datasets::id.eq_any(dataset_ids))
|
||
|
|
.filter(datasets::deleted_at.is_null())
|
||
|
|
.select((
|
||
|
|
datasets::id,
|
||
|
|
datasets::data_source_id,
|
||
|
|
datasets::name,
|
||
|
|
))
|
||
|
|
.load::<(Uuid, Uuid, String)>(&mut conn)
|
||
|
|
.await
|
||
|
|
.map_err(|e| anyhow!("Error querying datasets: {}", e))?;
|
||
|
|
|
||
|
|
// Map the results to the response structure
|
||
|
|
let result = datasets_result
|
||
|
|
.into_iter()
|
||
|
|
.map(|(dataset_id, data_source_id, name)| {
|
||
|
|
DatasetWithDataSource {
|
||
|
|
dataset_id,
|
||
|
|
data_source_id,
|
||
|
|
name,
|
||
|
|
}
|
||
|
|
})
|
||
|
|
.collect();
|
||
|
|
|
||
|
|
Ok(result)
|
||
|
|
}
|
||
|
|
|
||
|
|
/// Helper function that returns a HashMap mapping dataset IDs to their data source IDs
|
||
|
|
///
|
||
|
|
/// # Arguments
|
||
|
|
///
|
||
|
|
/// * `dataset_ids` - An array of dataset UUIDs to lookup
|
||
|
|
///
|
||
|
|
/// # Returns
|
||
|
|
///
|
||
|
|
/// A Result containing a HashMap where keys are dataset IDs and values are data source IDs
|
||
|
|
pub async fn get_data_source_map_for_datasets(dataset_ids: &[Uuid]) -> Result<HashMap<Uuid, Uuid>> {
|
||
|
|
let datasets = get_data_sources_for_datasets(dataset_ids).await?;
|
||
|
|
|
||
|
|
let map = datasets
|
||
|
|
.into_iter()
|
||
|
|
.map(|ds| (ds.dataset_id, ds.data_source_id))
|
||
|
|
.collect();
|
||
|
|
|
||
|
|
Ok(map)
|
||
|
|
}
|