prds for tasks

This commit is contained in:
dal 2025-04-01 12:46:47 -06:00
parent d255c5a719
commit 79f9e2a352
No known key found for this signature in database
GPG Key ID: 16F4B0E1E9F61122
4 changed files with 911 additions and 0 deletions

View File

@ -0,0 +1,257 @@
# PRD: Add Most Recent File ID and Type to Chat Listing
## Problem Statement ✅
Currently, the chat listing (`list_chats_handler.rs` and `list_chats.rs`) does not efficiently display information about the most recent file associated with a chat. Users need to open individual chats to determine if relevant files (and their types) have been added or updated recently. This requires fetching and joining data across multiple tables (`chats`, `messages`, `messages_to_files`, `metric_files`, `dashboard_files`) on each request, which is inefficient, especially for users with many chats.
Key issues:
- Inefficient retrieval of the most recent file ID and type for the chat list.
- Users cannot quickly identify chats with recent file activity and the type of file without opening them.
- Potential performance degradation as chat and message volume increases.
### Current Limitations
- No direct fields on the `chats` table to indicate the most recent file ID and type.
- Requires complex joins or multiple queries to derive this information.
### Impact
- User Impact: Slower loading times for the chat list, reduced productivity when searching for chats with specific file activity.
- System Impact: Increased database load due to repeated complex queries for file information.
## Requirements
### Functional Requirements ✅
- Add optional `most_recent_file_id` and `most_recent_file_type` columns to the `chats` table.
- Details: These columns will store the `UUID` and type (e.g., 'metric', 'dashboard') of the file associated with the most recently created message within that chat that has a file attached.
- Acceptance Criteria: The `chats` table has new nullable `most_recent_file_id` (UUID) and `most_recent_file_type` (VARCHAR or ENUM) columns.
- Dependencies: Database migration capabilities.
- Update the `chats` table whenever a message with a file is added.
- Details: Modify the logic where messages are created (`post_chat_handler.rs` or similar) to update the corresponding chat's `most_recent_file_id` and `most_recent_file_type`. Only update if the new message is more recent than the one currently associated with the chat's `most_recent_file_id`.
- Acceptance Criteria: The `chats.most_recent_file_id` and `chats.most_recent_file_type` are accurately updated upon the creation of a new message with an associated file.
- Dependencies: Access to message creation logic and file type information during creation.
- Include `latest_file_id` and `latest_file_type` in the `ChatListItem` response.
- Details: Modify `list_chats_handler.rs` to select the `most_recent_file_id` and `most_recent_file_type` from the `chats` table and include them in the `ChatListItem` struct.
- Acceptance Criteria: The `/chats` API response includes optional `latest_file_id` and `latest_file_type` fields for each chat item.
- Dependencies: Existing `ChatListItem` struct and `list_chats_handler.rs`.
### Non-Functional Requirements ✅
- Performance: Chat list retrieval should be fast, leveraging the direct column lookup. Query time should not increase significantly.
- Data Integrity: Ensure the `most_recent_file_id` and `most_recent_file_type` are reliably updated and kept consistent.
- Maintainability: The update logic should be centralized or clearly documented.
## Technical Design ✅
### System Architecture
Add denormalized columns to the `chats` table to optimize read performance. Updates occur during message/file creation.
### Core Components ✅
#### Component 1: Chat Model Update
```rust
// libs/database/src/models.rs
#[derive(Queryable, Insertable, Identifiable, Associations, Debug, Clone, Serialize)]
#[diesel(belongs_to(Organization))]
#[diesel(belongs_to(User, foreign_key = created_by))]
#[diesel(table_name = chats)]
pub struct Chat {
pub id: Uuid,
pub title: String,
pub organization_id: Uuid,
pub created_at: DateTime<Utc>,
pub updated_at: DateTime<Utc>,
pub deleted_at: Option<DateTime<Utc>>,
pub created_by: Uuid,
pub updated_by: Uuid,
pub publicly_accessible: bool,
pub publicly_enabled_by: Option<Uuid>,
pub public_expiry_date: Option<DateTime<Utc>>,
pub most_recent_file_id: Option<Uuid>,
pub most_recent_file_type: Option<String>,
}
```
#### Component 2: ChatListItem Update
```rust
// libs/handlers/src/chats/list_chats_handler.rs
#[derive(Debug, Serialize, Deserialize)]
pub struct ChatListItem {
pub id: String,
pub name: String,
pub is_favorited: bool,
pub updated_at: String,
pub created_at: String,
pub created_by: String,
pub created_by_id: String,
pub created_by_name: String,
pub created_by_avatar: Option<String>,
pub last_edited: String,
pub latest_file_id: Option<String>,
pub latest_file_type: Option<String>,
}
// Updated Queryable struct
#[derive(Queryable)]
struct ChatWithUser {
pub id: Uuid,
pub name: String,
pub created_at: DateTime<Utc>,
pub updated_at: DateTime<Utc>,
pub created_by: Uuid,
pub most_recent_file_id: Option<Uuid>,
pub most_recent_file_type: Option<String>,
pub user_name: Option<String>,
pub user_attributes: Value,
}
```
### Database Changes (If applicable) ✅
- **New Columns**: Add `most_recent_file_id` (UUID, nullable) and `most_recent_file_type` (VARCHAR, nullable) to the `chats` table.
- **Foreign Key**: Omitted for `most_recent_file_id` for simplicity as discussed.
- **Index**: Add indices on `most_recent_file_id` and potentially `most_recent_file_type` if filtering/sorting by type is anticipated.
```sql
-- Migration: up.sql
ALTER TABLE chats
ADD COLUMN most_recent_file_id UUID NULL,
ADD COLUMN most_recent_file_type VARCHAR(255) NULL; -- Adjust VARCHAR size or use ENUM if available
-- Optional: Add indices
CREATE INDEX idx_chats_most_recent_file_id ON chats(most_recent_file_id);
CREATE INDEX idx_chats_most_recent_file_type ON chats(most_recent_file_type);
-- Data Migration (Run once after adding the columns)
-- This needs refinement based on how file type is determined.
-- Assuming file type can be retrieved by joining messages_to_files.file_id
-- with the actual file tables (e.g., metric_files, dashboard_files)
WITH LatestFilePerChat AS (
SELECT DISTINCT ON (m.chat_id)
m.chat_id,
mtf.file_id,
-- Need to determine file type here, example placeholder:
CASE
WHEN mf.id IS NOT NULL THEN 'metric'
WHEN df.id IS NOT NULL THEN 'dashboard'
ELSE NULL
END AS file_type
FROM messages m
JOIN messages_to_files mtf ON m.id = mtf.message_id
LEFT JOIN metric_files mf ON mtf.file_id = mf.id -- Join to determine type
LEFT JOIN dashboard_files df ON mtf.file_id = df.id -- Join to determine type
WHERE m.deleted_at IS NULL AND mtf.deleted_at IS NULL
ORDER BY m.chat_id, m.created_at DESC
)
UPDATE chats c
SET
most_recent_file_id = lfpc.file_id,
most_recent_file_type = lfpc.file_type
FROM LatestFilePerChat lfpc
WHERE c.id = lfpc.chat_id;
-- Migration: down.sql
ALTER TABLE chats
DROP COLUMN IF EXISTS most_recent_file_id,
DROP COLUMN IF EXISTS most_recent_file_type;
-- Optional: Drop indices
DROP INDEX IF EXISTS idx_chats_most_recent_file_id;
DROP INDEX IF EXISTS idx_chats_most_recent_file_type;
```
### API Changes (If applicable) ✅
- **Response Modification**: The `GET /chats` endpoint response will include `latest_file_id: Option<String>` and `latest_file_type: Option<String>` in each `ChatListItem`.
### File Changes (If applicable) ✅
#### New Files
- `migrations/YYYY-MM-DD-HHMMSS_add_file_info_to_chats/up.sql`
- `migrations/YYYY-MM-DD-HHMMSS_add_file_info_to_chats/down.sql`
#### Modified Files
- `libs/database/src/schema.rs` (via `diesel print-schema`)
- `libs/database/src/models.rs` (Add fields to `Chat` struct)
- `libs/handlers/src/chats/post_chat_handler.rs` (Update chat on `MessageToFile` creation)
- `libs/handlers/src/chats/list_chats_handler.rs` (Select new fields, update `ChatWithUser`, update mapping to `ChatListItem`)
## Implementation Plan
### Phase 1: Database Migration & Model Update ✅
1. Generate Diesel migration: `diesel migration generate add_file_info_to_chats`
2. Implement `up.sql`: Add columns and indices.
3. Implement `down.sql`: Drop columns and indices.
4. Run `diesel print-schema`.
5. Update `Chat` struct in `models.rs`.
6. Run the migration: `diesel migration run`.
7. Run the data migration SQL (adjusting file type logic as needed).
### Phase 2: Update Message Creation Logic ✅
1. Modify `post_chat_handler.rs`:
- Within the logic that processes completed files (e.g., `process_completed_files` function) after successfully inserting a `MessageToFile` record:
- Retrieve the `chat_id` associated with the `message_id`.
- Determine the `file_type` (e.g., from `BusterReasoningFile.file_type`).
- Execute a Diesel update statement to set the `most_recent_file_id` and `most_recent_file_type` on the `chats` table for the specific `chat_id`.
```rust
// Example Diesel update logic within post_chat_handler.rs
// (inside the loop processing completed files after MessageToFile insert)
use crate::schema::chats;
use diesel::prelude::*;
use diesel_async::RunQueryDsl;
let file_id_to_set = Uuid::parse_str(file_id)?;
let file_type_to_set = file_content.file_type.clone(); // Assuming file_content is available
let target_chat_id = message.chat_id;
// Ensure this update only happens if the current message is indeed the latest
// This might require fetching the current most_recent_file timestamp or message timestamp
// For simplicity, this example updates unconditionally, but production code should verify.
diesel::update(chats::table.find(target_chat_id))
.set((
chats::most_recent_file_id.eq(Some(file_id_to_set)),
chats::most_recent_file_type.eq(Some(file_type_to_set)),
chats::updated_at.eq(Utc::now()), // Also update the chat's updated_at timestamp
))
.execute(conn) // Assuming conn is the AsyncPgConnection
.await?;
```
### Phase 3: Update Chat Listing Logic ✅
1. Modify `list_chats_handler.rs`:
- Add `chats::most_recent_file_id`, `chats::most_recent_file_type` to the `select` clause in the main Diesel query.
- Add `most_recent_file_type: Option<String>` field to the `ChatWithUser` struct.
- Update the `.map(|chat: ChatWithUser| { ... })` block to populate `ChatListItem.latest_file_id` and `ChatListItem.latest_file_type` from the `ChatWithUser` fields.
### Phase 4: Testing & Documentation ✅
1. Add unit tests for `post_chat_handler.rs` update logic (including type).
2. Update integration tests for `list_chats.rs` to verify `latest_file_id` and `latest_file_type`.
3. Test scenarios: chats with no files, different file types.
4. Document the new fields and update mechanism.
## Testing Strategy ✅
### Unit Tests
- Test the `post_chat_handler` logic:
- Verify `most_recent_file_id` and `most_recent_file_type` are updated correctly.
- Test with different file types.
### Integration Tests
- Test the `GET /chats` endpoint:
- Verify `latest_file_id` and `latest_file_type` are present and correct.
- Verify `null` values for chats without files.
## Security Considerations
- Ensure file type information is handled safely and doesn't introduce vulnerabilities (e.g., if using ENUMs, ensure proper validation).
## References
- @libs/handlers/src/chats/list_chats_handler.rs
- @src/routes/rest/routes/chats/list_chats.rs
- @libs/database/src/schema.rs
- @libs/database/src/models.rs
- @libs/handlers/src/chats/post_chat_handler.rs
- @libs/database/src/models.rs (MessageToFile struct)
- @libs/database/src/schema.rs (messages_to_files table)
- @libs/handlers/src/chats/post_chat_handler.rs (BusterReasoningFile struct for type info)

View File

@ -0,0 +1,214 @@
# PRD: Enhance Get Dashboard Handler with Collections
## Problem Statement ✅
The current `get_dashboard_handler` in `libs/handlers/src/dashboards/get_dashboard_handler.rs` retrieves detailed information about a specific dashboard but does not include information about which collections the dashboard belongs to. Users often need this context to understand how a dashboard is organized and grouped with other assets.
Adding this information directly to the `BusterDashboardResponse` type will improve usability and provide a more complete picture of the dashboard's organizational context.
### Current Limitations
- The `BusterDashboardResponse` struct currently includes an empty vector for `collections`.
- Users must perform separate queries or manually search to find which collections contain a specific dashboard.
### Impact
- User Impact: Increased effort required to understand dashboard organization. Potential for missing context about how dashboards are grouped.
- System Impact: More API calls needed to gather complete dashboard context.
- Business Impact: Slower analysis and navigation for users trying to understand dashboard relationships within collections.
## Requirements
### Functional Requirements ✅
#### Core Functionality
- Requirement 1: Fetch Associated Collections
- Details: When `get_dashboard_handler` is called, it should query the database to find all collections that include the requested dashboard.
- Acceptance Criteria: The `collections` field in the `BusterDashboardResponse` should contain a list of objects, each including the `id` and `name` of an associated collection.
- Dependencies: Requires access to `collections_to_assets` and `collections` tables.
### Non-Functional Requirements ✅
- Performance Requirements: The additional query should not significantly degrade the performance of the `get_dashboard_handler`. Aim for <100ms added latency. Consider optimizing the query.
- Security Requirements: Ensure existing permission checks are respected. Users should only see associations with collections they have permission to view (or at least know exist). *Initial implementation may return all associated collection IDs and names, assuming visibility is handled by subsequent requests for those assets.*
- Scalability Requirements: The query should scale efficiently as the number of dashboards and collections grows. Use appropriate indexing.
## Technical Design ✅
### System Architecture
No changes to the overall system architecture. This change enhances an existing handler by adding concurrent data fetching.
### Core Components ✅
#### Component 1: `libs/handlers/src/dashboards/get_dashboard_handler.rs`
```rust
// ... existing imports ...
use database::schema::{collections, collections_to_assets};
use database::models::Collection; // Assuming Collection model exists
use tokio::task::JoinHandle;
// Define struct for associated collection (assuming defined, possibly in types.rs or shared location)
// #[derive(Debug, Serialize, Deserialize, Clone)]
// pub struct AssociatedCollection { ... }
// Modify BusterDashboardResponse struct (assuming defined elsewhere)
// #[derive(Debug, Serialize, Deserialize, Clone)]
// pub struct BusterDashboardResponse { ... collections: Vec<AssociatedCollection> ... }
// --- NEW HELPER FUNCTION START ---
async fn fetch_associated_collections_for_dashboard(dashboard_id: Uuid) -> Result<Vec<AssociatedCollection>> {
let mut conn = get_pg_pool().get().await?;
let associated_collections = collections_to_assets::table
.inner_join(collections::table.on(collections::id.eq(collections_to_assets::collection_id)))
.filter(collections_to_assets::asset_id.eq(dashboard_id))
.filter(collections_to_assets::asset_type.eq(AssetType::DashboardFile)) // Filter by asset type
.filter(collections::deleted_at.is_null()) // Ensure collection isn't deleted
.select((collections::id, collections::name))
.load::<(Uuid, String)>(&mut conn)
.await?
.into_iter()
.map(|(id, name)| AssociatedCollection { id, name })
.collect();
Ok(associated_collections)
}
// --- NEW HELPER FUNCTION END ---
pub async fn get_dashboard_handler(
dashboard_id: &Uuid,
user: &AuthenticatedUser,
version_number: Option<i32>,
) -> Result<BusterDashboardResponse> {
// ... existing logic to fetch dashboard_with_permission, parse content, check permissions ...
// Clone dashboard_id for use in spawned task
let d_id = *dashboard_id;
// --- UPDATED LOGIC START ---
// Spawn task to fetch collections concurrently
let collections_handle: JoinHandle<Result<Vec<AssociatedCollection>>> =
tokio::spawn(async move { fetch_associated_collections_for_dashboard(d_id).await });
// Fetch metrics concurrently (assuming this is already happening or will be added)
// Example: let metrics_handle = tokio::spawn(async move { fetch_metrics_for_dashboard(...).await });
// Await results - potentially join multiple handles
// Example: let (collections_result, metrics_result) = tokio::join!(collections_handle, metrics_handle);
let collections_result = collections_handle.await;
// Handle collections result
let collections = match collections_result {
Ok(Ok(c)) => c,
Ok(Err(e)) => {
tracing::error!("Failed to fetch associated collections for dashboard {}: {}", dashboard_id, e);
vec![]
}
Err(e) => { // JoinError
tracing::error!("Task join error fetching collections for dashboard {}: {}", dashboard_id, e);
vec![]
}
};
// Handle metrics result (if fetched concurrently)
// Example: let metrics = handle_metrics_result(metrics_result, dashboard_id);
// Fetch metrics sequentially for now if not concurrent
// Collect all metric IDs from the rows
let metric_ids: Vec<Uuid> = config // Assuming config is parsed earlier
.rows
.iter()
.flat_map(|row| {
row.items
.iter()
.filter_map(|item| Uuid::parse_str(&item.id).ok())
})
.collect();
// Fetch all metrics (latest versions)
let metric_futures: Vec<_> = metric_ids
.iter()
.map(|metric_id| get_metric_handler(metric_id, &user, None)) // Assuming get_metric_handler exists
.collect();
let metric_results = futures::future::join_all(metric_futures).await;
let metrics: std::collections::HashMap<Uuid, BusterMetric> = metric_results // Assuming BusterMetric type
.into_iter()
.filter_map(|result| result.ok())
.map(|metric| (metric.id, metric))
.collect();
// --- UPDATED LOGIC END ---
// ... existing logic to fetch individual permissions, versions, construct final dashboard object ...
// Ensure a db connection is available if needed for sequential tasks
// let mut conn = get_pg_pool().get().await?;
Ok(BusterDashboardResponse {
// ... existing field assignments ...
metrics, // Assign fetched metrics
collections, // Assign fetched collections
// ... rest of the fields ...
})
}
```
### Database Changes (If applicable)
No schema changes required. Ensure appropriate indexes exist:
- `collections_to_assets`: Composite index on `(asset_id, asset_type)`.
- `collections`: Index on `id`.
### API Changes (If applicable)
The response structure of the endpoint(s) using `get_dashboard_handler` will change. The `collections` field within the `BusterDashboardResponse` object will now be populated with a list of `{id: Uuid, name: String}` objects instead of an empty array.
### File Changes (If applicable)
#### Modified Files
- `libs/handlers/src/dashboards/get_dashboard_handler.rs`
- Changes:
- Added a new private async function: `fetch_associated_collections_for_dashboard`.
- Modified `get_dashboard_handler` to use `tokio::spawn` to call this helper function concurrently with other potential async operations (like fetching metrics).
- Updated error handling for the concurrent task.
- Updated the `BusterDashboardResponse` struct field `collections`.
- Note: Assumes metric fetching might also be made concurrent or remains sequential as shown.
- Impact: Handler now performs collection fetch concurrently. Introduces dependency on `tokio`.
- Dependencies: `database` crate, `tokio`, `futures`, `anyhow`, `tracing`, `uuid`.
- `libs/handlers/src/dashboards/types.rs` (or wherever `BusterDashboardResponse`, `AssociatedCollection` are defined)
- Changes: Update/Define `BusterDashboardResponse` and `AssociatedCollection` structs.
- Impact: Changes the structure returned by the handler.
- Dependencies: `uuid::Uuid`, `serde::{Serialize, Deserialize}`.
## Implementation Plan
### Phase 1: Implement Backend Logic & Update Types ⏳ (Not Started)
1. Define or reuse `AssociatedCollection` struct (e.g., in `libs/handlers/src/dashboards/types.rs` or a shared location).
2. Update `BusterDashboardResponse` struct definition to use `Vec<AssociatedCollection>` for the `collections` field.
3. Implement the database query within `get_dashboard_handler.rs` to fetch associated collections.
4. Populate the `collections` field in the returned `BusterDashboardResponse` object.
5. Add tracing for errors during the fetch operation.
### Phase 2: Testing & Documentation 🔜 (Not Started)
1. Add unit tests for the new logic in `get_dashboard_handler.rs`. Mock database interactions to test:
- Correctly fetching collections.
- Handling cases where a dashboard has no collections.
- Handling database errors gracefully.
2. Add integration tests to verify the API endpoint returns the expected associations.
3. Update relevant documentation (e.g., API docs, internal handler docs) to reflect the change in the response structure.
## Testing Strategy
- Unit Tests: Mock database responses in `get_dashboard_handler_test.rs` to ensure the query logic is correct and handles various scenarios (associations found, none found, errors).
- Integration Tests: Create test data (dashboard, collections, associations) and call the relevant API endpoint(s). Verify the response contains the correct `collections` list. Test with different permission levels if applicable.
## Rollback Plan
- Revert the changes to `get_dashboard_handler.rs` and the `BusterDashboardResponse` type definition.
- Redeploy the previous version of the code.
## Security Considerations
- Ensure that fetching associated collection IDs and names does not leak sensitive information. The current design assumes IDs and names are not sensitive, but access to the full collection content is protected by separate permission checks when fetching those assets directly. Re-evaluate if names could be considered sensitive.
## Dependencies
- Database tables: `collections_to_assets`, `collections`.
- Enums: `database::enums::AssetType`.
- Crates: `diesel`, `diesel-async`, `anyhow`, `uuid`, `serde`, `tracing`.
## File References
- Handler: `libs/handlers/src/dashboards/get_dashboard_handler.rs`
- Types: `libs/handlers/src/dashboards/types.rs` (or wherever `BusterDashboardResponse` is defined)
- Schema: `libs/database/src/schema.rs`
- Models: `libs/database/src/models.rs`
- Enums: `libs/database/src/enums.rs`
</rewritten_file>

View File

@ -0,0 +1,235 @@
# PRD: Enhance Get Metric Handler with Associations
## Problem Statement ✅
The current `get_metric_handler` in `libs/handlers/src/metrics/get_metric_handler.rs` retrieves detailed information about a specific metric but does not include information about which dashboards or collections the metric belongs to. Users often need this context to understand where a metric is being used and how it relates to other assets. Adding this information directly to the `BusterMetric` response type will improve usability and provide a more complete picture of the metric's context within the system.
### Current Limitations
- The `BusterMetric` response struct currently includes empty vectors for `dashboards` and `collections`.
- Users must perform separate queries or manually search to find which dashboards and collections utilize a specific metric.
### Impact
- User Impact: Increased effort required to understand metric usage and relationships. Potential for overlooking metric usage in certain dashboards or collections.
- System Impact: More API calls needed to gather complete metric context.
- Business Impact: Slower analysis and navigation for users trying to understand metric relationships.
## Requirements
### Functional Requirements ✅
#### Core Functionality
- Requirement 1: Fetch Associated Dashboards
- Details: When `get_metric_handler` is called, it should query the database to find all dashboards that include the requested metric.
- Acceptance Criteria: The `dashboards` field in the `BusterMetric` response should contain a list of objects, each including the `id` and `name` of an associated dashboard.
- Dependencies: Requires access to `metric_files_to_dashboard_files` and `dashboard_files` tables.
- Requirement 2: Fetch Associated Collections
- Details: When `get_metric_handler` is called, it should query the database to find all collections that include the requested metric.
- Acceptance Criteria: The `collections` field in the `BusterMetric` response should contain a list of objects, each including the `id` and `name` of an associated collection.
- Dependencies: Requires access to `collections_to_assets` and `collections` tables.
### Non-Functional Requirements ✅
- Performance Requirements: The additional queries should not significantly degrade the performance of the `get_metric_handler`. Aim for <100ms added latency per query. Consider optimizing queries, potentially joining them if feasible.
- Security Requirements: Ensure existing permission checks are respected. Users should only see associations with dashboards/collections they have permission to view (or at least know exist). *Initial implementation may return all associated dashboards/collections IDs and names, assuming visibility is handled by subsequent requests for those assets.*
- Scalability Requirements: Queries should scale efficiently as the number of metrics, dashboards, and collections grows. Use appropriate indexing.
## Technical Design ✅
### System Architecture
No changes to the overall system architecture. This change enhances an existing handler by adding concurrent data fetching.
### Core Components ✅
#### Component 1: `libs/handlers/src/metrics/get_metric_handler.rs`
```rust
// ... existing imports ...
use database::schema::{collections, collections_to_assets, dashboard_files, metric_files_to_dashboard_files};
use database::models::{Collection, DashboardFile}; // Assuming Collection model exists
use tokio::task::JoinHandle;
use futures::future::try_join_all;
// Define structs for associated assets
#[derive(Debug, Serialize, Deserialize, Clone)]
pub struct AssociatedDashboard {
pub id: Uuid,
pub name: String,
}
#[derive(Debug, Serialize, Deserialize, Clone)]
pub struct AssociatedCollection {
pub id: Uuid,
pub name: String,
}
// Modify BusterMetric struct
#[derive(Debug, Serialize, Deserialize, Clone)]
pub struct BusterMetric {
// ... existing fields ...
pub dashboards: Vec<AssociatedDashboard>, // Updated type
pub collections: Vec<AssociatedCollection>, // Updated type
// ... existing fields ...
}
// --- NEW HELPER FUNCTIONS START ---
async fn fetch_associated_dashboards_for_metric(metric_id: Uuid) -> Result<Vec<AssociatedDashboard>> {
let mut conn = get_pg_pool().get().await?;
let associated_dashboards = metric_files_to_dashboard_files::table
.inner_join(dashboard_files::table.on(dashboard_files::id.eq(metric_files_to_dashboard_files::dashboard_file_id)))
.filter(metric_files_to_dashboard_files::metric_file_id.eq(metric_id))
.filter(dashboard_files::deleted_at.is_null())
.select((dashboard_files::id, dashboard_files::name))
.load::<(Uuid, String)>(&mut conn)
.await?
.into_iter()
.map(|(id, name)| AssociatedDashboard { id, name })
.collect();
Ok(associated_dashboards)
}
async fn fetch_associated_collections_for_metric(metric_id: Uuid) -> Result<Vec<AssociatedCollection>> {
let mut conn = get_pg_pool().get().await?;
let associated_collections = collections_to_assets::table
.inner_join(collections::table.on(collections::id.eq(collections_to_assets::collection_id)))
.filter(collections_to_assets::asset_id.eq(metric_id))
.filter(collections_to_assets::asset_type.eq(AssetType::MetricFile))
.filter(collections::deleted_at.is_null())
.select((collections::id, collections::name))
.load::<(Uuid, String)>(&mut conn)
.await?
.into_iter()
.map(|(id, name)| AssociatedCollection { id, name })
.collect();
Ok(associated_collections)
}
// --- NEW HELPER FUNCTIONS END ---
pub async fn get_metric_handler(
metric_id: &Uuid,
user: &AuthenticatedUser,
version_number: Option<i32>,
) -> Result<BusterMetric> {
// ... existing logic to fetch metric_file, check permissions ...
// Clone metric_id for use in spawned tasks
let m_id = *metric_id;
// --- UPDATED LOGIC START ---
// Spawn tasks to fetch associations concurrently
let dashboards_handle: JoinHandle<Result<Vec<AssociatedDashboard>>> =
tokio::spawn(async move { fetch_associated_dashboards_for_metric(m_id).await });
let collections_handle: JoinHandle<Result<Vec<AssociatedCollection>>> =
tokio::spawn(async move { fetch_associated_collections_for_metric(m_id).await });
// ... other potential concurrent tasks can be added here ...
// Await results
// Use join! or try_join! depending on whether other essential concurrent tasks exist
let (dashboards_result, collections_result) = tokio::join!(dashboards_handle, collections_handle);
// Handle results, logging errors but returning empty Vecs for failed tasks
let dashboards = match dashboards_result {
Ok(Ok(d)) => d,
Ok(Err(e)) => {
tracing::error!("Failed to fetch associated dashboards for metric {}: {}", metric_id, e);
vec![]
}
Err(e) => { // JoinError
tracing::error!("Task join error fetching dashboards for metric {}: {}", metric_id, e);
vec![]
}
};
let collections = match collections_result {
Ok(Ok(c)) => c,
Ok(Err(e)) => {
tracing::error!("Failed to fetch associated collections for metric {}: {}", metric_id, e);
vec![]
}
Err(e) => { // JoinError
tracing::error!("Task join error fetching collections for metric {}: {}", metric_id, e);
vec![]
}
};
// --- UPDATED LOGIC END ---
// ... existing logic to parse metric content, fetch versions, individual permissions etc. ...
// Note: The main db connection `conn` might still be needed for sequential operations here.
// Ensure it's acquired if necessary, e.g., let mut conn = get_pg_pool().get().await?;
Ok(BusterMetric {
// ... existing field assignments ...
dashboards, // Assign fetched dashboards
collections, // Assign fetched collections
// ... rest of the fields ...
})
}
```
### Database Changes (If applicable)
No schema changes required. Ensure appropriate indexes exist:
- `metric_files_to_dashboard_files`: Index on `metric_file_id`.
- `collections_to_assets`: Composite index on `(asset_id, asset_type)`.
- `dashboard_files`: Index on `id`.
- `collections`: Index on `id`.
### API Changes (If applicable)
The response structure of the endpoint(s) using `get_metric_handler` will change. The `dashboards` and `collections` fields within the `BusterMetric` object will now be populated with lists of `{id: Uuid, name: String}` objects instead of empty arrays.
### File Changes (If applicable)
#### Modified Files
- `libs/handlers/src/metrics/get_metric_handler.rs`
- Changes:
- Added two new private async functions: `fetch_associated_dashboards_for_metric` and `fetch_associated_collections_for_metric`.
- Modified `get_metric_handler` to use `tokio::spawn` and `tokio::join!` to call these helper functions concurrently.
- Updated error handling for concurrent tasks.
- Updated the `BusterMetric` struct fields `dashboards` and `collections`.
- Impact: Handler now performs association fetches concurrently. Introduces dependency on `tokio` for spawning tasks and `futures` potentially for joining.
- Dependencies: `database` crate, `tokio`, `futures`, `anyhow`, `tracing`, `uuid`.
- `libs/handlers/src/metrics/types.rs` (or wherever `BusterMetric`, `AssociatedDashboard`, `AssociatedCollection` are defined)
- Changes: Update/Define `BusterMetric`, `AssociatedDashboard`, `AssociatedCollection` structs.
- Impact: Changes the structure returned by the handler.
- Dependencies: `uuid::Uuid`, `serde::{Serialize, Deserialize}`.
## Implementation Plan
### Phase 1: Implement Backend Logic & Update Types ⏳ (Not Started)
1. Define `AssociatedDashboard` and `AssociatedCollection` structs in `libs/handlers/src/metrics/types.rs` (or appropriate location).
2. Update `BusterMetric` struct definition to use these new types for `dashboards` and `collections`.
3. Implement the database queries within `get_metric_handler.rs` to fetch associated dashboards and collections.
4. Populate the `dashboards` and `collections` fields in the returned `BusterMetric` object.
5. Add tracing for errors during the fetch operations.
### Phase 2: Testing & Documentation 🔜 (Not Started)
1. Add unit tests for the new logic in `get_metric_handler.rs`. Mock database interactions to test:
- Correctly fetching dashboards.
- Correctly fetching collections.
- Handling cases where a metric has no associations.
- Handling database errors gracefully.
2. Add integration tests to verify the API endpoint returns the expected associations.
3. Update relevant documentation (e.g., API docs, internal handler docs) to reflect the change in the response structure.
## Testing Strategy
- Unit Tests: Mock database responses in `get_metric_handler_test.rs` to ensure the query logic is correct and handles various scenarios (associations found, none found, errors).
- Integration Tests: Create test data (metric, dashboards, collections, associations) and call the relevant API endpoint(s). Verify the response contains the correct `dashboards` and `collections` lists. Test with different permission levels if applicable.
## Rollback Plan
- Revert the changes to `get_metric_handler.rs` and the `BusterMetric` type definition.
- Redeploy the previous version of the code.
## Security Considerations
- Ensure that fetching associated IDs and names does not leak sensitive information. The current design assumes IDs and names are not sensitive, but access to the full dashboard/collection content is protected by separate permission checks when fetching those assets directly. Re-evaluate if names could be considered sensitive.
## Dependencies
- Database tables: `metric_files_to_dashboard_files`, `dashboard_files`, `collections_to_assets`, `collections`.
- Enums: `database::enums::AssetType`.
- Crates: `diesel`, `diesel-async`, `anyhow`, `uuid`, `serde`, `tracing`.
## File References
- Handler: `libs/handlers/src/metrics/get_metric_handler.rs`
- Types: `libs/handlers/src/metrics/types.rs` (or wherever `BusterMetric` is defined)
- Schema: `libs/database/src/schema.rs`
- Models: `libs/database/src/models.rs`
- Enums: `libs/database/src/enums.rs`

View File

@ -0,0 +1,205 @@
# PRD: Query Engine Performance and Limit Optimization
*Parent Project PRD:* (If applicable, link here)
*Status:* ⏳ In Progress
*Author:* AI Assistant
*Date Created:* YYYY-MM-DD
*Last Updated:* YYYY-MM-DD
## 1. Problem Statement ✅
The `query_engine` library provides a unified interface for querying various data sources (Postgres, MySQL, Redshift, SQL Server, Databricks, BigQuery, Snowflake). However, current implementations exhibit performance bottlenecks, particularly with large datasets, and inconsistent handling of row limits across different database connectors.
### Current Limitations
* **Inefficient Limiting:** Several connectors (Postgres, MySQL, Redshift, SQL Server) fetch potentially large result sets and apply limits *after* fetching, leading to unnecessary data transfer and processing.
* **Row-by-Row Processing:** Some implementations (Postgres, MySQL, SQL Server) process results row-by-row, often spawning an asynchronous task per row, which introduces significant overhead for large results.
* **Inconsistent Limit Implementation:** The default 5000 row limit is sometimes hardcoded (MySQL, Redshift) or applied post-fetch instead of being consistently applied at the database query level.
* **Type Conversion Overhead:** Each implementation performs type conversions for every column of every row, potentially using inefficient string comparisons.
* **Postgres Parsing Overhead:** The Postgres connector parses and manipulates the SQL AST for every query to handle identifier quoting, adding potentially unnecessary overhead.
### Impact
* **User Impact:** Users querying large tables experience slow response times. Inconsistent limit behavior can lead to unexpected results or errors.
* **System Impact:** High CPU and memory usage on the application server during query processing for large results. Increased network traffic between the application and the database for connectors that don't limit at the source.
* **Business Impact:** Degraded application performance affects user satisfaction and potentially impacts downstream processes relying on query results.
## 2. Requirements
### Functional Requirements ✅
* **FR1: Standardized Limit Handling:** All query functions MUST apply the row limit directly within the SQL query or API request sent to the database.
* *Details:* The default limit MUST be 5000 rows. The `limit: Option<i64>` parameter MUST override the default when provided. Post-fetch filtering/limiting in the application code MUST be removed.
* *Acceptance Criteria:* Queries without an explicit limit return max 5000 rows. Queries with an explicit `limit: Some(N)` return max N rows. Limits are visible in database query logs/profiles where applicable.
* *Dependencies:* None.
* **FR2: Consistent Implementation Pattern:** Ensure the `limit` parameter handling, default value application (5000), and overall function signature patterns are consistent across all query engine implementations.
* *Details:* Refactor implementations to follow a similar pattern for limit application and result processing where possible.
* *Acceptance Criteria:* Code review confirms consistent patterns.
* *Dependencies:* None.
### Non-Functional Requirements ✅
* **NFR1: Performance:** Significantly reduce query latency and server-side resource consumption (CPU, Memory) for queries that previously returned large datasets before limiting.
* *Details:* Optimize row processing and type conversion logic.
* *Acceptance Criteria:* Benchmark tests show measurable improvement in latency (average, p95) and reduced peak memory usage for affected queries compared to the baseline.
* **NFR2: Efficient Row Processing:** Implementations MUST avoid excessive overhead during result fetching and processing.
* *Details:* Eliminate patterns that spawn a new asynchronous task for each individual row. Process results using efficient streaming or batching mechanisms provided by the underlying database driver/client library.
* *Acceptance Criteria:* Code review confirms removal of per-row task spawning. Performance benchmarks (NFR1) validate efficiency gains.
* **NFR3: Optimized Type Conversion:** Minimize the computational cost of converting database types to the internal `DataType` enum.
* *Details:* Avoid costly string comparisons for type identification where alternative methods (e.g., type IDs, driver-provided types) are available. Explore caching type information if needed.
* *Acceptance Criteria:* Code review confirms optimization efforts. Profiling may be used to verify reduced overhead if significant issues are suspected.
* **NFR4: Memory Efficiency:** Reduce heap allocations and overall memory usage during result processing.
* *Details:* Pre-allocate collections like `Vec` and `IndexMap` with estimated capacity where possible. Reuse buffers or data structures if feasible within the processing loop.
* *Acceptance Criteria:* Memory usage benchmarks (NFR1) show improvement. Code review confirms best practices for memory allocation.
* **NFR5: Postgres SQL Parsing Optimization:** The SQL parsing step in `postgres_query.rs` MUST be reviewed for necessity and performance.
* *Details:* Determine if identifier quoting can be handled directly by `sqlx` or connection settings. If parsing is unavoidable, optimize it.
* *Acceptance Criteria:* The parsing step is removed OR profiling/benchmarking shows its performance impact is negligible or significantly reduced.
## 3. Technical Design ✅
### System Architecture
The `query_engine` library acts as a facade. A client (e.g., API handler) calls the main query function, which dispatches the request to the appropriate database-specific implementation based on the data source configuration. The specific implementation interacts with the database and transforms the results into a standardized `Vec<IndexMap<String, DataType>>`. This refactor focuses on optimizing the database-specific implementation modules.
```mermaid
graph TD
Client --> QE[Query Engine Facade]
QE -->|Dispatch| Postgres[postgres_query.rs]
QE -->|Dispatch| MySQL[mysql_query.rs]
QE -->|Dispatch| Redshift[redshift_query.rs]
QE -->|Dispatch| SQLServer[sql_server_query.rs]
QE -->|Dispatch| Databricks[databricks_query.rs]
QE -->|Dispatch| BigQuery[bigquery_query.rs]
QE -->|Dispatch| Snowflake[snowflake_query.rs]
Postgres --> DB[(Postgres DB)]
MySQL --> DB2[(MySQL DB)]
Redshift --> DB3[(Redshift DB)]
SQLServer --> DB4[(SQL Server DB)]
Databricks --> DB5[(Databricks)]
BigQuery --> DB6[(BigQuery API)]
Snowflake --> DB7[(Snowflake API)]
DB --> Postgres
DB2 --> MySQL
DB3 --> Redshift
DB4 --> SQLServer
DB5 --> Databricks
DB6 --> BigQuery
DB7 --> Snowflake
Postgres --> QE
MySQL --> QE
Redshift --> QE
SQLServer --> QE
Databricks --> QE
BigQuery --> QE
Snowflake --> QE
QE --> Client
```
### Core Components ✅
This refactoring primarily modifies existing components (database-specific query handlers). No new core components are introduced.
### Database Changes (If applicable)
N/A - No schema changes are required.
### API Changes (If applicable)
N/A - The public function signatures of the query handlers in `query_engine` remain unchanged.
### File Changes (If applicable) ✅
#### Modified Files
* `libs/query_engine/src/data_source_query_routes/postgres_query.rs`
* *Purpose:* Apply query limit at the database level using `LIMIT $1`. Remove post-fetch limiting. Refactor row processing (`process_batch`) to eliminate per-row tasks. Optimize type handling. Review/optimize/remove SQL parsing for identifier quoting.
* `libs/query_engine/src/data_source_query_routes/mysql_query.rs`
* *Purpose:* Apply query limit at the database level using `LIMIT ?`. Remove post-fetch limiting and hardcoded limit check. Remove per-row task spawning; process rows sequentially. Optimize type handling.
* `libs/query_engine/src/data_source_query_routes/redshift_query.rs`
* *Purpose:* Apply query limit at the database level using `LIMIT $1`. Remove post-fetch limiting and hardcoded limit check. Optimize type handling.
* `libs/query_engine/src/data_source_query_routes/sql_server_query.rs`
* *Purpose:* Apply query limit at the database level using parameter binding (`FETCH FIRST` or `TOP`). Remove post-fetch limiting (`.take()`). Remove per-row task spawning; process rows sequentially. Optimize type handling.
* `libs/query_engine/src/data_source_query_routes/databricks_query.rs`
* *Purpose:* Apply query limit directly in the submitted query string. Remove post-fetch limiting (`.take()`). Optimize type handling.
* `libs/query_engine/src/data_source_query_routes/bigquery_query.rs`
* *Purpose:* Review existing implementation (limit already applied correctly) for potential micro-optimizations in type parsing (`parse_string_to_datatype`, `parse_number_to_datatype`).
* `libs/query_engine/src/data_source_query_routes/snowflake_query.rs`
* *Purpose:* Review existing implementation (limit already applied correctly) for potential micro-optimizations in Arrow data processing and type handling (`match column.data_type()`).
#### New Files
N/A
#### Deleted Files
N/A
## 4. Implementation Plan ✅
### Phase 1: Database-Level Limit Implementation
* **Goal:** Ensure all connectors apply the 5000 default limit (or override) at the database level.
* **Tasks:** Modify `postgres_query.rs`, `mysql_query.rs`, `redshift_query.rs`, `sql_server_query.rs`, `databricks_query.rs` to apply limits via query modification/parameters. Remove post-fetch limit logic.
* **Success Criteria:** Unit/integration tests pass confirming correct limit behavior for default and override cases across all modified connectors. Code review confirms removal of post-fetch logic.
* **Dependencies:** None.
### Phase 2: Row Processing & Type Conversion Optimization
* **Goal:** Improve performance and memory efficiency by optimizing result processing.
* **Tasks:** Refactor `postgres_query.rs`, `mysql_query.rs`, `sql_server_query.rs` to remove per-row task spawning. Optimize type conversion logic across all relevant connectors (investigate alternatives to string matching). Review and optimize/remove Postgres SQL parsing. Review BigQuery and Snowflake implementations for micro-optimizations.
* **Success Criteria:** Benchmark tests show significant improvement in latency and memory usage for affected connectors. Code review confirms removal of per-row tasks and optimized type handling. Postgres parsing overhead is eliminated or demonstrably negligible.
* **Dependencies:** Phase 1 (optional, can be done concurrently but limit changes are higher priority).
## 5. Testing Strategy ✅
* **Unit Tests:**
* For each modified connector, add/update unit tests verifying the correct construction of the limited SQL query or API request based on the `limit` parameter (None, Some(N)).
* Test helper functions used for type conversion or processing.
* **Integration Tests:**
* Set up test databases/endpoints for each connector type.
* Create test tables with > 5000 rows.
* Run queries against test tables:
* Without explicit limit (assert exactly 5000 rows returned).
* With `limit: Some(100)` (assert exactly 100 rows returned).
* With `limit: Some(7000)` (assert exactly 7000 rows returned, if supported by test setup).
* With `limit: Some(0)` (assert 0 rows returned).
* Verify data integrity and types of returned data.
* **Benchmarking:**
* Establish baseline performance (latency, peak memory) for queries returning large datasets (>5000 rows) on affected connectors (Postgres, MySQL, SQL Server, Redshift) *before* changes.
* Re-run benchmarks *after* Phase 1 and Phase 2 changes.
* Compare results to validate performance improvements (Success Criteria in NFR1).
## 6. Security Considerations ✅
* The primary changes involve modifying query construction (adding LIMIT/FETCH) and optimizing internal processing.
* Ensure that query modification logic does not introduce SQL injection vulnerabilities (continue using parameter binding provided by `sqlx` and other clients).
* Changes to row processing should not expose sensitive data or alter existing security contexts.
## 7. Monitoring & Rollback ✅
* **Monitoring:**
* Monitor application-level metrics for query latency (average, p95) potentially filterable by data source type if available.
* Monitor application server CPU and memory usage.
* Observe database-level monitoring for query performance changes if accessible.
* **Rollback Plan:**
* Changes will be deployed via standard Git workflow.
* If significant issues (major performance degradation, incorrect results) are detected post-deployment, revert the relevant commit(s) and redeploy the previous version.
## 8. Dependencies & Files ✅
* **Modified Files:** (Listed again for clarity)
* `libs/query_engine/src/data_source_query_routes/postgres_query.rs`
* `libs/query_engine/src/data_source_query_routes/mysql_query.rs`
* `libs/query_engine/src/data_source_query_routes/redshift_query.rs`
* `libs/query_engine/src/data_source_query_routes/sql_server_query.rs`
* `libs/query_engine/src/data_source_query_routes/databricks_query.rs`
* `libs/query_engine/src/data_source_query_routes/bigquery_query.rs`
* `libs/query_engine/src/data_source_query_routes/snowflake_query.rs`
* **Dependencies:** Requires underlying database drivers/clients (`sqlx`, `gcp_bigquery_client`, `snowflake_api`, `tiberius`, etc.).
## 9. Open Questions
* Is the performance overhead of the `sqlparser` dependency in the Postgres connector significant enough to warrant removal/complex optimization, or is focusing on limit application and row processing sufficient? (Requires benchmarking/profiling in Phase 2).
* Are there more efficient type identification methods available in `sqlx` (Postgres, MySQL, Redshift), `tiberius` (SQL Server) than string matching type names?
## 10. Future Considerations
* **Streaming API:** Explore modifying the `query_engine` interface to return a stream (`impl Stream<Item = Result<IndexMap<String, DataType>>>`) instead of `Vec` for potential further memory savings and lower latency-to-first-result, especially for consumers that can process results incrementally.
* **Advanced Type Caching:** Implement more sophisticated caching for type information if basic optimizations in Phase 2 prove insufficient.
## 11. References
* `libs/query_engine/`
* `prds/template.md`
* `prds.mdc` rule