mirror of
https://github.com/ruvnet/RuView
synced 2026-08-10 20:31:42 +00:00
feat: vendor midstream and sublinear-time-solver libraries (#109)
Add ruvnet/midstream (AIMDS real-time inference) and ruvnet/sublinear-time-solver (sublinear optimization algorithms) as vendored dependencies under vendor/.
This commit is contained in:
rUv
+210
@@ -0,0 +1,210 @@
|
||||
//! Core aggregation framework for time-series data.
|
||||
//!
|
||||
//! This module provides the foundational types and functionality for aggregating
|
||||
//! time-series data. It defines:
|
||||
//! - Generic aggregation functions (Sum, Avg, Min, Max, Count)
|
||||
//! - Time window specifications (None, Fixed, Sliding)
|
||||
//! - Grouping operations
|
||||
//! - SQL query generation
|
||||
//!
|
||||
//! This framework is used by more specific aggregation implementations, such as
|
||||
//! the metric-specific aggregation in `crate::metrics::aggregation`.
|
||||
|
||||
use std::time::Duration;
|
||||
use serde::{Serialize, Deserialize};
|
||||
use std::fmt::{Display, Formatter};
|
||||
|
||||
/// Time window for aggregation
|
||||
#[derive(Debug, Clone, Copy, Serialize, Deserialize)]
|
||||
pub enum TimeWindow {
|
||||
/// No time window, aggregate all data
|
||||
None,
|
||||
/// Fixed time window (e.g., 5 minutes, 1 hour)
|
||||
Fixed(Duration),
|
||||
/// Sliding time window with window size and slide interval
|
||||
Sliding {
|
||||
window: Duration,
|
||||
slide: Duration,
|
||||
},
|
||||
}
|
||||
|
||||
/// Generic aggregation functions that can be applied to time-series data
|
||||
#[derive(Debug, Clone, Copy, Serialize, Deserialize)]
|
||||
pub enum AggregateFunction {
|
||||
/// Count the number of values
|
||||
Count,
|
||||
/// Sum all values
|
||||
Sum,
|
||||
/// Calculate the average
|
||||
Avg,
|
||||
/// Find the minimum value
|
||||
Min,
|
||||
/// Find the maximum value
|
||||
Max,
|
||||
}
|
||||
|
||||
impl Display for AggregateFunction {
|
||||
fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
|
||||
match self {
|
||||
AggregateFunction::Count => write!(f, "COUNT"),
|
||||
AggregateFunction::Sum => write!(f, "SUM"),
|
||||
AggregateFunction::Avg => write!(f, "AVG"),
|
||||
AggregateFunction::Min => write!(f, "MIN"),
|
||||
AggregateFunction::Max => write!(f, "MAX"),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Grouping specification for aggregation
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct GroupBy {
|
||||
pub columns: Vec<String>,
|
||||
pub time_column: Option<String>,
|
||||
}
|
||||
|
||||
/// Result of an aggregation operation
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct AggregateResult {
|
||||
pub value: f64,
|
||||
pub timestamp: i64,
|
||||
}
|
||||
|
||||
impl TimeWindow {
|
||||
/// Calculates the window boundaries for a given timestamp
|
||||
pub fn window_bounds(&self, timestamp: i64) -> (i64, i64) {
|
||||
match *self {
|
||||
TimeWindow::None => (i64::MIN, i64::MAX),
|
||||
TimeWindow::Fixed(duration) => {
|
||||
let window_size = duration.as_secs() as i64;
|
||||
let window_start = (timestamp / window_size) * window_size;
|
||||
(window_start, window_start + window_size)
|
||||
},
|
||||
TimeWindow::Sliding { window, slide } => {
|
||||
let window_size = window.as_secs() as i64;
|
||||
let slide_size = slide.as_secs() as i64;
|
||||
let current_slide = (timestamp / slide_size) * slide_size;
|
||||
(current_slide, current_slide + window_size)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Generates SQL expressions for window boundaries
|
||||
pub fn to_sql(&self) -> Option<String> {
|
||||
match *self {
|
||||
TimeWindow::None => None,
|
||||
TimeWindow::Fixed(duration) => {
|
||||
let window_size = duration.as_secs();
|
||||
Some(format!(
|
||||
"(timestamp / {}) * {} as window_start,
|
||||
((timestamp / {}) + 1) * {} as window_end",
|
||||
window_size, window_size, window_size, window_size
|
||||
))
|
||||
},
|
||||
TimeWindow::Sliding { window, slide } => {
|
||||
let window_size = window.as_secs();
|
||||
let slide_size = slide.as_secs();
|
||||
Some(format!(
|
||||
"(timestamp / {}) * {} as window_start,
|
||||
((timestamp / {}) * {} + {}) as window_end",
|
||||
slide_size, slide_size, slide_size, slide_size, window_size
|
||||
))
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl AggregateFunction {
|
||||
/// Generates SQL for the aggregation function
|
||||
pub fn to_sql(&self, column: &str) -> String {
|
||||
match self {
|
||||
AggregateFunction::Sum => format!("SUM({})", column),
|
||||
AggregateFunction::Avg => format!("AVG({})", column),
|
||||
AggregateFunction::Min => format!("MIN({})", column),
|
||||
AggregateFunction::Max => format!("MAX({})", column),
|
||||
AggregateFunction::Count => format!("COUNT({})", column),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Builds a SQL query for aggregation.
|
||||
///
|
||||
/// This is the core query builder used by specific aggregation implementations.
|
||||
/// It provides a flexible way to build SQL queries for different types of
|
||||
/// time-series data aggregation.
|
||||
///
|
||||
/// # Arguments
|
||||
///
|
||||
/// * `table_name` - The source table name
|
||||
/// * `function` - The aggregation function to apply
|
||||
/// * `group_by` - The grouping specification
|
||||
/// * `columns` - The columns to aggregate
|
||||
/// * `from_timestamp` - Optional start of the time range
|
||||
/// * `to_timestamp` - Optional end of the time range
|
||||
///
|
||||
/// # Returns
|
||||
///
|
||||
/// A SQL query string for the specified aggregation
|
||||
pub fn build_aggregate_query(
|
||||
table_name: &str,
|
||||
function: AggregateFunction,
|
||||
group_by: &GroupBy,
|
||||
columns: &[&str],
|
||||
from_timestamp: Option<i64>,
|
||||
to_timestamp: Option<i64>,
|
||||
) -> String {
|
||||
let mut query = String::new();
|
||||
|
||||
// Build SELECT clause
|
||||
query.push_str("SELECT ");
|
||||
|
||||
// Add group by columns
|
||||
if !group_by.columns.is_empty() {
|
||||
let cols: Vec<&str> = group_by.columns.iter().map(|s| s.as_str()).collect();
|
||||
query.push_str(&cols.join(", "));
|
||||
query.push_str(", ");
|
||||
}
|
||||
|
||||
// Add time column if present
|
||||
if let Some(time_col) = &group_by.time_column {
|
||||
query.push_str(&format!("{}, ", time_col));
|
||||
}
|
||||
|
||||
// Add aggregation function
|
||||
match function {
|
||||
AggregateFunction::Sum => query.push_str("SUM(value)"),
|
||||
AggregateFunction::Avg => query.push_str("AVG(value)"),
|
||||
AggregateFunction::Count => query.push_str("COUNT(*)"),
|
||||
AggregateFunction::Min => query.push_str("MIN(value)"),
|
||||
AggregateFunction::Max => query.push_str("MAX(value)"),
|
||||
}
|
||||
|
||||
// Add FROM clause
|
||||
query.push_str(&format!(" FROM {}", table_name));
|
||||
|
||||
// Add WHERE clause for timestamp range
|
||||
if let Some(from_ts) = from_timestamp {
|
||||
query.push_str(&format!(" WHERE timestamp >= {}", from_ts));
|
||||
if let Some(to_ts) = to_timestamp {
|
||||
query.push_str(&format!(" AND timestamp <= {}", to_ts));
|
||||
}
|
||||
}
|
||||
|
||||
// Add GROUP BY clause
|
||||
if !group_by.columns.is_empty() || group_by.time_column.is_some() {
|
||||
query.push_str(" GROUP BY ");
|
||||
let mut group_cols = Vec::new();
|
||||
|
||||
if !group_by.columns.is_empty() {
|
||||
let cols: Vec<&str> = group_by.columns.iter().map(|s| s.as_str()).collect();
|
||||
group_cols.extend(cols);
|
||||
}
|
||||
|
||||
if let Some(time_col) = &group_by.time_column {
|
||||
group_cols.push(time_col.as_str());
|
||||
}
|
||||
|
||||
query.push_str(&group_cols.join(", "));
|
||||
}
|
||||
|
||||
query
|
||||
}
|
||||
+247
@@ -0,0 +1,247 @@
|
||||
//! Hyprstream server binary.
|
||||
//!
|
||||
//! This binary provides the main entry point for the Hyprstream service, a next-generation application
|
||||
//! for real-time data ingestion, windowed aggregation, caching, and serving.
|
||||
//!
|
||||
//! # Features
|
||||
//!
|
||||
//! - **Data Ingestion**: Ingest data efficiently using Arrow Flight
|
||||
//! - **Intelligent Caching**: High-performance caching with DuckDB
|
||||
//! - **Real-time Aggregation**: Dynamic metrics and time-windowed aggregates
|
||||
//! - **ADBC Integration**: Seamless connection to external databases
|
||||
//!
|
||||
//! # Configuration
|
||||
//!
|
||||
//! Configuration can be provided through multiple sources, in order of precedence:
|
||||
//!
|
||||
//! 1. Command-line arguments (highest precedence)
|
||||
//! 2. Environment variables (prefixed with `HYPRSTREAM_`)
|
||||
//! 3. User-specified configuration file (via `--config`)
|
||||
//! 4. System-wide configuration (`/etc/hyprstream/config.toml`)
|
||||
//! 5. Default configuration (embedded in binary)
|
||||
//!
|
||||
//! ## Command-line Options
|
||||
//!
|
||||
//! ```text
|
||||
//! Options:
|
||||
//! -c, --config <FILE> Path to configuration file
|
||||
//! --host <HOST> Server host address [env: HYPRSTREAM_SERVER_HOST]
|
||||
//! --port <PORT> Server port [env: HYPRSTREAM_SERVER_PORT]
|
||||
//! --engine <TYPE> Primary storage engine type [env: HYPRSTREAM_ENGINE]
|
||||
//! --engine-connection <STR> Engine connection string [env: HYPRSTREAM_ENGINE_CONNECTION]
|
||||
//! --engine-options <KEY=VAL> Engine options (can be specified multiple times) [env: HYPRSTREAM_ENGINE_OPTIONS]
|
||||
//! --enable-cache Enable caching [env: HYPRSTREAM_ENABLE_CACHE]
|
||||
//! --cache-engine <TYPE> Cache engine type [env: HYPRSTREAM_CACHE_ENGINE]
|
||||
//! --cache-connection <STR> Cache connection string [env: HYPRSTREAM_CACHE_CONNECTION]
|
||||
//! --cache-options <KEY=VAL> Cache options (can be specified multiple times) [env: HYPRSTREAM_CACHE_OPTIONS]
|
||||
//! --cache-max-duration <SECS> Cache max duration in seconds [env: HYPRSTREAM_CACHE_MAX_DURATION]
|
||||
//! ```
|
||||
//!
|
||||
//! ## Configuration File Format (TOML)
|
||||
//!
|
||||
//! ```toml
|
||||
//! # Server Configuration
|
||||
//! [server]
|
||||
//! host = "127.0.0.1" # Server host address
|
||||
//! port = 50051 # Server port number
|
||||
//!
|
||||
//! # Engine Configuration
|
||||
//! [engine]
|
||||
//! engine = "duckdb" # Options: "duckdb", "adbc"
|
||||
//! connection = ":memory:"
|
||||
//! options = { } # Engine-specific options
|
||||
//!
|
||||
//! # Cache Configuration
|
||||
//! [cache]
|
||||
//! enabled = true # Enable caching
|
||||
//! engine = "duckdb" # Options: "duckdb", "adbc"
|
||||
//! connection = ":memory:"
|
||||
//! max_duration_secs = 3600
|
||||
//! options = { } # Cache-specific options
|
||||
//! ```
|
||||
//!
|
||||
//! ## Storage Backends
|
||||
//!
|
||||
//! ### DuckDB Backend
|
||||
//!
|
||||
//! The DuckDB backend provides high-performance embedded storage:
|
||||
//!
|
||||
//! ```toml
|
||||
//! [engine]
|
||||
//! engine = "duckdb"
|
||||
//! connection = ":memory:" # Use ":memory:" for in-memory or file path
|
||||
//! options = {
|
||||
//! threads = "4", # Number of threads (optional)
|
||||
//! read_only = "false" # Read-only mode (optional)
|
||||
//! }
|
||||
//! ```
|
||||
//!
|
||||
//! ### ADBC Backend
|
||||
//!
|
||||
//! The ADBC backend enables connection to external databases:
|
||||
//!
|
||||
//! ```toml
|
||||
//! [engine]
|
||||
//! engine = "adbc"
|
||||
//! connection = "postgresql://localhost:5432"
|
||||
//! options = {
|
||||
//! driver_path = "/usr/local/lib/libadbc_driver_postgresql.so", # Required
|
||||
//! username = "postgres", # Optional
|
||||
//! password = "secret", # Optional
|
||||
//! database = "metrics", # Optional
|
||||
//! pool_max = "10", # Optional
|
||||
//! pool_min = "1", # Optional
|
||||
//! connect_timeout = "30" # Optional
|
||||
//! }
|
||||
//! ```
|
||||
//!
|
||||
//! ### Cached Backend
|
||||
//!
|
||||
//! The cached backend implements a two-tier storage system:
|
||||
//!
|
||||
//! ```toml
|
||||
//! [engine]
|
||||
//! engine = "duckdb"
|
||||
//! connection = "data.db"
|
||||
//! options = { }
|
||||
//!
|
||||
//! [cache]
|
||||
//! enabled = true
|
||||
//! engine = "duckdb"
|
||||
//! connection = ":memory:"
|
||||
//! max_duration_secs = 3600
|
||||
//! options = {
|
||||
//! threads = "2"
|
||||
//! }
|
||||
//! ```
|
||||
//!
|
||||
//! # Examples
|
||||
//!
|
||||
//! ## Basic Usage
|
||||
//!
|
||||
//! ```bash
|
||||
//! # Run with default configuration
|
||||
//! hyprstream
|
||||
//!
|
||||
//! # Run with custom configuration file
|
||||
//! hyprstream --config /path/to/config.toml
|
||||
//!
|
||||
//! # Run with environment variables
|
||||
//! HYPRSTREAM_SERVER_HOST=0.0.0.0 HYPRSTREAM_SERVER_PORT=50051 hyprstream
|
||||
//! ```
|
||||
//!
|
||||
//! ## Advanced Configuration
|
||||
//!
|
||||
//! ```bash
|
||||
//! # Run with DuckDB storage and caching
|
||||
//! hyprstream \
|
||||
//! --engine duckdb \
|
||||
//! --engine-connection ":memory:" \
|
||||
//! --engine-options threads=4 \
|
||||
//! --enable-cache \
|
||||
//! --cache-engine duckdb \
|
||||
//! --cache-connection ":memory:" \
|
||||
//! --cache-max-duration 3600
|
||||
//!
|
||||
//! # Run with ADBC PostgreSQL backend
|
||||
//! hyprstream \
|
||||
//! --engine adbc \
|
||||
//! --engine-connection "postgresql://localhost:5432" \
|
||||
//! --engine-options driver_path=/usr/local/lib/libadbc_driver_postgresql.so \
|
||||
//! --engine-options username=postgres \
|
||||
//! --engine-options database=metrics \
|
||||
//! --engine-options pool_max=10
|
||||
//!
|
||||
//! # Run with ADBC backend and DuckDB cache
|
||||
//! hyprstream \
|
||||
//! --engine adbc \
|
||||
//! --engine-connection "postgresql://localhost:5432" \
|
||||
//! --engine-options driver_path=/usr/local/lib/libadbc_driver_postgresql.so \
|
||||
//! --engine-options username=postgres \
|
||||
//! --enable-cache \
|
||||
//! --cache-engine duckdb \
|
||||
//! --cache-connection ":memory:" \
|
||||
//! --cache-options threads=2 \
|
||||
//! --cache-max-duration 3600
|
||||
//! ```
|
||||
//!
|
||||
//! For more examples and detailed API documentation, visit the
|
||||
//! [Hyprstream documentation](https://docs.rs/hyprstream).
|
||||
|
||||
use clap::Parser;
|
||||
use hyprstream_core::{
|
||||
config::{CliArgs, Settings},
|
||||
service::FlightSqlService,
|
||||
storage::{StorageBackend, adbc::AdbcBackend, duckdb::DuckDbBackend},
|
||||
};
|
||||
use std::sync::Arc;
|
||||
use tonic::transport::Server;
|
||||
|
||||
#[tokio::main]
|
||||
async fn main() -> Result<(), Box<dyn std::error::Error>> {
|
||||
let cli_args = CliArgs::parse();
|
||||
|
||||
// Load settings from config file and CLI args
|
||||
let settings = Settings::new(cli_args)?;
|
||||
|
||||
// Create the storage backend based on configuration
|
||||
let engine_backend: Box<dyn StorageBackend> = match settings.engine.engine.as_str() {
|
||||
"adbc" => {
|
||||
Box::new(AdbcBackend::new_with_options(
|
||||
&settings.engine.connection,
|
||||
&settings.engine.options,
|
||||
settings.engine.credentials.as_ref(),
|
||||
)?)
|
||||
}
|
||||
"duckdb" => {
|
||||
Box::new(DuckDbBackend::new_with_options(
|
||||
&settings.engine.connection,
|
||||
&settings.engine.options,
|
||||
settings.engine.credentials.as_ref(),
|
||||
)?)
|
||||
}
|
||||
_ => return Err("Unsupported engine type".into()),
|
||||
};
|
||||
|
||||
// Initialize the storage backend
|
||||
engine_backend.init().await?;
|
||||
|
||||
// Create cache backend if configured
|
||||
let cache_backend = if settings.cache.enabled {
|
||||
let cache_config = &settings.cache;
|
||||
let backend: Box<dyn StorageBackend> = match cache_config.engine.as_str() {
|
||||
"adbc" => {
|
||||
Box::new(AdbcBackend::new_with_options(
|
||||
&cache_config.connection,
|
||||
&cache_config.options,
|
||||
cache_config.credentials.as_ref(),
|
||||
)?)
|
||||
}
|
||||
"duckdb" => {
|
||||
Box::new(DuckDbBackend::new_with_options(
|
||||
&cache_config.connection,
|
||||
&cache_config.options,
|
||||
cache_config.credentials.as_ref(),
|
||||
)?)
|
||||
}
|
||||
_ => return Err("Unsupported cache engine type".into()),
|
||||
};
|
||||
Some(backend)
|
||||
} else {
|
||||
None
|
||||
};
|
||||
|
||||
// Create the Flight SQL service
|
||||
let service = FlightSqlService::new(engine_backend);
|
||||
|
||||
// Start the server
|
||||
let addr = format!("{}:{}", settings.server.host, settings.server.port).parse()?;
|
||||
println!("Starting server on {}", addr);
|
||||
|
||||
Server::builder()
|
||||
.add_service(arrow_flight::flight_service_server::FlightServiceServer::new(service))
|
||||
.serve(addr)
|
||||
.await?;
|
||||
|
||||
Ok(())
|
||||
}
|
||||
+329
@@ -0,0 +1,329 @@
|
||||
//! Configuration management for Hyprstream service.
|
||||
//!
|
||||
//! This module provides configuration handling through multiple sources:
|
||||
//! 1. Default configuration (embedded in binary)
|
||||
//! 2. System-wide configuration file (`/etc/hyprstream/config.toml`)
|
||||
//! 3. User-specified configuration file
|
||||
//! 4. Environment variables (prefixed with `HYPRSTREAM_`)
|
||||
//! 5. Command-line arguments
|
||||
//!
|
||||
//! Configuration options are loaded in order of precedence, with later sources
|
||||
//! overriding earlier ones.
|
||||
//!
|
||||
//! # Environment Variables
|
||||
//!
|
||||
//! Backend-specific credentials should be provided via environment variables:
|
||||
//! - `HYPRSTREAM_ENGINE_USERNAME` - Primary storage backend username
|
||||
//! - `HYPRSTREAM_ENGINE_PASSWORD` - Primary storage backend password
|
||||
//! - `HYPRSTREAM_CACHE_USERNAME` - Cache backend username (if needed)
|
||||
//! - `HYPRSTREAM_CACHE_PASSWORD` - Cache backend password (if needed)
|
||||
|
||||
use clap::Parser;
|
||||
use config::{Config, ConfigError};
|
||||
use serde::Deserialize;
|
||||
use std::env;
|
||||
use std::path::PathBuf;
|
||||
use std::collections::HashMap;
|
||||
|
||||
const DEFAULT_CONFIG: &str = include_str!("../config/default.toml");
|
||||
const DEFAULT_CONFIG_PATH: &str = "/etc/hyprstream/config.toml";
|
||||
|
||||
/// Command-line arguments parser.
|
||||
///
|
||||
/// This structure defines all available command-line options and their
|
||||
/// corresponding environment variables. It uses clap for parsing and
|
||||
/// supports both short and long option forms.
|
||||
#[derive(Parser, Debug)]
|
||||
#[command(author, version, about)]
|
||||
pub struct CliArgs {
|
||||
/// Path to the configuration file
|
||||
#[arg(short, long, value_name = "FILE")]
|
||||
config: Option<PathBuf>,
|
||||
|
||||
/// Server host address
|
||||
#[arg(long, env = "HYPRSTREAM_SERVER_HOST")]
|
||||
host: Option<String>,
|
||||
|
||||
/// Server port
|
||||
#[arg(long, env = "HYPRSTREAM_SERVER_PORT")]
|
||||
port: Option<u16>,
|
||||
|
||||
/// Primary storage engine type
|
||||
#[arg(long, env = "HYPRSTREAM_ENGINE")]
|
||||
engine: Option<String>,
|
||||
|
||||
/// Primary storage engine connection string
|
||||
#[arg(long, env = "HYPRSTREAM_ENGINE_CONNECTION")]
|
||||
engine_connection: Option<String>,
|
||||
|
||||
/// Primary storage engine options (key=value pairs)
|
||||
#[arg(long, env = "HYPRSTREAM_ENGINE_OPTIONS")]
|
||||
engine_options: Option<Vec<String>>,
|
||||
|
||||
/// Enable caching
|
||||
#[arg(long, env = "HYPRSTREAM_ENABLE_CACHE")]
|
||||
enable_cache: Option<bool>,
|
||||
|
||||
/// Cache engine type
|
||||
#[arg(long, env = "HYPRSTREAM_CACHE_ENGINE")]
|
||||
cache_engine: Option<String>,
|
||||
|
||||
/// Cache engine connection string
|
||||
#[arg(long, env = "HYPRSTREAM_CACHE_CONNECTION")]
|
||||
cache_connection: Option<String>,
|
||||
|
||||
/// Cache engine options (key=value pairs)
|
||||
#[arg(long, env = "HYPRSTREAM_CACHE_OPTIONS")]
|
||||
cache_options: Option<Vec<String>>,
|
||||
|
||||
/// Cache maximum duration in seconds
|
||||
#[arg(long, env = "HYPRSTREAM_CACHE_MAX_DURATION")]
|
||||
cache_max_duration: Option<u64>,
|
||||
|
||||
/// Primary storage engine username
|
||||
#[arg(long, env = "HYPRSTREAM_ENGINE_USERNAME")]
|
||||
engine_username: Option<String>,
|
||||
|
||||
/// Primary storage engine password
|
||||
#[arg(long, env = "HYPRSTREAM_ENGINE_PASSWORD")]
|
||||
engine_password: Option<String>,
|
||||
|
||||
/// Cache engine username
|
||||
#[arg(long, env = "HYPRSTREAM_CACHE_USERNAME")]
|
||||
cache_username: Option<String>,
|
||||
|
||||
/// Cache engine password
|
||||
#[arg(long, env = "HYPRSTREAM_CACHE_PASSWORD")]
|
||||
cache_password: Option<String>,
|
||||
}
|
||||
|
||||
/// Complete service configuration.
|
||||
///
|
||||
/// This structure holds all configuration options for the service,
|
||||
/// including server settings, storage backend configuration, and
|
||||
/// cache settings.
|
||||
#[derive(Debug, Deserialize)]
|
||||
pub struct Settings {
|
||||
/// Server configuration
|
||||
pub server: ServerConfig,
|
||||
/// Engine configuration
|
||||
pub engine: EngineConfig,
|
||||
/// Cache configuration
|
||||
pub cache: CacheConfig,
|
||||
}
|
||||
|
||||
/// Server configuration options.
|
||||
///
|
||||
/// Defines the network interface and port for the Flight SQL service.
|
||||
#[derive(Debug, Deserialize)]
|
||||
pub struct ServerConfig {
|
||||
/// Host address to bind to
|
||||
pub host: String,
|
||||
/// Port number to listen on
|
||||
pub port: u16,
|
||||
}
|
||||
|
||||
/// Engine configuration.
|
||||
///
|
||||
/// Specifies the primary storage engine to use for metric data.
|
||||
#[derive(Debug, Deserialize)]
|
||||
pub struct EngineConfig {
|
||||
/// Engine type ("duckdb" or "adbc")
|
||||
pub engine: String,
|
||||
/// Connection string for the engine
|
||||
pub connection: String,
|
||||
/// Engine-specific options
|
||||
#[serde(default)]
|
||||
pub options: std::collections::HashMap<String, String>,
|
||||
/// Authentication credentials (not serialized)
|
||||
#[serde(skip)]
|
||||
pub credentials: Option<Credentials>,
|
||||
}
|
||||
|
||||
/// Authentication credentials for storage backends.
|
||||
#[derive(Debug, Clone, Deserialize)]
|
||||
pub struct Credentials {
|
||||
/// Username for authentication
|
||||
pub username: String,
|
||||
/// Password for authentication
|
||||
pub password: String,
|
||||
}
|
||||
|
||||
/// Cache configuration for the storage backend.
|
||||
#[derive(Debug, Deserialize)]
|
||||
pub struct CacheConfig {
|
||||
/// Whether caching is enabled
|
||||
pub enabled: bool,
|
||||
/// Cache storage engine type (e.g., "duckdb", "adbc")
|
||||
pub engine: String,
|
||||
/// Cache connection string
|
||||
pub connection: String,
|
||||
/// Cache engine options
|
||||
pub options: HashMap<String, String>,
|
||||
/// Cache credentials (optional)
|
||||
#[serde(default)]
|
||||
pub credentials: Option<Credentials>,
|
||||
/// Maximum duration to keep entries in cache (in seconds)
|
||||
#[serde(default = "default_ttl")]
|
||||
pub ttl: Option<u64>,
|
||||
}
|
||||
|
||||
fn default_ttl() -> Option<u64> {
|
||||
Some(3600) // Default 1 hour TTL
|
||||
}
|
||||
|
||||
impl Default for CacheConfig {
|
||||
fn default() -> Self {
|
||||
Self {
|
||||
enabled: false,
|
||||
engine: "duckdb".to_string(),
|
||||
connection: ":memory:".to_string(),
|
||||
options: HashMap::new(),
|
||||
credentials: None,
|
||||
ttl: default_ttl(),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl Settings {
|
||||
/// Loads configuration from all available sources.
|
||||
pub fn new(cli: CliArgs) -> Result<Self, ConfigError> {
|
||||
let mut builder = Config::builder();
|
||||
|
||||
// Load default configuration
|
||||
builder = builder.add_source(config::File::from_str(
|
||||
DEFAULT_CONFIG,
|
||||
config::FileFormat::Toml,
|
||||
));
|
||||
|
||||
// Load system configuration if it exists
|
||||
if let Ok(metadata) = std::fs::metadata(DEFAULT_CONFIG_PATH) {
|
||||
if metadata.is_file() {
|
||||
builder = builder.add_source(config::File::from(PathBuf::from(DEFAULT_CONFIG_PATH)));
|
||||
}
|
||||
}
|
||||
|
||||
// Load user configuration if specified
|
||||
if let Some(ref config_path) = cli.config {
|
||||
builder = builder.add_source(config::File::from(config_path.clone()));
|
||||
}
|
||||
|
||||
// Add environment variables (prefixed with HYPRSTREAM_)
|
||||
builder = builder.add_source(config::Environment::with_prefix("HYPRSTREAM"));
|
||||
|
||||
// Override with command line arguments
|
||||
if let Some(ref host) = cli.host {
|
||||
builder = builder.set_override("server.host", host.as_str())?;
|
||||
}
|
||||
if let Some(port) = cli.port {
|
||||
builder = builder.set_override("server.port", port)?;
|
||||
}
|
||||
|
||||
// Engine settings
|
||||
if let Some(ref engine) = cli.engine {
|
||||
builder = builder.set_override("engine.engine", engine.as_str())?;
|
||||
}
|
||||
if let Some(ref connection) = cli.engine_connection {
|
||||
builder = builder.set_override("engine.connection", connection.as_str())?;
|
||||
}
|
||||
if let Some(ref options) = cli.engine_options {
|
||||
let options: std::collections::HashMap<String, String> = options
|
||||
.iter()
|
||||
.filter_map(|opt| {
|
||||
let parts: Vec<&str> = opt.split('=').collect();
|
||||
if parts.len() == 2 {
|
||||
Some((parts[0].to_string(), parts[1].to_string()))
|
||||
} else {
|
||||
None
|
||||
}
|
||||
})
|
||||
.collect();
|
||||
builder = builder.set_override("engine.options", options)?;
|
||||
}
|
||||
|
||||
// Cache settings
|
||||
if let Some(enabled) = cli.enable_cache {
|
||||
builder = builder.set_override("cache.enabled", enabled)?;
|
||||
}
|
||||
if let Some(ref engine) = cli.cache_engine {
|
||||
builder = builder.set_override("cache.engine", engine.as_str())?;
|
||||
}
|
||||
if let Some(ref connection) = cli.cache_connection {
|
||||
builder = builder.set_override("cache.connection", connection.as_str())?;
|
||||
}
|
||||
if let Some(ref options) = cli.cache_options {
|
||||
let options: std::collections::HashMap<String, String> = options
|
||||
.iter()
|
||||
.filter_map(|opt| {
|
||||
let parts: Vec<&str> = opt.split('=').collect();
|
||||
if parts.len() == 2 {
|
||||
Some((parts[0].to_string(), parts[1].to_string()))
|
||||
} else {
|
||||
None
|
||||
}
|
||||
})
|
||||
.collect();
|
||||
builder = builder.set_override("cache.options", options)?;
|
||||
}
|
||||
if let Some(duration) = cli.cache_max_duration {
|
||||
builder = builder.set_override("cache.max_duration_secs", duration)?;
|
||||
}
|
||||
|
||||
// Build initial settings
|
||||
let mut settings: Settings = builder.build()?.try_deserialize()?;
|
||||
|
||||
// Load credentials from environment variables and CLI args
|
||||
settings.engine.credentials = Self::load_engine_credentials(&cli);
|
||||
settings.cache.credentials = Self::load_cache_credentials(&cli);
|
||||
|
||||
Ok(settings)
|
||||
}
|
||||
|
||||
/// Load engine credentials from all available sources.
|
||||
/// Priority order (highest to lowest):
|
||||
/// 1. Environment variables
|
||||
/// 2. Command line arguments
|
||||
fn load_engine_credentials(cli: &CliArgs) -> Option<Credentials> {
|
||||
// Try environment variables first
|
||||
if let (Some(username), Some(password)) = (
|
||||
env::var("HYPRSTREAM_ENGINE_USERNAME").ok(),
|
||||
env::var("HYPRSTREAM_ENGINE_PASSWORD").ok(),
|
||||
) {
|
||||
return Some(Credentials { username, password });
|
||||
}
|
||||
|
||||
// Try command line arguments
|
||||
if let (Some(username), Some(password)) = (&cli.engine_username, &cli.engine_password) {
|
||||
return Some(Credentials {
|
||||
username: username.clone(),
|
||||
password: password.clone(),
|
||||
});
|
||||
}
|
||||
|
||||
None
|
||||
}
|
||||
|
||||
/// Load cache credentials from all available sources.
|
||||
/// Priority order (highest to lowest):
|
||||
/// 1. Environment variables
|
||||
/// 2. Command line arguments
|
||||
fn load_cache_credentials(cli: &CliArgs) -> Option<Credentials> {
|
||||
// Try environment variables first
|
||||
if let (Some(username), Some(password)) = (
|
||||
env::var("HYPRSTREAM_CACHE_USERNAME").ok(),
|
||||
env::var("HYPRSTREAM_CACHE_PASSWORD").ok(),
|
||||
) {
|
||||
return Some(Credentials { username, password });
|
||||
}
|
||||
|
||||
// Try command line arguments
|
||||
if let (Some(username), Some(password)) = (&cli.cache_username, &cli.cache_password) {
|
||||
return Some(Credentials {
|
||||
username: username.clone(),
|
||||
password: password.clone(),
|
||||
});
|
||||
}
|
||||
|
||||
None
|
||||
}
|
||||
}
|
||||
+83
@@ -0,0 +1,83 @@
|
||||
/*!
|
||||
# Hyprstream: Real-time Aggregation Windows and High-Performance Cache for Apache Arrow Flight SQL
|
||||
|
||||
Hyprstream is a next-generation application for real-time data ingestion, windowed aggregation, caching, and serving.
|
||||
Built on Apache Arrow Flight and DuckDB, and developed in Rust, Hyprstream dynamically calculates metrics like running
|
||||
sums, counts, and averages, enabling blazing-fast data workflows, intelligent caching, and seamless integration with
|
||||
ADBC-compliant datastores.
|
||||
|
||||
## Key Features
|
||||
|
||||
### Data Ingestion via Apache Arrow Flight
|
||||
- Streamlined ingestion using Arrow Flight for efficient columnar data transport
|
||||
- Real-time streaming support for metrics, datasets, and vectorized data
|
||||
- Seamless integration with data producers for high-throughput ingestion
|
||||
- Write-through to ADBC datastores for eventual data consistency
|
||||
|
||||
### Intelligent Read Caching with DuckDB
|
||||
- In-memory performance using DuckDB for lightning-fast caching
|
||||
- Optimized querying for analytics workloads
|
||||
- Automatic cache management with configurable expiry policies
|
||||
- Time-based expiry with future support for LRU/LFU policies
|
||||
|
||||
### Data Serving with Arrow Flight SQL
|
||||
- High-performance queries via Arrow Flight SQL
|
||||
- Support for vectorized data and analytical queries
|
||||
- Seamless integration with analytics and visualization tools
|
||||
|
||||
### Real-Time Aggregation
|
||||
- Dynamic metrics with running sums, counts, and averages
|
||||
- Lightweight state management for aggregate calculations
|
||||
- Dynamic weight computation for AI/ML pipelines
|
||||
- Time window partitioning for granular analysis
|
||||
|
||||
## Usage
|
||||
|
||||
Basic usage example with programmatic configuration:
|
||||
|
||||
```rust,no_run
|
||||
use hyprstream::config::{Settings, EngineConfig, CacheConfig};
|
||||
use hyprstream::service::FlightServiceImpl;
|
||||
use std::sync::Arc;
|
||||
use std::collections::HashMap;
|
||||
|
||||
#[tokio::main]
|
||||
async fn main() -> Result<(), Box<dyn std::error::Error>> {
|
||||
// Create configuration programmatically
|
||||
let mut settings = Settings::default();
|
||||
|
||||
// Configure primary storage engine
|
||||
settings.engine.engine = "duckdb".to_string();
|
||||
settings.engine.connection = ":memory:".to_string();
|
||||
settings.engine.options.insert("threads".to_string(), "4".to_string());
|
||||
|
||||
// Configure caching (optional)
|
||||
settings.cache.enabled = true;
|
||||
settings.cache.engine = "duckdb".to_string();
|
||||
settings.cache.connection = ":memory:".to_string();
|
||||
settings.cache.max_duration_secs = 3600;
|
||||
|
||||
// Create and initialize the service
|
||||
let service = FlightServiceImpl::from_settings(&settings).await?;
|
||||
|
||||
// Use the service in your application...
|
||||
Ok(())
|
||||
}
|
||||
```
|
||||
|
||||
For detailed configuration options and examples, see:
|
||||
- [`config`](crate::config) module for configuration options
|
||||
- [`storage`](crate::storage) module for storage backend details
|
||||
- [`examples`](examples/) directory for more usage examples
|
||||
*/
|
||||
|
||||
pub mod metrics;
|
||||
pub mod storage;
|
||||
pub mod service;
|
||||
pub mod config;
|
||||
pub mod aggregation;
|
||||
|
||||
pub use service::FlightSqlService;
|
||||
pub use storage::StorageBackend;
|
||||
pub use metrics::MetricRecord;
|
||||
pub use aggregation::{TimeWindow, AggregateFunction, GroupBy, AggregateResult};
|
||||
@@ -0,0 +1,97 @@
|
||||
//! Metric-specific aggregation functionality.
|
||||
//!
|
||||
//! This module provides metric-specific implementations of aggregation functions
|
||||
//! while leveraging the core aggregation framework from `crate::aggregation`.
|
||||
//! It handles the specific requirements of metric aggregation, such as:
|
||||
//! - Running window calculations (sum, avg, count)
|
||||
//! - Metric-specific SQL query generation
|
||||
//! - Direct aggregation of MetricRecord instances
|
||||
//!
|
||||
//! The implementation reuses the generic aggregation types and query building
|
||||
//! functionality from the core aggregation module while adding metric-specific
|
||||
//! logic and optimizations.
|
||||
|
||||
use crate::metrics::MetricRecord;
|
||||
use crate::aggregation::{AggregateFunction, GroupBy, build_aggregate_query};
|
||||
use tonic::Status;
|
||||
|
||||
/// The standard metric value columns used in aggregation queries
|
||||
const METRIC_VALUE_COLUMNS: [&str; 3] = [
|
||||
"value_running_window_sum",
|
||||
"value_running_window_avg",
|
||||
"value_running_window_count"
|
||||
];
|
||||
|
||||
/// Applies the aggregation function to a set of metrics.
|
||||
///
|
||||
/// This function implements metric-specific aggregation by operating directly
|
||||
/// on MetricRecord instances. It uses the appropriate running window value
|
||||
/// based on the aggregation function type.
|
||||
///
|
||||
/// # Arguments
|
||||
///
|
||||
/// * `function` - The aggregation function to apply
|
||||
/// * `metrics` - The set of metrics to aggregate
|
||||
///
|
||||
/// # Returns
|
||||
///
|
||||
/// The aggregated value as a float, or an error if the operation fails
|
||||
pub fn apply_function(function: AggregateFunction, metrics: &[MetricRecord]) -> Result<f64, Status> {
|
||||
if metrics.is_empty() {
|
||||
return Ok(0.0);
|
||||
}
|
||||
|
||||
match function {
|
||||
AggregateFunction::Sum => Ok(metrics.iter().map(|m| m.value_running_window_sum).sum()),
|
||||
AggregateFunction::Avg => {
|
||||
let sum: f64 = metrics.iter().map(|m| m.value_running_window_avg).sum();
|
||||
Ok(sum / metrics.len() as f64)
|
||||
},
|
||||
AggregateFunction::Min => Ok(metrics
|
||||
.iter()
|
||||
.map(|m| m.value_running_window_sum)
|
||||
.fold(f64::INFINITY, f64::min)),
|
||||
AggregateFunction::Max => Ok(metrics
|
||||
.iter()
|
||||
.map(|m| m.value_running_window_sum)
|
||||
.fold(f64::NEG_INFINITY, f64::max)),
|
||||
AggregateFunction::Count => Ok(metrics.len() as f64),
|
||||
}
|
||||
}
|
||||
|
||||
/// Builds a SQL query for metrics aggregation.
|
||||
///
|
||||
/// This function specializes the generic aggregate query builder for metrics
|
||||
/// by providing the metric-specific value columns and table name. It reuses
|
||||
/// the core query building logic while adding metric-specific context.
|
||||
///
|
||||
/// # Arguments
|
||||
///
|
||||
/// * `function` - The aggregation function to apply
|
||||
/// * `group_by` - The grouping specification
|
||||
/// * `from_timestamp` - The start of the time range
|
||||
/// * `to_timestamp` - The optional end of the time range
|
||||
///
|
||||
/// # Returns
|
||||
///
|
||||
/// A SQL query string optimized for metric aggregation
|
||||
pub fn build_metrics_query(
|
||||
function: AggregateFunction,
|
||||
group_by: &GroupBy,
|
||||
from_timestamp: i64,
|
||||
to_timestamp: Option<i64>,
|
||||
) -> String {
|
||||
let columns = METRIC_VALUE_COLUMNS.iter()
|
||||
.map(|&c| c.to_string())
|
||||
.collect::<Vec<_>>();
|
||||
|
||||
let column_refs: Vec<&str> = columns.iter().map(|s| s.as_str()).collect();
|
||||
build_aggregate_query(
|
||||
"metrics",
|
||||
function,
|
||||
group_by,
|
||||
&column_refs,
|
||||
Some(from_timestamp),
|
||||
to_timestamp,
|
||||
)
|
||||
}
|
||||
@@ -0,0 +1,90 @@
|
||||
pub mod aggregation;
|
||||
|
||||
use arrow_array::{ArrayRef, Float64Array, Int64Array, RecordBatch, StringArray};
|
||||
use arrow_schema::{DataType, Field, Schema};
|
||||
use std::sync::Arc;
|
||||
use tonic::Status;
|
||||
|
||||
/// A single metric record with running window calculations.
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct MetricRecord {
|
||||
/// Unique identifier for the metric
|
||||
pub metric_id: String,
|
||||
/// Unix timestamp in seconds
|
||||
pub timestamp: i64,
|
||||
/// Running sum within the window
|
||||
pub value_running_window_sum: f64,
|
||||
/// Running average within the window
|
||||
pub value_running_window_avg: f64,
|
||||
/// Running count within the window
|
||||
pub value_running_window_count: i64,
|
||||
}
|
||||
|
||||
/// Gets the schema for metric records in Arrow format.
|
||||
pub fn get_metrics_schema() -> Schema {
|
||||
Schema::new(vec![
|
||||
Field::new("metric_id", DataType::Utf8, false),
|
||||
Field::new("timestamp", DataType::Int64, false),
|
||||
Field::new("value_running_window_sum", DataType::Float64, false),
|
||||
Field::new("value_running_window_avg", DataType::Float64, false),
|
||||
Field::new("value_running_window_count", DataType::Int64, false),
|
||||
])
|
||||
}
|
||||
|
||||
/// Creates a RecordBatch from a vector of MetricRecords.
|
||||
pub fn create_record_batch(metrics: &[MetricRecord]) -> Result<RecordBatch, Status> {
|
||||
let schema = get_metrics_schema();
|
||||
|
||||
let metric_ids = StringArray::from_iter_values(metrics.iter().map(|m| m.metric_id.as_str()));
|
||||
let timestamps = Int64Array::from_iter_values(metrics.iter().map(|m| m.timestamp));
|
||||
let sums = Float64Array::from_iter_values(metrics.iter().map(|m| m.value_running_window_sum));
|
||||
let avgs = Float64Array::from_iter_values(metrics.iter().map(|m| m.value_running_window_avg));
|
||||
let counts = Int64Array::from_iter_values(metrics.iter().map(|m| m.value_running_window_count));
|
||||
|
||||
let arrays: Vec<ArrayRef> = vec![
|
||||
Arc::new(metric_ids),
|
||||
Arc::new(timestamps),
|
||||
Arc::new(sums),
|
||||
Arc::new(avgs),
|
||||
Arc::new(counts),
|
||||
];
|
||||
|
||||
RecordBatch::try_new(Arc::new(schema), arrays)
|
||||
.map_err(|e| Status::internal(format!("Failed to create record batch: {}", e)))
|
||||
}
|
||||
|
||||
/// Encodes a RecordBatch into a vector of MetricRecords.
|
||||
pub fn encode_record_batch(batch: &RecordBatch) -> Result<Vec<MetricRecord>, Status> {
|
||||
let metric_ids = batch.column_by_name("metric_id")
|
||||
.and_then(|col| col.as_any().downcast_ref::<StringArray>())
|
||||
.ok_or_else(|| Status::internal("Invalid metric_id column"))?;
|
||||
|
||||
let timestamps = batch.column_by_name("timestamp")
|
||||
.and_then(|col| col.as_any().downcast_ref::<Int64Array>())
|
||||
.ok_or_else(|| Status::internal("Invalid timestamp column"))?;
|
||||
|
||||
let sums = batch.column_by_name("value_running_window_sum")
|
||||
.and_then(|col| col.as_any().downcast_ref::<Float64Array>())
|
||||
.ok_or_else(|| Status::internal("Invalid value_running_window_sum column"))?;
|
||||
|
||||
let avgs = batch.column_by_name("value_running_window_avg")
|
||||
.and_then(|col| col.as_any().downcast_ref::<Float64Array>())
|
||||
.ok_or_else(|| Status::internal("Invalid value_running_window_avg column"))?;
|
||||
|
||||
let counts = batch.column_by_name("value_running_window_count")
|
||||
.and_then(|col| col.as_any().downcast_ref::<Int64Array>())
|
||||
.ok_or_else(|| Status::internal("Invalid value_running_window_count column"))?;
|
||||
|
||||
let mut metrics = Vec::with_capacity(batch.num_rows());
|
||||
for i in 0..batch.num_rows() {
|
||||
metrics.push(MetricRecord {
|
||||
metric_id: metric_ids.value(i).to_string(),
|
||||
timestamp: timestamps.value(i),
|
||||
value_running_window_sum: sums.value(i),
|
||||
value_running_window_avg: avgs.value(i),
|
||||
value_running_window_count: counts.value(i),
|
||||
});
|
||||
}
|
||||
|
||||
Ok(metrics)
|
||||
}
|
||||
+194
@@ -0,0 +1,194 @@
|
||||
//! Arrow Flight SQL service implementation for high-performance data transport.
|
||||
//!
|
||||
//! This module provides the core Flight SQL service implementation that enables:
|
||||
//! - High-performance data queries via Arrow Flight protocol
|
||||
//! - Support for vectorized data operations
|
||||
//! - Real-time metric aggregation queries
|
||||
//! - Time-windowed data access
|
||||
//!
|
||||
//! The service implementation is designed to work with multiple storage backends
|
||||
//! while maintaining consistent query semantics and high performance.
|
||||
|
||||
use crate::storage::StorageBackend;
|
||||
use arrow_flight::{
|
||||
flight_service_server::FlightService,
|
||||
Action, ActionType, Criteria, FlightData, FlightDescriptor, FlightInfo,
|
||||
HandshakeRequest, HandshakeResponse, PutResult, SchemaResult, Ticket,
|
||||
Empty, PollInfo,
|
||||
};
|
||||
use arrow_schema::Schema;
|
||||
use bytes::Bytes;
|
||||
use futures::Stream;
|
||||
use std::pin::Pin;
|
||||
use std::sync::Arc;
|
||||
use tonic::{Request, Response, Status, Streaming};
|
||||
use crate::storage::table_manager::AggregationView;
|
||||
use serde::Deserialize;
|
||||
use serde_json;
|
||||
|
||||
/// Command types for table and view operations
|
||||
#[derive(Debug)]
|
||||
enum TableCommand {
|
||||
CreateTable {
|
||||
name: String,
|
||||
schema: Arc<Schema>,
|
||||
},
|
||||
CreateAggregationView(AggregationView),
|
||||
DropTable(String),
|
||||
DropAggregationView(String),
|
||||
}
|
||||
|
||||
impl TableCommand {
|
||||
fn from_json(cmd: &[u8]) -> Result<Self, Status> {
|
||||
#[derive(Deserialize)]
|
||||
struct CreateTableCmd {
|
||||
name: String,
|
||||
schema_bytes: Vec<u8>,
|
||||
}
|
||||
|
||||
let value: serde_json::Value = serde_json::from_slice(cmd)
|
||||
.map_err(|e| Status::invalid_argument(format!("Invalid JSON: {}", e)))?;
|
||||
|
||||
match value.get("type").and_then(|t| t.as_str()) {
|
||||
Some("create_table") => {
|
||||
let cmd: CreateTableCmd = serde_json::from_value(value["data"].clone())
|
||||
.map_err(|e| Status::invalid_argument(format!("Invalid create table command: {}", e)))?;
|
||||
|
||||
// Convert schema bytes to Schema using Arrow IPC
|
||||
let message = arrow_ipc::root_as_message(&cmd.schema_bytes[..])
|
||||
.map_err(|e| Status::invalid_argument(format!("Invalid schema bytes: {}", e)))?;
|
||||
let schema = message.header_as_schema()
|
||||
.ok_or_else(|| Status::invalid_argument("Message is not a schema"))?;
|
||||
let schema = arrow_ipc::convert::fb_to_schema(schema);
|
||||
|
||||
Ok(TableCommand::CreateTable {
|
||||
name: cmd.name,
|
||||
schema: Arc::new(schema),
|
||||
})
|
||||
}
|
||||
Some("create_aggregation_view") => {
|
||||
let view: AggregationView = serde_json::from_value(value["data"].clone())
|
||||
.map_err(|e| Status::invalid_argument(format!("Invalid view command: {}", e)))?;
|
||||
Ok(TableCommand::CreateAggregationView(view))
|
||||
}
|
||||
Some("drop_table") => {
|
||||
let name = value["data"]["name"].as_str()
|
||||
.ok_or_else(|| Status::invalid_argument("Missing table name"))?;
|
||||
Ok(TableCommand::DropTable(name.to_string()))
|
||||
}
|
||||
Some("drop_aggregation_view") => {
|
||||
let name = value["data"]["name"].as_str()
|
||||
.ok_or_else(|| Status::invalid_argument("Missing view name"))?;
|
||||
Ok(TableCommand::DropAggregationView(name.to_string()))
|
||||
}
|
||||
_ => Err(Status::invalid_argument("Invalid command type")),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
pub struct FlightSqlService {
|
||||
backend: Box<dyn StorageBackend>,
|
||||
}
|
||||
|
||||
impl FlightSqlService {
|
||||
pub fn new(backend: Box<dyn StorageBackend>) -> Self {
|
||||
Self { backend }
|
||||
}
|
||||
}
|
||||
|
||||
#[tonic::async_trait]
|
||||
impl FlightService for FlightSqlService {
|
||||
type HandshakeStream = Pin<Box<dyn Stream<Item = Result<HandshakeResponse, Status>> + Send + 'static>>;
|
||||
type ListFlightsStream = Pin<Box<dyn Stream<Item = Result<FlightInfo, Status>> + Send + 'static>>;
|
||||
type DoGetStream = Pin<Box<dyn Stream<Item = Result<FlightData, Status>> + Send + 'static>>;
|
||||
type DoPutStream = Pin<Box<dyn Stream<Item = Result<PutResult, Status>> + Send + 'static>>;
|
||||
type DoActionStream = Pin<Box<dyn Stream<Item = Result<arrow_flight::Result, Status>> + Send + 'static>>;
|
||||
type ListActionsStream = Pin<Box<dyn Stream<Item = Result<ActionType, Status>> + Send + 'static>>;
|
||||
type DoExchangeStream = Pin<Box<dyn Stream<Item = Result<FlightData, Status>> + Send + 'static>>;
|
||||
|
||||
async fn get_schema(
|
||||
&self,
|
||||
_request: Request<FlightDescriptor>,
|
||||
) -> Result<Response<SchemaResult>, Status> {
|
||||
// Implementation here
|
||||
todo!()
|
||||
}
|
||||
|
||||
async fn do_get(
|
||||
&self,
|
||||
_request: Request<Ticket>,
|
||||
) -> Result<Response<Self::DoGetStream>, Status> {
|
||||
// Implementation here
|
||||
todo!()
|
||||
}
|
||||
|
||||
async fn handshake(
|
||||
&self,
|
||||
_request: Request<Streaming<HandshakeRequest>>,
|
||||
) -> Result<Response<Self::HandshakeStream>, Status> {
|
||||
// Implementation here
|
||||
todo!()
|
||||
}
|
||||
|
||||
async fn list_flights(
|
||||
&self,
|
||||
_request: Request<Criteria>,
|
||||
) -> Result<Response<Self::ListFlightsStream>, Status> {
|
||||
// Implementation here
|
||||
todo!()
|
||||
}
|
||||
|
||||
async fn get_flight_info(
|
||||
&self,
|
||||
_request: Request<FlightDescriptor>,
|
||||
) -> Result<Response<FlightInfo>, Status> {
|
||||
Ok(Response::new(FlightInfo {
|
||||
schema: Bytes::new(),
|
||||
flight_descriptor: None,
|
||||
endpoint: vec![],
|
||||
total_records: -1,
|
||||
total_bytes: -1,
|
||||
app_metadata: Bytes::new(),
|
||||
ordered: false,
|
||||
}))
|
||||
}
|
||||
|
||||
async fn poll_flight_info(
|
||||
&self,
|
||||
_request: Request<FlightDescriptor>,
|
||||
) -> Result<Response<PollInfo>, Status> {
|
||||
Err(Status::unimplemented("poll_flight_info not implemented"))
|
||||
}
|
||||
|
||||
async fn do_put(
|
||||
&self,
|
||||
_request: Request<Streaming<FlightData>>,
|
||||
) -> Result<Response<Self::DoPutStream>, Status> {
|
||||
// Implementation here
|
||||
todo!()
|
||||
}
|
||||
|
||||
async fn do_action(
|
||||
&self,
|
||||
_request: Request<Action>,
|
||||
) -> Result<Response<Self::DoActionStream>, Status> {
|
||||
// Implementation here
|
||||
todo!()
|
||||
}
|
||||
|
||||
async fn list_actions(
|
||||
&self,
|
||||
_request: Request<Empty>,
|
||||
) -> Result<Response<Self::ListActionsStream>, Status> {
|
||||
// Implementation here
|
||||
todo!()
|
||||
}
|
||||
|
||||
async fn do_exchange(
|
||||
&self,
|
||||
_request: Request<Streaming<FlightData>>,
|
||||
) -> Result<Response<Self::DoExchangeStream>, Status> {
|
||||
// Implementation here
|
||||
todo!()
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,870 @@
|
||||
//! ADBC (Arrow Database Connectivity) storage backend implementation.
|
||||
//!
|
||||
//! This module provides a storage backend using ADBC, enabling:
|
||||
//! - Connection to any ADBC-compliant database
|
||||
//! - High-performance data transport using Arrow's columnar format
|
||||
//! - Connection pooling and prepared statements
|
||||
//! - Support for various database systems (PostgreSQL, MySQL, etc.)
|
||||
//!
|
||||
//! # Configuration
|
||||
//!
|
||||
//! The ADBC backend can be configured using the following options:
|
||||
//!
|
||||
//! ```toml
|
||||
//! [engine]
|
||||
//! engine = "adbc"
|
||||
//! # Base connection without credentials
|
||||
//! connection = "postgresql://localhost:5432/metrics"
|
||||
//! options = {
|
||||
//! driver_path = "/usr/local/lib/libadbc_driver_postgresql.so", # Required: Path to ADBC driver
|
||||
//! pool_max = "10", # Optional: Maximum pool connections
|
||||
//! pool_min = "1", # Optional: Minimum pool connections
|
||||
//! connect_timeout = "30" # Optional: Connection timeout in seconds
|
||||
//! }
|
||||
//! ```
|
||||
//!
|
||||
//! For security, credentials should be provided via environment variables:
|
||||
//! ```bash
|
||||
//! export HYPRSTREAM_DB_USERNAME=postgres
|
||||
//! export HYPRSTREAM_DB_PASSWORD=secret
|
||||
//! ```
|
||||
//!
|
||||
//! Or via command line:
|
||||
//!
|
||||
//! ```bash
|
||||
//! hyprstream \
|
||||
//! --engine adbc \
|
||||
//! --engine-connection "postgresql://localhost:5432/metrics" \
|
||||
//! --engine-options driver_path=/usr/local/lib/libadbc_driver_postgresql.so \
|
||||
//! --engine-options pool_max=10
|
||||
//! ```
|
||||
//!
|
||||
//! The implementation is optimized for efficient data transfer and
|
||||
//! query execution using Arrow's native formats.
|
||||
|
||||
use adbc_core::{
|
||||
driver_manager::{ManagedConnection, ManagedDriver},
|
||||
options::{AdbcVersion, OptionDatabase, OptionValue},
|
||||
Connection, Database, Driver, Statement, Optionable,
|
||||
};
|
||||
use arrow_array::{
|
||||
Array, Int8Array, Int16Array, Int32Array, Int64Array,
|
||||
Float32Array, Float64Array, BooleanArray, StringArray,
|
||||
BinaryArray, TimestampNanosecondArray,
|
||||
};
|
||||
use arrow_schema::{Schema, DataType, Field};
|
||||
use async_trait::async_trait;
|
||||
use std::collections::HashMap;
|
||||
use std::sync::Arc;
|
||||
use std::sync::atomic::{AtomicU64, Ordering};
|
||||
use tokio::sync::Mutex;
|
||||
use tonic::Status;
|
||||
use crate::aggregation::{AggregateFunction, GroupBy, AggregateResult, build_aggregate_query};
|
||||
use crate::storage::table_manager::{TableManager, AggregationView};
|
||||
use crate::config::Credentials;
|
||||
use crate::metrics::MetricRecord;
|
||||
use crate::storage::StorageBackend;
|
||||
use crate::storage::cache::{CacheManager, CacheEviction};
|
||||
use arrow_array::ArrayRef;
|
||||
use arrow_array::RecordBatch;
|
||||
use crate::aggregation::TimeWindow;
|
||||
use crate::storage::BatchAggregation;
|
||||
use std::time::Duration;
|
||||
use hex;
|
||||
|
||||
pub struct AdbcBackend {
|
||||
conn: Arc<Mutex<ManagedConnection>>,
|
||||
statement_counter: AtomicU64,
|
||||
prepared_statements: Arc<Mutex<Vec<(u64, String)>>>,
|
||||
cache_manager: CacheManager,
|
||||
table_manager: TableManager,
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl CacheEviction for AdbcBackend {
|
||||
async fn execute_eviction(&self, query: &str) -> Result<(), Status> {
|
||||
let conn = self.conn.clone();
|
||||
let query = query.to_string(); // Clone for background task
|
||||
tokio::spawn(async move {
|
||||
let mut conn_guard = conn.lock().await;
|
||||
if let Err(e) = conn_guard.new_statement()
|
||||
.and_then(|mut stmt| {
|
||||
stmt.set_sql_query(&query)?;
|
||||
stmt.execute_update()
|
||||
}) {
|
||||
eprintln!("Background eviction error: {}", e);
|
||||
}
|
||||
});
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
impl AdbcBackend {
|
||||
pub fn new(driver_path: &str, connection: Option<&str>, credentials: Option<&Credentials>) -> Result<Self, Status> {
|
||||
let mut driver = ManagedDriver::load_dynamic_from_filename(
|
||||
driver_path,
|
||||
None,
|
||||
AdbcVersion::V100,
|
||||
).map_err(|e| Status::internal(format!("Failed to load ADBC driver: {}", e)))?;
|
||||
|
||||
let mut database = driver.new_database()
|
||||
.map_err(|e| Status::internal(format!("Failed to create database: {}", e)))?;
|
||||
|
||||
// Set connection string if provided
|
||||
if let Some(conn_str) = connection {
|
||||
database.set_option(OptionDatabase::Uri, OptionValue::String(conn_str.to_string()))
|
||||
.map_err(|e| Status::internal(format!("Failed to set connection string: {}", e)))?;
|
||||
}
|
||||
|
||||
// Set credentials if provided
|
||||
if let Some(creds) = credentials {
|
||||
database.set_option(OptionDatabase::Username, OptionValue::String(creds.username.clone()))
|
||||
.map_err(|e| Status::internal(format!("Failed to set username: {}", e)))?;
|
||||
|
||||
database.set_option(OptionDatabase::Password, OptionValue::String(creds.password.clone()))
|
||||
.map_err(|e| Status::internal(format!("Failed to set password: {}", e)))?;
|
||||
}
|
||||
|
||||
let connection = database.new_connection()
|
||||
.map_err(|e| Status::internal(format!("Failed to create connection: {}", e)))?;
|
||||
|
||||
Ok(Self {
|
||||
conn: Arc::new(Mutex::new(connection)),
|
||||
statement_counter: AtomicU64::new(0),
|
||||
prepared_statements: Arc::new(Mutex::new(Vec::new())),
|
||||
cache_manager: CacheManager::new(None), // Initialize without TTL
|
||||
table_manager: TableManager::new(),
|
||||
})
|
||||
}
|
||||
|
||||
async fn get_connection(&self) -> Result<tokio::sync::MutexGuard<'_, ManagedConnection>, Status> {
|
||||
Ok(self.conn.lock().await)
|
||||
}
|
||||
|
||||
async fn execute_statement(&self, conn: &mut ManagedConnection, query: &str) -> Result<(), Status> {
|
||||
let mut stmt = conn.new_statement()
|
||||
.map_err(|e| Status::internal(format!("Failed to create statement: {}", e)))?;
|
||||
|
||||
stmt.set_sql_query(query)
|
||||
.map_err(|e| Status::internal(format!("Failed to set query: {}", e)))?;
|
||||
|
||||
stmt.execute_update()
|
||||
.map_err(|e| Status::internal(format!("Failed to execute statement: {}", e)))?;
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn execute_query(&self, conn: &mut ManagedConnection, query: &str, params: Option<RecordBatch>) -> Result<Vec<MetricRecord>, Status> {
|
||||
let mut stmt = conn.new_statement()
|
||||
.map_err(|e| Status::internal(format!("Failed to create statement: {}", e)))?;
|
||||
|
||||
stmt.set_sql_query(query)
|
||||
.map_err(|e| Status::internal(format!("Failed to set query: {}", e)))?;
|
||||
|
||||
if let Some(batch) = params {
|
||||
// Create a new statement for binding parameters
|
||||
let mut bind_stmt = conn.new_statement()
|
||||
.map_err(|e| Status::internal(format!("Failed to create bind statement: {}", e)))?;
|
||||
|
||||
// Set the parameters using SQL directly
|
||||
let mut param_values = Vec::new();
|
||||
for i in 0..batch.num_rows() {
|
||||
for j in 0..batch.num_columns() {
|
||||
let col = batch.column(j);
|
||||
match col.data_type() {
|
||||
DataType::Int64 => {
|
||||
let array = col.as_any().downcast_ref::<Int64Array>().unwrap();
|
||||
param_values.push(array.value(i).to_string());
|
||||
}
|
||||
DataType::Float64 => {
|
||||
let array = col.as_any().downcast_ref::<Float64Array>().unwrap();
|
||||
param_values.push(array.value(i).to_string());
|
||||
}
|
||||
DataType::Utf8 => {
|
||||
let array = col.as_any().downcast_ref::<StringArray>().unwrap();
|
||||
param_values.push(format!("'{}'", array.value(i)));
|
||||
}
|
||||
_ => return Err(Status::internal("Unsupported parameter type")),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
let params_sql = format!("VALUES ({})", param_values.join(", "));
|
||||
bind_stmt.set_sql_query(¶ms_sql)
|
||||
.map_err(|e| Status::internal(format!("Failed to set parameters: {}", e)))?;
|
||||
|
||||
let mut bind_result = bind_stmt.execute()
|
||||
.map_err(|e| Status::internal(format!("Failed to execute parameter binding: {}", e)))?;
|
||||
|
||||
while let Some(batch_result) = bind_result.next() {
|
||||
let _ = batch_result.map_err(|e| Status::internal(format!("Failed to bind parameters: {}", e)))?;
|
||||
}
|
||||
}
|
||||
|
||||
let mut reader = stmt.execute()
|
||||
.map_err(|e| Status::internal(format!("Failed to execute query: {}", e)))?;
|
||||
|
||||
let mut metrics = Vec::new();
|
||||
while let Some(batch_result) = reader.next() {
|
||||
let batch = batch_result.map_err(|e| Status::internal(format!("Failed to get next batch: {}", e)))?;
|
||||
|
||||
let metric_ids = batch.column_by_name("metric_id")
|
||||
.and_then(|col| col.as_any().downcast_ref::<StringArray>())
|
||||
.ok_or_else(|| Status::internal("Invalid metric_id column"))?;
|
||||
|
||||
let timestamps = batch.column_by_name("timestamp")
|
||||
.and_then(|col| col.as_any().downcast_ref::<Int64Array>())
|
||||
.ok_or_else(|| Status::internal("Invalid timestamp column"))?;
|
||||
|
||||
let sums = batch.column_by_name("value_running_window_sum")
|
||||
.and_then(|col| col.as_any().downcast_ref::<Float64Array>())
|
||||
.ok_or_else(|| Status::internal("Invalid value_running_window_sum column"))?;
|
||||
|
||||
let avgs = batch.column_by_name("value_running_window_avg")
|
||||
.and_then(|col| col.as_any().downcast_ref::<Float64Array>())
|
||||
.ok_or_else(|| Status::internal("Invalid value_running_window_avg column"))?;
|
||||
|
||||
let counts = batch.column_by_name("value_running_window_count")
|
||||
.and_then(|col| col.as_any().downcast_ref::<Int64Array>())
|
||||
.ok_or_else(|| Status::internal("Invalid value_running_window_count column"))?;
|
||||
|
||||
for i in 0..batch.num_rows() {
|
||||
metrics.push(MetricRecord {
|
||||
metric_id: metric_ids.value(i).to_string(),
|
||||
timestamp: timestamps.value(i),
|
||||
value_running_window_sum: sums.value(i),
|
||||
value_running_window_avg: avgs.value(i),
|
||||
value_running_window_count: counts.value(i),
|
||||
});
|
||||
}
|
||||
}
|
||||
|
||||
Ok(metrics)
|
||||
}
|
||||
|
||||
fn prepare_timestamp_param(timestamp: i64) -> Result<RecordBatch, Status> {
|
||||
let schema = Arc::new(Schema::new(vec![
|
||||
Field::new("timestamp", DataType::Int64, false),
|
||||
]));
|
||||
|
||||
let timestamps: ArrayRef = Arc::new(Int64Array::from(vec![timestamp]));
|
||||
|
||||
RecordBatch::try_new(schema, vec![timestamps])
|
||||
.map_err(|e| Status::internal(format!("Failed to create parameter batch: {}", e)))
|
||||
}
|
||||
|
||||
fn prepare_params(metrics: &[MetricRecord]) -> Result<RecordBatch, Status> {
|
||||
let schema = Arc::new(Schema::new(vec![
|
||||
Field::new("metric_id", DataType::Utf8, false),
|
||||
Field::new("timestamp", DataType::Int64, false),
|
||||
Field::new("value_running_window_sum", DataType::Float64, false),
|
||||
Field::new("value_running_window_avg", DataType::Float64, false),
|
||||
Field::new("value_running_window_count", DataType::Int64, false),
|
||||
]));
|
||||
|
||||
let metric_ids = StringArray::from_iter_values(metrics.iter().map(|m| m.metric_id.as_str()));
|
||||
let timestamps = Int64Array::from_iter_values(metrics.iter().map(|m| m.timestamp));
|
||||
let sums = Float64Array::from_iter_values(metrics.iter().map(|m| m.value_running_window_sum));
|
||||
let avgs = Float64Array::from_iter_values(metrics.iter().map(|m| m.value_running_window_avg));
|
||||
let counts = Int64Array::from_iter_values(metrics.iter().map(|m| m.value_running_window_count));
|
||||
|
||||
let arrays: Vec<ArrayRef> = vec![
|
||||
Arc::new(metric_ids),
|
||||
Arc::new(timestamps),
|
||||
Arc::new(sums),
|
||||
Arc::new(avgs),
|
||||
Arc::new(counts),
|
||||
];
|
||||
|
||||
RecordBatch::try_new(schema, arrays)
|
||||
.map_err(|e| Status::internal(format!("Failed to create parameter batch: {}", e)))
|
||||
}
|
||||
|
||||
/// Inserts a batch of metrics with optimized aggregation updates.
|
||||
async fn insert_batch_optimized(&self, metrics: &[MetricRecord], _window: TimeWindow) -> Result<(), Status> {
|
||||
// Begin transaction
|
||||
self.begin_transaction().await?;
|
||||
let mut conn = self.conn.lock().await;
|
||||
|
||||
// Insert metrics
|
||||
let batch = Self::prepare_params(metrics)?;
|
||||
let sql = self.build_insert_sql("metrics", &batch);
|
||||
let mut stmt = conn.new_statement()
|
||||
.map_err(|e| Status::internal(format!("Failed to create statement: {}", e)))?;
|
||||
|
||||
stmt.set_sql_query(&sql)
|
||||
.map_err(|e| Status::internal(format!("Failed to set query: {}", e)))?;
|
||||
|
||||
// Bind parameters
|
||||
let mut bind_stmt = conn.new_statement()
|
||||
.map_err(|e| Status::internal(format!("Failed to create bind statement: {}", e)))?;
|
||||
|
||||
let mut param_values = Vec::new();
|
||||
for i in 0..batch.num_rows() {
|
||||
for j in 0..batch.num_columns() {
|
||||
let col = batch.column(j);
|
||||
match col.data_type() {
|
||||
DataType::Int64 => {
|
||||
let array = col.as_any().downcast_ref::<Int64Array>().unwrap();
|
||||
param_values.push(array.value(i).to_string());
|
||||
}
|
||||
DataType::Float64 => {
|
||||
let array = col.as_any().downcast_ref::<Float64Array>().unwrap();
|
||||
param_values.push(array.value(i).to_string());
|
||||
}
|
||||
DataType::Utf8 => {
|
||||
let array = col.as_any().downcast_ref::<StringArray>().unwrap();
|
||||
param_values.push(format!("'{}'", array.value(i)));
|
||||
}
|
||||
_ => return Err(Status::internal("Unsupported parameter type")),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
let params_sql = format!("VALUES ({})", param_values.join(", "));
|
||||
bind_stmt.set_sql_query(¶ms_sql)
|
||||
.map_err(|e| Status::internal(format!("Failed to set parameters: {}", e)))?;
|
||||
|
||||
let mut bind_result = bind_stmt.execute()
|
||||
.map_err(|e| Status::internal(format!("Failed to execute parameter binding: {}", e)))?;
|
||||
|
||||
while let Some(batch_result) = bind_result.next() {
|
||||
let _ = batch_result.map_err(|e| Status::internal(format!("Failed to bind parameters: {}", e)))?;
|
||||
}
|
||||
|
||||
stmt.execute_update()
|
||||
.map_err(|e| Status::internal(format!("Failed to insert metrics: {}", e)))?;
|
||||
|
||||
// Commit transaction
|
||||
self.commit_transaction().await?;
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Prepares parameters for aggregation insertion
|
||||
fn prepare_aggregation_params(agg: &BatchAggregation) -> Result<RecordBatch, Status> {
|
||||
let schema = Arc::new(Schema::new(vec![
|
||||
Field::new("metric_id", DataType::Utf8, false),
|
||||
Field::new("window_start", DataType::Int64, false),
|
||||
Field::new("window_end", DataType::Int64, false),
|
||||
Field::new("running_sum", DataType::Float64, false),
|
||||
Field::new("running_count", DataType::Int64, false),
|
||||
Field::new("min_value", DataType::Float64, false),
|
||||
Field::new("max_value", DataType::Float64, false),
|
||||
]));
|
||||
|
||||
let arrays: Vec<ArrayRef> = vec![
|
||||
Arc::new(StringArray::from(vec![agg.metric_id.as_str()])),
|
||||
Arc::new(Int64Array::from(vec![agg.window_start])),
|
||||
Arc::new(Int64Array::from(vec![agg.window_end])),
|
||||
Arc::new(Float64Array::from(vec![agg.running_sum])),
|
||||
Arc::new(Int64Array::from(vec![agg.running_count])),
|
||||
Arc::new(Float64Array::from(vec![agg.min_value])),
|
||||
Arc::new(Float64Array::from(vec![agg.max_value])),
|
||||
];
|
||||
|
||||
RecordBatch::try_new(schema, arrays)
|
||||
.map_err(|e| Status::internal(format!("Failed to create aggregation batch: {}", e)))
|
||||
}
|
||||
|
||||
async fn begin_transaction(&self) -> Result<(), Status> {
|
||||
let mut conn = self.conn.lock().await;
|
||||
let mut stmt = conn.new_statement()
|
||||
.map_err(|e| Status::internal(format!("Failed to create statement: {}", e)))?;
|
||||
|
||||
stmt.set_sql_query("BEGIN")
|
||||
.map_err(|e| Status::internal(format!("Failed to set SQL query: {}", e)))?;
|
||||
|
||||
stmt.execute_update()
|
||||
.map_err(|e| Status::internal(format!("Failed to execute query: {}", e)))?;
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn commit_transaction(&self) -> Result<(), Status> {
|
||||
let mut conn = self.conn.lock().await;
|
||||
let mut stmt = conn.new_statement()
|
||||
.map_err(|e| Status::internal(format!("Failed to create statement: {}", e)))?;
|
||||
|
||||
stmt.set_sql_query("COMMIT")
|
||||
.map_err(|e| Status::internal(format!("Failed to set SQL query: {}", e)))?;
|
||||
|
||||
stmt.execute_update()
|
||||
.map_err(|e| Status::internal(format!("Failed to execute query: {}", e)))?;
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn rollback_transaction(&self, conn: &mut ManagedConnection) -> Result<(), Status> {
|
||||
self.execute_statement(conn, "ROLLBACK").await
|
||||
}
|
||||
|
||||
async fn create_table(&self, table_name: &str, schema: &Schema) -> Result<(), Status> {
|
||||
let mut conn = self.conn.lock().await;
|
||||
let sql = self.build_create_table_sql(table_name, schema);
|
||||
self.execute_statement(&mut conn, &sql).await
|
||||
}
|
||||
|
||||
async fn create_view(&self, view: &AggregationView, sql: &str) -> Result<(), Status> {
|
||||
let mut conn = self.conn.lock().await;
|
||||
let create_view_sql = format!("CREATE VIEW {} AS {}", view.source_table, sql);
|
||||
self.execute_statement(&mut conn, &create_view_sql).await
|
||||
}
|
||||
|
||||
async fn drop_table(&self, table_name: &str) -> Result<(), Status> {
|
||||
let mut conn = self.conn.lock().await;
|
||||
let sql = format!("DROP TABLE IF EXISTS {}", table_name);
|
||||
self.execute_statement(&mut conn, &sql).await
|
||||
}
|
||||
|
||||
async fn drop_view(&self, view_name: &str) -> Result<(), Status> {
|
||||
let mut conn = self.conn.lock().await;
|
||||
let sql = format!("DROP VIEW IF EXISTS {}", view_name);
|
||||
self.execute_statement(&mut conn, &sql).await
|
||||
}
|
||||
|
||||
fn build_create_table_sql(&self, table_name: &str, schema: &Schema) -> String {
|
||||
let mut sql = format!("CREATE TABLE IF NOT EXISTS {} (", table_name);
|
||||
let mut first = true;
|
||||
|
||||
for field in schema.fields() {
|
||||
if !first {
|
||||
sql.push_str(", ");
|
||||
}
|
||||
first = false;
|
||||
|
||||
sql.push_str(&format!("{} {}", field.name(), self.arrow_type_to_sql_type(field.data_type())));
|
||||
}
|
||||
|
||||
sql.push_str(")");
|
||||
sql
|
||||
}
|
||||
|
||||
fn build_insert_sql(&self, table_name: &str, batch: &RecordBatch) -> String {
|
||||
let mut sql = format!("INSERT INTO {} (", table_name);
|
||||
let mut first = true;
|
||||
|
||||
for field in batch.schema().fields() {
|
||||
if !first {
|
||||
sql.push_str(", ");
|
||||
}
|
||||
first = false;
|
||||
sql.push_str(field.name());
|
||||
}
|
||||
|
||||
sql.push_str(") VALUES (");
|
||||
first = true;
|
||||
|
||||
for _i in 0..batch.num_columns() {
|
||||
if !first {
|
||||
sql.push_str(", ");
|
||||
}
|
||||
first = false;
|
||||
sql.push('?');
|
||||
}
|
||||
|
||||
sql.push(')');
|
||||
sql
|
||||
}
|
||||
|
||||
fn arrow_type_to_sql_type(&self, data_type: &DataType) -> &'static str {
|
||||
match data_type {
|
||||
DataType::Boolean => "BOOLEAN",
|
||||
DataType::Int8 => "TINYINT",
|
||||
DataType::Int16 => "SMALLINT",
|
||||
DataType::Int32 => "INTEGER",
|
||||
DataType::Int64 => "BIGINT",
|
||||
DataType::UInt8 => "TINYINT UNSIGNED",
|
||||
DataType::UInt16 => "SMALLINT UNSIGNED",
|
||||
DataType::UInt32 => "INTEGER UNSIGNED",
|
||||
DataType::UInt64 => "BIGINT UNSIGNED",
|
||||
DataType::Float32 => "FLOAT",
|
||||
DataType::Float64 => "DOUBLE",
|
||||
DataType::Utf8 => "VARCHAR",
|
||||
DataType::Binary => "BLOB",
|
||||
DataType::Date32 => "DATE",
|
||||
DataType::Date64 => "DATE",
|
||||
DataType::Time32(_) => "TIME",
|
||||
DataType::Time64(_) => "TIME",
|
||||
DataType::Timestamp(_, _) => "TIMESTAMP",
|
||||
_ => "VARCHAR",
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl StorageBackend for AdbcBackend {
|
||||
async fn init(&self) -> Result<(), Status> {
|
||||
let mut conn = self.conn.lock().await;
|
||||
|
||||
// Create metrics table
|
||||
let mut stmt = conn.new_statement()
|
||||
.map_err(|e| Status::internal(format!("Failed to create statement: {}", e)))?;
|
||||
|
||||
stmt.set_sql_query(r#"
|
||||
CREATE TABLE IF NOT EXISTS metrics (
|
||||
metric_id VARCHAR NOT NULL,
|
||||
timestamp BIGINT NOT NULL,
|
||||
value_running_window_sum DOUBLE PRECISION NOT NULL,
|
||||
value_running_window_avg DOUBLE PRECISION NOT NULL,
|
||||
value_running_window_count BIGINT NOT NULL,
|
||||
PRIMARY KEY (metric_id, timestamp)
|
||||
);
|
||||
|
||||
CREATE INDEX IF NOT EXISTS idx_metrics_timestamp ON metrics(timestamp);
|
||||
|
||||
CREATE TABLE IF NOT EXISTS metric_aggregations (
|
||||
metric_id VARCHAR NOT NULL,
|
||||
window_start BIGINT NOT NULL,
|
||||
window_end BIGINT NOT NULL,
|
||||
running_sum DOUBLE PRECISION NOT NULL,
|
||||
running_count BIGINT NOT NULL,
|
||||
min_value DOUBLE PRECISION NOT NULL,
|
||||
max_value DOUBLE PRECISION NOT NULL,
|
||||
PRIMARY KEY (metric_id, window_start, window_end)
|
||||
);
|
||||
|
||||
CREATE INDEX IF NOT EXISTS idx_aggregations_window
|
||||
ON metric_aggregations(window_start, window_end);
|
||||
"#).map_err(|e| Status::internal(format!("Failed to set query: {}", e)))?;
|
||||
|
||||
stmt.execute_update()
|
||||
.map_err(|e| Status::internal(format!("Failed to create tables: {}", e)))?;
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn insert_metrics(&self, metrics: Vec<MetricRecord>) -> Result<(), Status> {
|
||||
if metrics.is_empty() {
|
||||
return Ok(());
|
||||
}
|
||||
|
||||
// Check if eviction is needed
|
||||
if let Some(cutoff) = self.cache_manager.should_evict().await? {
|
||||
let query = self.cache_manager.eviction_query(cutoff);
|
||||
self.execute_eviction(&query).await?;
|
||||
}
|
||||
|
||||
// Use sliding window for batch-level aggregations
|
||||
let window = TimeWindow::Sliding {
|
||||
window: Duration::from_secs(3600), // 1 hour window
|
||||
slide: Duration::from_secs(60), // 1 minute slide
|
||||
};
|
||||
|
||||
// Use optimized batch insertion
|
||||
self.insert_batch_optimized(&metrics, window).await
|
||||
}
|
||||
|
||||
async fn query_metrics(&self, from_timestamp: i64) -> Result<Vec<MetricRecord>, Status> {
|
||||
// Check if eviction is needed
|
||||
if let Some(cutoff) = self.cache_manager.should_evict().await? {
|
||||
let query = self.cache_manager.eviction_query(cutoff);
|
||||
self.execute_eviction(&query).await?;
|
||||
}
|
||||
|
||||
let mut conn = self.conn.lock().await;
|
||||
|
||||
let query = r#"
|
||||
SELECT
|
||||
metric_id,
|
||||
timestamp,
|
||||
value_running_window_sum,
|
||||
value_running_window_avg,
|
||||
value_running_window_count
|
||||
FROM metrics
|
||||
WHERE timestamp >= ?
|
||||
ORDER BY timestamp ASC
|
||||
"#;
|
||||
|
||||
let params = Self::prepare_timestamp_param(from_timestamp)?;
|
||||
self.execute_query(&mut conn, query, Some(params)).await
|
||||
}
|
||||
|
||||
async fn prepare_sql(&self, query: &str) -> Result<Vec<u8>, Status> {
|
||||
let handle = self.statement_counter.fetch_add(1, Ordering::SeqCst);
|
||||
let mut statements = self.prepared_statements.lock().await;
|
||||
statements.push((handle, query.to_string()));
|
||||
Ok(handle.to_le_bytes().to_vec())
|
||||
}
|
||||
|
||||
async fn query_sql(&self, statement_handle: &[u8]) -> Result<Vec<MetricRecord>, Status> {
|
||||
let handle = u64::from_le_bytes(
|
||||
statement_handle.try_into()
|
||||
.map_err(|_| Status::invalid_argument("Invalid statement handle"))?
|
||||
);
|
||||
|
||||
let statements = self.prepared_statements.lock().await;
|
||||
let sql = statements
|
||||
.iter()
|
||||
.find(|(h, _)| *h == handle)
|
||||
.map(|(_, sql)| sql.as_str())
|
||||
.ok_or_else(|| Status::invalid_argument("Statement handle not found"))?;
|
||||
|
||||
let mut conn = self.conn.lock().await;
|
||||
self.execute_query(&mut conn, sql, None).await
|
||||
}
|
||||
|
||||
async fn aggregate_metrics(
|
||||
&self,
|
||||
function: AggregateFunction,
|
||||
group_by: &GroupBy,
|
||||
from_timestamp: i64,
|
||||
to_timestamp: Option<i64>,
|
||||
) -> Result<Vec<AggregateResult>, Status> {
|
||||
// Check if eviction is needed
|
||||
if let Some(cutoff) = self.cache_manager.should_evict().await? {
|
||||
let query = self.cache_manager.eviction_query(cutoff);
|
||||
self.execute_eviction(&query).await?;
|
||||
}
|
||||
|
||||
const DEFAULT_COLUMNS: [&str; 5] = [
|
||||
"metric_id",
|
||||
"timestamp",
|
||||
"value_running_window_sum",
|
||||
"value_running_window_avg",
|
||||
"value_running_window_count"
|
||||
];
|
||||
|
||||
let query = build_aggregate_query(
|
||||
"metrics",
|
||||
function,
|
||||
group_by,
|
||||
&DEFAULT_COLUMNS,
|
||||
Some(from_timestamp),
|
||||
to_timestamp,
|
||||
);
|
||||
let mut conn = self.conn.lock().await;
|
||||
let metrics = self.execute_query(&mut conn, &query, None).await?;
|
||||
|
||||
let mut results = Vec::new();
|
||||
for metric in metrics {
|
||||
let result = AggregateResult {
|
||||
value: metric.value_running_window_sum,
|
||||
timestamp: metric.timestamp,
|
||||
// Add any other fields required by AggregateResult
|
||||
};
|
||||
results.push(result);
|
||||
}
|
||||
|
||||
Ok(results)
|
||||
}
|
||||
|
||||
fn new_with_options(
|
||||
connection_string: &str,
|
||||
options: &HashMap<String, String>,
|
||||
credentials: Option<&Credentials>,
|
||||
) -> Result<Self, Status> {
|
||||
let driver_path = options.get("driver_path")
|
||||
.ok_or_else(|| Status::invalid_argument("driver_path is required"))?;
|
||||
|
||||
let mut driver = ManagedDriver::load_dynamic_from_filename(
|
||||
driver_path,
|
||||
None,
|
||||
AdbcVersion::V100,
|
||||
).map_err(|e| Status::internal(format!("Failed to load ADBC driver: {}", e)))?;
|
||||
|
||||
let mut database = driver.new_database()
|
||||
.map_err(|e| Status::internal(format!("Failed to create database: {}", e)))?;
|
||||
|
||||
// Set connection string
|
||||
database.set_option(OptionDatabase::Uri, OptionValue::String(connection_string.to_string()))
|
||||
.map_err(|e| Status::internal(format!("Failed to set connection string: {}", e)))?;
|
||||
|
||||
// Set credentials if provided
|
||||
if let Some(creds) = credentials {
|
||||
database.set_option(OptionDatabase::Username, OptionValue::String(creds.username.clone()))
|
||||
.map_err(|e| Status::internal(format!("Failed to set username: {}", e)))?;
|
||||
|
||||
database.set_option(OptionDatabase::Password, OptionValue::String(creds.password.clone()))
|
||||
.map_err(|e| Status::internal(format!("Failed to set password: {}", e)))?;
|
||||
}
|
||||
|
||||
let connection = database.new_connection()
|
||||
.map_err(|e| Status::internal(format!("Failed to create connection: {}", e)))?;
|
||||
|
||||
Ok(Self {
|
||||
conn: Arc::new(Mutex::new(connection)),
|
||||
statement_counter: AtomicU64::new(0),
|
||||
prepared_statements: Arc::new(Mutex::new(Vec::new())),
|
||||
cache_manager: CacheManager::new(None), // Initialize without TTL
|
||||
table_manager: TableManager::new(),
|
||||
})
|
||||
}
|
||||
|
||||
async fn create_table(&self, table_name: &str, schema: &Schema) -> Result<(), Status> {
|
||||
let mut conn = self.conn.lock().await;
|
||||
let sql = self.build_create_table_sql(table_name, schema);
|
||||
self.execute_statement(&mut conn, &sql).await
|
||||
}
|
||||
|
||||
async fn insert_into_table(&self, table_name: &str, batch: RecordBatch) -> Result<(), Status> {
|
||||
let mut conn = self.conn.lock().await;
|
||||
let sql = self.build_insert_sql(table_name, &batch);
|
||||
self.execute_statement(&mut conn, &sql).await
|
||||
}
|
||||
|
||||
async fn query_table(&self, table_name: &str, projection: Option<Vec<String>>) -> Result<RecordBatch, Status> {
|
||||
let mut conn = self.conn.lock().await;
|
||||
let mut stmt = conn.new_statement()
|
||||
.map_err(|e| Status::internal(format!("Failed to create statement: {}", e)))?;
|
||||
|
||||
let columns = projection.map(|cols| cols.join(", ")).unwrap_or_else(|| "*".to_string());
|
||||
let sql = format!("SELECT {} FROM {}", columns, table_name);
|
||||
|
||||
stmt.set_sql_query(&sql)
|
||||
.map_err(|e| Status::internal(format!("Failed to set query: {}", e)))?;
|
||||
|
||||
let mut reader = stmt.execute()
|
||||
.map_err(|e| Status::internal(format!("Failed to execute query: {}", e)))?;
|
||||
|
||||
let batch = reader.next()
|
||||
.ok_or_else(|| Status::internal("No data returned"))?
|
||||
.map_err(|e| Status::internal(format!("Failed to read record batch: {}", e)))?;
|
||||
|
||||
// Convert ADBC RecordBatch to Arrow RecordBatch
|
||||
let schema = batch.schema();
|
||||
let mut arrays = Vec::with_capacity(batch.num_columns());
|
||||
|
||||
for i in 0..batch.num_columns() {
|
||||
let col = batch.column(i);
|
||||
let array: ArrayRef = match col.data_type() {
|
||||
&duckdb::arrow::datatypes::DataType::Int64 => {
|
||||
Arc::new(col.as_any().downcast_ref::<Int64Array>().unwrap().clone())
|
||||
},
|
||||
&duckdb::arrow::datatypes::DataType::Float64 => {
|
||||
Arc::new(col.as_any().downcast_ref::<Float64Array>().unwrap().clone())
|
||||
},
|
||||
&duckdb::arrow::datatypes::DataType::Utf8 => {
|
||||
Arc::new(col.as_any().downcast_ref::<StringArray>().unwrap().clone())
|
||||
},
|
||||
_ => return Err(Status::internal("Unsupported column type")),
|
||||
};
|
||||
arrays.push(array);
|
||||
}
|
||||
|
||||
// Convert DuckDB schema to Arrow schema
|
||||
let fields: Vec<Field> = schema.fields().iter().map(|f| {
|
||||
Field::new(
|
||||
f.name(),
|
||||
match f.data_type() {
|
||||
&duckdb::arrow::datatypes::DataType::Int64 => DataType::Int64,
|
||||
&duckdb::arrow::datatypes::DataType::Float64 => DataType::Float64,
|
||||
&duckdb::arrow::datatypes::DataType::Utf8 => DataType::Utf8,
|
||||
_ => DataType::Utf8, // Default to string for unsupported types
|
||||
},
|
||||
f.is_nullable()
|
||||
)
|
||||
}).collect();
|
||||
|
||||
let arrow_schema = Schema::new(fields);
|
||||
RecordBatch::try_new(Arc::new(arrow_schema), arrays)
|
||||
.map_err(|e| Status::internal(format!("Failed to create record batch: {}", e)))
|
||||
}
|
||||
|
||||
async fn create_aggregation_view(&self, view: &AggregationView) -> Result<(), Status> {
|
||||
let columns: Vec<&str> = view.aggregate_columns.iter()
|
||||
.map(|s| s.as_str())
|
||||
.collect();
|
||||
|
||||
let sql = build_aggregate_query(
|
||||
&view.source_table,
|
||||
view.function,
|
||||
&view.group_by,
|
||||
&columns,
|
||||
None,
|
||||
None
|
||||
);
|
||||
|
||||
let mut conn = self.conn.lock().await;
|
||||
let mut stmt = conn.new_statement()
|
||||
.map_err(|e| Status::internal(format!("Failed to create statement: {}", e)))?;
|
||||
|
||||
stmt.set_sql_query(&format!("CREATE VIEW {} AS {}", view.source_table, sql))
|
||||
.map_err(|e| Status::internal(format!("Failed to set SQL query: {}", e)))?;
|
||||
|
||||
stmt.execute_update()
|
||||
.map_err(|e| Status::internal(format!("Failed to execute query: {}", e)))?;
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn query_aggregation_view(&self, view_name: &str) -> Result<RecordBatch, Status> {
|
||||
let sql = format!("SELECT * FROM {}", view_name);
|
||||
|
||||
let mut conn = self.conn.lock().await;
|
||||
let mut stmt = conn.new_statement()
|
||||
.map_err(|e| Status::internal(format!("Failed to create statement: {}", e)))?;
|
||||
|
||||
stmt.set_sql_query(&sql)
|
||||
.map_err(|e| Status::internal(format!("Failed to set SQL query: {}", e)))?;
|
||||
|
||||
let mut reader = stmt.execute()
|
||||
.map_err(|e| Status::internal(format!("Failed to execute query: {}", e)))?;
|
||||
|
||||
let batch = reader.next()
|
||||
.ok_or_else(|| Status::internal("No data returned"))?
|
||||
.map_err(|e| Status::internal(format!("Failed to read record batch: {}", e)))?;
|
||||
|
||||
// Convert DuckDB RecordBatch to Arrow RecordBatch
|
||||
let schema = batch.schema();
|
||||
let mut arrays = Vec::with_capacity(batch.num_columns());
|
||||
|
||||
for i in 0..batch.num_columns() {
|
||||
let col = batch.column(i);
|
||||
let array: ArrayRef = match col.data_type() {
|
||||
&duckdb::arrow::datatypes::DataType::Int64 => {
|
||||
Arc::new(col.as_any().downcast_ref::<Int64Array>().unwrap().clone())
|
||||
},
|
||||
&duckdb::arrow::datatypes::DataType::Float64 => {
|
||||
Arc::new(col.as_any().downcast_ref::<Float64Array>().unwrap().clone())
|
||||
},
|
||||
&duckdb::arrow::datatypes::DataType::Utf8 => {
|
||||
Arc::new(col.as_any().downcast_ref::<StringArray>().unwrap().clone())
|
||||
},
|
||||
_ => return Err(Status::internal("Unsupported column type")),
|
||||
};
|
||||
arrays.push(array);
|
||||
}
|
||||
|
||||
// Convert DuckDB schema to Arrow schema
|
||||
let fields: Vec<Field> = schema.fields().iter().map(|f| {
|
||||
Field::new(
|
||||
f.name(),
|
||||
match f.data_type() {
|
||||
&duckdb::arrow::datatypes::DataType::Int64 => DataType::Int64,
|
||||
&duckdb::arrow::datatypes::DataType::Float64 => DataType::Float64,
|
||||
&duckdb::arrow::datatypes::DataType::Utf8 => DataType::Utf8,
|
||||
_ => DataType::Utf8, // Default to string for unsupported types
|
||||
},
|
||||
f.is_nullable()
|
||||
)
|
||||
}).collect();
|
||||
|
||||
let arrow_schema = Schema::new(fields);
|
||||
RecordBatch::try_new(Arc::new(arrow_schema), arrays)
|
||||
.map_err(|e| Status::internal(format!("Failed to create record batch: {}", e)))
|
||||
}
|
||||
|
||||
async fn drop_table(&self, table_name: &str) -> Result<(), Status> {
|
||||
let mut conn = self.conn.lock().await;
|
||||
let mut stmt = conn.new_statement().map_err(|e| Status::internal(format!("Failed to create statement: {}", e)))?;
|
||||
stmt.set_sql_query(&format!("DROP TABLE IF EXISTS {}", table_name))
|
||||
.map_err(|e| Status::internal(format!("Failed to set SQL query: {}", e)))?;
|
||||
stmt.execute_update().map_err(|e| Status::internal(format!("Failed to drop table: {}", e)))?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn drop_aggregation_view(&self, view_name: &str) -> Result<(), Status> {
|
||||
let mut conn = self.conn.lock().await;
|
||||
let mut stmt = conn.new_statement().map_err(|e| Status::internal(format!("Failed to create statement: {}", e)))?;
|
||||
stmt.set_sql_query(&format!("DROP VIEW IF EXISTS {}", view_name))
|
||||
.map_err(|e| Status::internal(format!("Failed to set SQL query: {}", e)))?;
|
||||
stmt.execute_update().map_err(|e| Status::internal(format!("Failed to drop view: {}", e)))?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn table_manager(&self) -> &TableManager {
|
||||
&self.table_manager
|
||||
}
|
||||
}
|
||||
|
||||
@@ -0,0 +1,78 @@
|
||||
use std::sync::Arc;
|
||||
use std::time::{SystemTime, Duration, UNIX_EPOCH};
|
||||
use tokio::sync::RwLock;
|
||||
use tonic::Status;
|
||||
|
||||
/// Shared cache eviction manager that can be used across different storage backends.
|
||||
#[derive(Clone)]
|
||||
pub struct CacheManager {
|
||||
ttl: Option<u64>,
|
||||
last_eviction: Arc<RwLock<SystemTime>>,
|
||||
min_eviction_interval: Duration,
|
||||
}
|
||||
|
||||
impl CacheManager {
|
||||
/// Creates a new cache manager with the specified TTL.
|
||||
pub fn new(ttl: Option<u64>) -> Self {
|
||||
Self {
|
||||
ttl,
|
||||
last_eviction: Arc::new(RwLock::new(SystemTime::now())),
|
||||
min_eviction_interval: Duration::from_secs(60), // Default 60s between evictions
|
||||
}
|
||||
}
|
||||
|
||||
/// Sets a custom minimum interval between evictions.
|
||||
pub fn set_min_eviction_interval(&mut self, interval: Duration) {
|
||||
self.min_eviction_interval = interval;
|
||||
}
|
||||
|
||||
/// Checks if eviction should be performed based on TTL and rate limiting.
|
||||
/// Returns the cutoff timestamp if eviction should proceed, None otherwise.
|
||||
pub async fn should_evict(&self) -> Result<Option<i64>, Status> {
|
||||
match self.ttl {
|
||||
None | Some(0) => Ok(None), // No TTL or TTL=0 means no eviction
|
||||
Some(ttl) => {
|
||||
let now = SystemTime::now();
|
||||
let last = *self.last_eviction.read().await;
|
||||
|
||||
// Check if enough time has passed since last eviction
|
||||
if now.duration_since(last).unwrap_or(Duration::from_secs(0)) < self.min_eviction_interval {
|
||||
return Ok(None);
|
||||
}
|
||||
|
||||
// Calculate cutoff timestamp
|
||||
let cutoff = now
|
||||
.duration_since(UNIX_EPOCH)
|
||||
.map_err(|e| Status::internal(e.to_string()))?
|
||||
.as_secs() as i64
|
||||
- ttl as i64;
|
||||
|
||||
// Update last eviction time
|
||||
*self.last_eviction.write().await = now;
|
||||
|
||||
Ok(Some(cutoff))
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Generates an optimized SQL query for evicting expired entries.
|
||||
pub fn eviction_query(&self, cutoff: i64) -> String {
|
||||
format!(
|
||||
"DELETE FROM metrics USING (
|
||||
SELECT timestamp
|
||||
FROM metrics
|
||||
WHERE timestamp < {}
|
||||
LIMIT 10000
|
||||
) as expired
|
||||
WHERE metrics.timestamp = expired.timestamp",
|
||||
cutoff
|
||||
)
|
||||
}
|
||||
}
|
||||
|
||||
/// Trait for storage backends that support cache eviction.
|
||||
#[async_trait::async_trait]
|
||||
pub trait CacheEviction {
|
||||
/// Executes the eviction query in the background.
|
||||
async fn execute_eviction(&self, query: &str) -> Result<(), Status>;
|
||||
}
|
||||
@@ -0,0 +1,245 @@
|
||||
//! Two-tier cached storage backend implementation.
|
||||
//!
|
||||
//! This module provides a caching layer on top of any storage backend:
|
||||
//! - Fast access to frequently used data
|
||||
//! - Write-through caching for consistency
|
||||
//! - Configurable cache duration
|
||||
//! - Support for any StorageBackend as cache or store
|
||||
//!
|
||||
//! # Configuration
|
||||
//!
|
||||
//! The cached backend can be configured using the following options:
|
||||
//!
|
||||
//! ```toml
|
||||
//! # Primary storage configuration
|
||||
//! [engine]
|
||||
//! engine = "adbc"
|
||||
//! connection = "postgresql://localhost:5432"
|
||||
//! options = {
|
||||
//! driver_path = "/usr/local/lib/libadbc_driver_postgresql.so",
|
||||
//! username = "postgres",
|
||||
//! database = "metrics"
|
||||
//! }
|
||||
//!
|
||||
//! # Cache configuration
|
||||
//! [cache]
|
||||
//! enabled = true
|
||||
//! engine = "duckdb"
|
||||
//! connection = ":memory:"
|
||||
//! max_duration_secs = 3600
|
||||
//! options = {
|
||||
//! threads = "2"
|
||||
//! }
|
||||
//! ```
|
||||
//!
|
||||
//! Or via command line:
|
||||
//!
|
||||
//! ```bash
|
||||
//! hyprstream \
|
||||
//! --engine adbc \
|
||||
//! --engine-connection "postgresql://localhost:5432" \
|
||||
//! --engine-options driver_path=/usr/local/lib/libadbc_driver_postgresql.so \
|
||||
//! --engine-options username=postgres \
|
||||
//! --enable-cache \
|
||||
//! --cache-engine duckdb \
|
||||
//! --cache-connection ":memory:" \
|
||||
//! --cache-options threads=2 \
|
||||
//! --cache-max-duration 3600
|
||||
//! ```
|
||||
//!
|
||||
//! The implementation follows standard caching patterns while ensuring
|
||||
//! data consistency between cache and backing store.
|
||||
|
||||
use crate::config::Credentials;
|
||||
use crate::metrics::MetricRecord;
|
||||
use crate::storage::{StorageBackend, adbc::AdbcBackend, duckdb::DuckDbBackend};
|
||||
use std::sync::Arc;
|
||||
use std::collections::HashMap;
|
||||
use tonic::Status;
|
||||
|
||||
/// Two-tier storage backend with caching support.
|
||||
///
|
||||
/// This backend provides:
|
||||
/// - Fast access to recent data through caching
|
||||
/// - Write-through caching for data consistency
|
||||
/// - Configurable cache duration
|
||||
/// - Support for any StorageBackend implementation
|
||||
///
|
||||
/// The implementation uses two storage backends:
|
||||
/// 1. A fast cache (e.g., in-memory DuckDB)
|
||||
/// 2. A persistent store (e.g., PostgreSQL via ADBC)
|
||||
pub struct CachedStorageBackend {
|
||||
/// Fast storage backend for caching
|
||||
cache: Arc<dyn StorageBackend>,
|
||||
/// Persistent storage backend for data
|
||||
store: Arc<dyn StorageBackend>,
|
||||
/// Maximum cache entry lifetime in seconds
|
||||
max_duration_secs: u64,
|
||||
}
|
||||
|
||||
impl CachedStorageBackend {
|
||||
/// Creates a new cached storage backend.
|
||||
///
|
||||
/// This method sets up a two-tier storage system with:
|
||||
/// - A fast cache layer for frequent access
|
||||
/// - A persistent backing store
|
||||
/// - Configurable cache duration
|
||||
///
|
||||
/// # Arguments
|
||||
///
|
||||
/// * `cache` - Fast storage backend for caching
|
||||
/// * `store` - Persistent storage backend
|
||||
/// * `max_duration_secs` - Maximum cache entry lifetime in seconds
|
||||
pub fn new(
|
||||
cache: Arc<dyn StorageBackend>,
|
||||
store: Arc<dyn StorageBackend>,
|
||||
max_duration_secs: u64,
|
||||
) -> Self {
|
||||
Self {
|
||||
cache,
|
||||
store,
|
||||
max_duration_secs,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait::async_trait]
|
||||
impl StorageBackend for CachedStorageBackend {
|
||||
/// Initializes both cache and backing store.
|
||||
///
|
||||
/// This method ensures both storage layers are properly
|
||||
/// initialized and ready for use.
|
||||
async fn init(&self) -> Result<(), Status> {
|
||||
// Initialize both cache and backing store
|
||||
self.cache.init().await?;
|
||||
self.store.init().await?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Inserts metrics into both cache and backing store.
|
||||
///
|
||||
/// This method implements write-through caching:
|
||||
/// 1. Writes to cache for fast access
|
||||
/// 2. Writes to backing store for persistence
|
||||
///
|
||||
/// # Arguments
|
||||
///
|
||||
/// * `metrics` - Vector of MetricRecord instances to insert
|
||||
async fn insert_metrics(&self, metrics: Vec<MetricRecord>) -> Result<(), Status> {
|
||||
// Insert into both cache and backing store
|
||||
self.cache.insert_metrics(metrics.clone()).await?;
|
||||
self.store.insert_metrics(metrics).await?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Queries metrics with caching support.
|
||||
///
|
||||
/// This method implements a cache-first query strategy:
|
||||
/// 1. Attempts to read from cache
|
||||
/// 2. On cache miss, reads from backing store
|
||||
/// 3. Updates cache with results from backing store
|
||||
///
|
||||
/// # Arguments
|
||||
///
|
||||
/// * `from_timestamp` - Unix timestamp to query from
|
||||
async fn query_metrics(&self, from_timestamp: i64) -> Result<Vec<MetricRecord>, Status> {
|
||||
// Calculate cache cutoff time
|
||||
let now = std::time::SystemTime::now()
|
||||
.duration_since(std::time::UNIX_EPOCH)
|
||||
.unwrap()
|
||||
.as_secs() as i64;
|
||||
let cache_cutoff = now - self.max_duration_secs as i64;
|
||||
|
||||
// Only use cache for data within cache window
|
||||
if from_timestamp >= cache_cutoff {
|
||||
match self.cache.query_metrics(from_timestamp).await {
|
||||
Ok(metrics) if !metrics.is_empty() => return Ok(metrics),
|
||||
_ => {}
|
||||
}
|
||||
}
|
||||
|
||||
// Cache miss or data too old, query backing store
|
||||
let metrics = self.store.query_metrics(from_timestamp).await?;
|
||||
|
||||
// Update cache with results if within cache window
|
||||
if !metrics.is_empty() && from_timestamp >= cache_cutoff {
|
||||
self.cache.insert_metrics(metrics.clone()).await?;
|
||||
}
|
||||
|
||||
Ok(metrics)
|
||||
}
|
||||
|
||||
/// Prepares a SQL statement on the backing store.
|
||||
///
|
||||
/// This method bypasses the cache and prepares statements
|
||||
/// directly on the backing store, as prepared statements
|
||||
/// are typically used for complex queries.
|
||||
///
|
||||
/// # Arguments
|
||||
///
|
||||
/// * `query` - SQL query to prepare
|
||||
async fn prepare_sql(&self, query: &str) -> Result<Vec<u8>, Status> {
|
||||
// Prepare on backing store only
|
||||
self.store.prepare_sql(query).await
|
||||
}
|
||||
|
||||
/// Executes a prepared SQL statement on the backing store.
|
||||
///
|
||||
/// This method bypasses the cache and executes statements
|
||||
/// directly on the backing store, ensuring consistent results
|
||||
/// for complex queries.
|
||||
///
|
||||
/// # Arguments
|
||||
///
|
||||
/// * `statement_handle` - Handle of the prepared statement
|
||||
async fn query_sql(&self, statement_handle: &[u8]) -> Result<Vec<MetricRecord>, Status> {
|
||||
// Execute on backing store only
|
||||
self.store.query_sql(statement_handle).await
|
||||
}
|
||||
|
||||
fn new_with_options(
|
||||
connection_string: &str,
|
||||
options: &HashMap<String, String>,
|
||||
credentials: Option<&Credentials>,
|
||||
) -> Result<Self, Status> {
|
||||
// Parse cache duration from options
|
||||
let max_duration_secs = options
|
||||
.get("max_duration_secs")
|
||||
.and_then(|s| s.parse().ok())
|
||||
.unwrap_or(3600);
|
||||
|
||||
// Create cache backend
|
||||
let default_engine = "duckdb".to_string();
|
||||
let default_connection = ":memory:".to_string();
|
||||
let cache_engine = options.get("cache_engine").unwrap_or(&default_engine);
|
||||
let cache_connection = options.get("cache_connection").unwrap_or(&default_connection);
|
||||
let cache_options: HashMap<String, String> = options
|
||||
.iter()
|
||||
.filter(|(k, _)| k.starts_with("cache_"))
|
||||
.map(|(k, v)| (k[6..].to_string(), v.clone()))
|
||||
.collect();
|
||||
|
||||
let cache: Arc<dyn StorageBackend> = match cache_engine.as_str() {
|
||||
"duckdb" => Arc::new(DuckDbBackend::new_with_options(
|
||||
cache_connection,
|
||||
&cache_options,
|
||||
None,
|
||||
)?),
|
||||
"adbc" => Arc::new(AdbcBackend::new_with_options(
|
||||
cache_connection,
|
||||
&cache_options,
|
||||
None,
|
||||
)?),
|
||||
_ => return Err(Status::invalid_argument("Invalid cache engine type")),
|
||||
};
|
||||
|
||||
// Create store backend
|
||||
let store = Arc::new(AdbcBackend::new_with_options(
|
||||
connection_string,
|
||||
options,
|
||||
credentials,
|
||||
)?);
|
||||
|
||||
Ok(Self::new(cache, store, max_duration_secs))
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,637 @@
|
||||
//! DuckDB storage backend implementation.
|
||||
//!
|
||||
//! This module provides a high-performance storage backend using DuckDB,
|
||||
//! an embedded analytical database. The implementation supports:
|
||||
//! - In-memory and persistent storage options
|
||||
//! - Efficient batch operations
|
||||
//! - SQL query capabilities
|
||||
//! - Time-based filtering
|
||||
//!
|
||||
//! # Configuration
|
||||
//!
|
||||
//! The DuckDB backend can be configured using the following options:
|
||||
//!
|
||||
//! ```toml
|
||||
//! [engine]
|
||||
//! engine = "duckdb"
|
||||
//! connection = ":memory:" # Use ":memory:" for in-memory or file path
|
||||
//! options = {
|
||||
//! threads = "4", # Optional: Number of threads (default: 4)
|
||||
//! read_only = "false" # Optional: Read-only mode (default: false)
|
||||
//! }
|
||||
//! ```
|
||||
//!
|
||||
//! Or via command line:
|
||||
//!
|
||||
//! ```bash
|
||||
//! hyprstream \
|
||||
//! --engine duckdb \
|
||||
//! --engine-connection ":memory:" \
|
||||
//! --engine-options threads=4 \
|
||||
//! --engine-options read_only=false
|
||||
//! ```
|
||||
//!
|
||||
//! DuckDB is particularly well-suited for analytics workloads and
|
||||
//! provides excellent performance for both caching and primary storage.
|
||||
|
||||
use std::collections::HashMap;
|
||||
use std::sync::Arc;
|
||||
use duckdb::{Connection, Config, params, ToSql};
|
||||
use tokio::sync::Mutex;
|
||||
use tonic::Status;
|
||||
use crate::metrics::MetricRecord;
|
||||
use crate::config::Credentials;
|
||||
use crate::storage::{StorageBackend, BatchAggregation};
|
||||
use crate::storage::cache::{CacheManager, CacheEviction};
|
||||
use crate::storage::table_manager::{TableManager, AggregationView};
|
||||
use crate::aggregation::{TimeWindow, AggregateFunction, GroupBy, AggregateResult, build_aggregate_query};
|
||||
use async_trait::async_trait;
|
||||
use arrow::datatypes::{DataType, Field, Schema};
|
||||
use arrow::array::{
|
||||
Array, ArrayRef, RecordBatch, Int64Array, Float64Array, StringArray,
|
||||
};
|
||||
use arrow::array::builder::{
|
||||
ArrayBuilder, Int64Builder, Float64Builder, StringBuilder,
|
||||
};
|
||||
use std::time::Duration;
|
||||
|
||||
/// DuckDB-based storage backend for metrics.
|
||||
#[derive(Clone)]
|
||||
pub struct DuckDbBackend {
|
||||
conn: Arc<Mutex<Connection>>,
|
||||
cache_manager: CacheManager,
|
||||
table_manager: TableManager,
|
||||
}
|
||||
|
||||
impl DuckDbBackend {
|
||||
/// Creates a new DuckDB backend instance.
|
||||
pub fn new(connection_string: String, _options: HashMap<String, String>, ttl: Option<u64>) -> Result<Self, Status> {
|
||||
let config = Config::default();
|
||||
let conn = Connection::open_with_flags(&connection_string, config)
|
||||
.map_err(|e| Status::internal(e.to_string()))?;
|
||||
|
||||
let backend = Self {
|
||||
conn: Arc::new(Mutex::new(conn)),
|
||||
cache_manager: CacheManager::new(ttl),
|
||||
table_manager: TableManager::new(),
|
||||
};
|
||||
|
||||
// Initialize tables
|
||||
let backend_clone = backend.clone();
|
||||
tokio::spawn(async move {
|
||||
if let Err(e) = backend_clone.init().await {
|
||||
eprintln!("Failed to initialize tables: {}", e);
|
||||
}
|
||||
});
|
||||
|
||||
Ok(backend)
|
||||
}
|
||||
|
||||
/// Creates a new DuckDB backend with an in-memory database.
|
||||
pub fn new_in_memory() -> Result<Self, Status> {
|
||||
Self::new(":memory:".to_string(), HashMap::new(), Some(0))
|
||||
}
|
||||
|
||||
/// Inserts a batch of metrics with optimized aggregation updates.
|
||||
async fn insert_batch_optimized(&self, metrics: &[MetricRecord], window: TimeWindow) -> Result<(), Status> {
|
||||
let conn = self.conn.lock().await;
|
||||
|
||||
// Begin transaction
|
||||
conn.execute("BEGIN TRANSACTION", params![])
|
||||
.map_err(|e| Status::internal(format!("Failed to begin transaction: {}", e)))?;
|
||||
|
||||
// Convert metrics to RecordBatch for efficient insertion
|
||||
let batch = Self::prepare_params(metrics)?;
|
||||
|
||||
// Insert metrics using prepared statement
|
||||
let mut stmt = conn.prepare(r#"
|
||||
INSERT INTO metrics (
|
||||
metric_id,
|
||||
timestamp,
|
||||
value_running_window_sum,
|
||||
value_running_window_avg,
|
||||
value_running_window_count
|
||||
) VALUES (?, ?, ?, ?, ?)
|
||||
"#).map_err(|e| Status::internal(format!("Failed to prepare statement: {}", e)))?;
|
||||
|
||||
// Bind and execute in batches
|
||||
for i in 0..batch.num_rows() {
|
||||
let metric_id = batch.column(0).as_any().downcast_ref::<StringArray>().unwrap().value(i);
|
||||
let timestamp = batch.column(1).as_any().downcast_ref::<Int64Array>().unwrap().value(i);
|
||||
let sum = batch.column(2).as_any().downcast_ref::<Float64Array>().unwrap().value(i);
|
||||
let avg = batch.column(3).as_any().downcast_ref::<Float64Array>().unwrap().value(i);
|
||||
let count = batch.column(4).as_any().downcast_ref::<Int64Array>().unwrap().value(i);
|
||||
|
||||
stmt.execute(params![
|
||||
metric_id,
|
||||
timestamp,
|
||||
sum,
|
||||
avg,
|
||||
count,
|
||||
]).map_err(|e| Status::internal(format!("Failed to insert metrics: {}", e)))?;
|
||||
}
|
||||
|
||||
// Update aggregations based on window
|
||||
let window_start = match window {
|
||||
TimeWindow::Sliding { window, slide: _ } => {
|
||||
let now = metrics.iter().map(|m| m.timestamp).max().unwrap_or(0);
|
||||
now - window.as_nanos() as i64
|
||||
}
|
||||
TimeWindow::Fixed(start) => start.as_nanos() as i64,
|
||||
TimeWindow::None => metrics.iter().map(|m| m.timestamp).min().unwrap_or(0),
|
||||
};
|
||||
|
||||
let window_end = match window {
|
||||
TimeWindow::Sliding { window: _, slide: _ } => {
|
||||
metrics.iter().map(|m| m.timestamp).max().unwrap_or(0)
|
||||
}
|
||||
TimeWindow::Fixed(end) => end.as_nanos() as i64,
|
||||
TimeWindow::None => metrics.iter().map(|m| m.timestamp).max().unwrap_or(0),
|
||||
};
|
||||
|
||||
// Group metrics by ID and calculate aggregations
|
||||
let mut aggregations = HashMap::new();
|
||||
for metric in metrics {
|
||||
let entry = aggregations.entry(metric.metric_id.clone()).or_insert_with(|| BatchAggregation {
|
||||
metric_id: metric.metric_id.clone(),
|
||||
window_start,
|
||||
window_end,
|
||||
running_sum: 0.0,
|
||||
running_count: 0,
|
||||
min_value: f64::INFINITY,
|
||||
max_value: f64::NEG_INFINITY,
|
||||
});
|
||||
|
||||
entry.running_sum += metric.value_running_window_sum;
|
||||
entry.running_count += metric.value_running_window_count as i64;
|
||||
entry.min_value = entry.min_value.min(metric.value_running_window_sum);
|
||||
entry.max_value = entry.max_value.max(metric.value_running_window_sum);
|
||||
}
|
||||
|
||||
// Update aggregations table using prepared statement with proper type handling
|
||||
let mut agg_stmt = conn.prepare(r#"
|
||||
INSERT INTO metric_aggregations (
|
||||
metric_id, window_start, window_end,
|
||||
running_sum, running_count, min_value, max_value
|
||||
) VALUES (?, ?, ?, ?, ?, ?, ?)
|
||||
ON CONFLICT (metric_id, window_start, window_end) DO UPDATE
|
||||
SET running_sum = metric_aggregations.running_sum + EXCLUDED.running_sum,
|
||||
running_count = metric_aggregations.running_count + EXCLUDED.running_count,
|
||||
min_value = LEAST(metric_aggregations.min_value, EXCLUDED.min_value),
|
||||
max_value = GREATEST(metric_aggregations.max_value, EXCLUDED.max_value)
|
||||
"#).map_err(|e| Status::internal(format!("Failed to prepare aggregation statement: {}", e)))?;
|
||||
|
||||
for agg in aggregations.values() {
|
||||
agg_stmt.execute(params![
|
||||
agg.metric_id,
|
||||
agg.window_start,
|
||||
agg.window_end,
|
||||
agg.running_sum,
|
||||
agg.running_count,
|
||||
agg.min_value,
|
||||
agg.max_value,
|
||||
]).map_err(|e| Status::internal(format!("Failed to update aggregations: {}", e)))?;
|
||||
}
|
||||
|
||||
// Commit transaction
|
||||
conn.execute("COMMIT", params![])
|
||||
.map_err(|e| Status::internal(format!("Failed to commit transaction: {}", e)))?;
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Prepares parameters for batch insertion
|
||||
fn prepare_params(metrics: &[MetricRecord]) -> Result<RecordBatch, Status> {
|
||||
let schema = Arc::new(Schema::new(vec![
|
||||
Field::new("metric_id", DataType::Utf8, false),
|
||||
Field::new("timestamp", DataType::Int64, false),
|
||||
Field::new("value_running_window_sum", DataType::Float64, false),
|
||||
Field::new("value_running_window_avg", DataType::Float64, false),
|
||||
Field::new("value_running_window_count", DataType::Int64, false),
|
||||
]));
|
||||
|
||||
let metric_ids = StringArray::from_iter_values(metrics.iter().map(|m| m.metric_id.as_str()));
|
||||
let timestamps = Int64Array::from_iter_values(metrics.iter().map(|m| m.timestamp));
|
||||
let sums = Float64Array::from_iter_values(metrics.iter().map(|m| m.value_running_window_sum));
|
||||
let avgs = Float64Array::from_iter_values(metrics.iter().map(|m| m.value_running_window_avg));
|
||||
let counts = Int64Array::from_iter_values(metrics.iter().map(|m| m.value_running_window_count));
|
||||
|
||||
let arrays: Vec<ArrayRef> = vec![
|
||||
Arc::new(metric_ids),
|
||||
Arc::new(timestamps),
|
||||
Arc::new(sums),
|
||||
Arc::new(avgs),
|
||||
Arc::new(counts),
|
||||
];
|
||||
|
||||
RecordBatch::try_new(schema, arrays)
|
||||
.map_err(|e| Status::internal(format!("Failed to create parameter batch: {}", e)))
|
||||
}
|
||||
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl CacheEviction for DuckDbBackend {
|
||||
async fn execute_eviction(&self, query: &str) -> Result<(), Status> {
|
||||
let conn = self.conn.clone();
|
||||
let query = query.to_string();
|
||||
tokio::spawn(async move {
|
||||
let conn_guard = conn.lock().await;
|
||||
if let Err(e) = conn_guard.execute_batch(&query) {
|
||||
eprintln!("Background eviction error: {}", e);
|
||||
}
|
||||
});
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl StorageBackend for DuckDbBackend {
|
||||
async fn init(&self) -> Result<(), Status> {
|
||||
let conn = self.conn.lock().await;
|
||||
|
||||
// Create metrics table with optimized schema
|
||||
conn.execute_batch(r#"
|
||||
CREATE TABLE IF NOT EXISTS metrics (
|
||||
metric_id VARCHAR NOT NULL,
|
||||
timestamp BIGINT NOT NULL,
|
||||
value_running_window_sum DOUBLE NOT NULL,
|
||||
value_running_window_avg DOUBLE NOT NULL,
|
||||
value_running_window_count BIGINT NOT NULL,
|
||||
PRIMARY KEY (metric_id, timestamp)
|
||||
);
|
||||
|
||||
CREATE INDEX IF NOT EXISTS idx_metrics_timestamp ON metrics(timestamp);
|
||||
|
||||
CREATE TABLE IF NOT EXISTS metric_aggregations (
|
||||
metric_id VARCHAR NOT NULL,
|
||||
window_start BIGINT NOT NULL,
|
||||
window_end BIGINT NOT NULL,
|
||||
running_sum DOUBLE NOT NULL,
|
||||
running_count BIGINT NOT NULL,
|
||||
min_value DOUBLE NOT NULL,
|
||||
max_value DOUBLE NOT NULL,
|
||||
PRIMARY KEY (metric_id, window_start, window_end)
|
||||
);
|
||||
|
||||
CREATE INDEX IF NOT EXISTS idx_aggregations_window
|
||||
ON metric_aggregations(window_start, window_end);
|
||||
"#).map_err(|e| Status::internal(format!("Failed to create tables: {}", e)))?;
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn insert_metrics(&self, metrics: Vec<MetricRecord>) -> Result<(), Status> {
|
||||
if metrics.is_empty() {
|
||||
return Ok(());
|
||||
}
|
||||
|
||||
// Check if eviction is needed
|
||||
if let Some(cutoff) = self.cache_manager.should_evict().await? {
|
||||
let query = self.cache_manager.eviction_query(cutoff);
|
||||
self.execute_eviction(&query).await?;
|
||||
}
|
||||
|
||||
// Use sliding window for batch-level aggregations
|
||||
let window = TimeWindow::Sliding {
|
||||
window: Duration::from_secs(3600), // 1 hour window
|
||||
slide: Duration::from_secs(60), // 1 minute slide
|
||||
};
|
||||
|
||||
// Use optimized batch insertion
|
||||
self.insert_batch_optimized(&metrics, window).await
|
||||
}
|
||||
|
||||
async fn query_metrics(&self, from_timestamp: i64) -> Result<Vec<MetricRecord>, Status> {
|
||||
// Check if eviction is needed
|
||||
if let Some(cutoff) = self.cache_manager.should_evict().await? {
|
||||
let query = self.cache_manager.eviction_query(cutoff);
|
||||
self.execute_eviction(&query).await?;
|
||||
}
|
||||
|
||||
let query = format!(
|
||||
"SELECT metric_id, timestamp, value_running_window_sum, value_running_window_avg, value_running_window_count \
|
||||
FROM metrics WHERE timestamp >= {} ORDER BY timestamp ASC",
|
||||
from_timestamp
|
||||
);
|
||||
|
||||
let conn = self.conn.lock().await;
|
||||
let mut stmt = conn.prepare(&query)
|
||||
.map_err(|e| Status::internal(e.to_string()))?;
|
||||
|
||||
let mut rows = stmt.query(params![])
|
||||
.map_err(|e| Status::internal(e.to_string()))?;
|
||||
|
||||
let mut metrics = Vec::new();
|
||||
while let Some(row) = rows.next().map_err(|e| Status::internal(e.to_string()))? {
|
||||
metrics.push(MetricRecord {
|
||||
metric_id: row.get(0).map_err(|e| Status::internal(e.to_string()))?,
|
||||
timestamp: row.get(1).map_err(|e| Status::internal(e.to_string()))?,
|
||||
value_running_window_sum: row.get(2).map_err(|e| Status::internal(e.to_string()))?,
|
||||
value_running_window_avg: row.get(3).map_err(|e| Status::internal(e.to_string()))?,
|
||||
value_running_window_count: row.get(4).map_err(|e| Status::internal(e.to_string()))?,
|
||||
});
|
||||
}
|
||||
|
||||
Ok(metrics)
|
||||
}
|
||||
|
||||
async fn prepare_sql(&self, query: &str) -> Result<Vec<u8>, Status> {
|
||||
Ok(query.as_bytes().to_vec())
|
||||
}
|
||||
|
||||
async fn query_sql(&self, statement_handle: &[u8]) -> Result<Vec<MetricRecord>, Status> {
|
||||
let sql = std::str::from_utf8(statement_handle)
|
||||
.map_err(|e| Status::internal(e.to_string()))?;
|
||||
self.query_metrics(sql.parse().unwrap_or(0)).await
|
||||
}
|
||||
|
||||
async fn aggregate_metrics(
|
||||
&self,
|
||||
function: AggregateFunction,
|
||||
group_by: &GroupBy,
|
||||
from_timestamp: i64,
|
||||
to_timestamp: Option<i64>,
|
||||
) -> Result<Vec<AggregateResult>, Status> {
|
||||
// Check if eviction is needed
|
||||
if let Some(cutoff) = self.cache_manager.should_evict().await? {
|
||||
let query = self.cache_manager.eviction_query(cutoff);
|
||||
self.execute_eviction(&query).await?;
|
||||
}
|
||||
|
||||
let query = build_aggregate_query(
|
||||
"metrics",
|
||||
function,
|
||||
group_by,
|
||||
&["value_running_window_sum"],
|
||||
Some(from_timestamp),
|
||||
to_timestamp,
|
||||
);
|
||||
|
||||
let conn = self.conn.lock().await;
|
||||
let mut stmt = conn.prepare(&query)
|
||||
.map_err(|e| Status::internal(e.to_string()))?;
|
||||
|
||||
let mut rows = stmt.query(params![])
|
||||
.map_err(|e| Status::internal(e.to_string()))?;
|
||||
|
||||
let mut results = Vec::new();
|
||||
while let Some(row) = rows.next().map_err(|e| Status::internal(e.to_string()))? {
|
||||
let value: f64 = row.get(0).map_err(|e| Status::internal(e.to_string()))?;
|
||||
let timestamp: i64 = row.get(1).map_err(|e| Status::internal(e.to_string()))?;
|
||||
|
||||
results.push(AggregateResult {
|
||||
value,
|
||||
timestamp,
|
||||
});
|
||||
}
|
||||
|
||||
Ok(results)
|
||||
}
|
||||
|
||||
fn new_with_options(
|
||||
connection_string: &str,
|
||||
options: &HashMap<String, String>,
|
||||
credentials: Option<&Credentials>,
|
||||
) -> Result<Self, Status> {
|
||||
let mut all_options = options.clone();
|
||||
if let Some(creds) = credentials {
|
||||
all_options.insert("username".to_string(), creds.username.clone());
|
||||
all_options.insert("password".to_string(), creds.password.clone());
|
||||
}
|
||||
|
||||
let ttl = all_options.get("ttl")
|
||||
.and_then(|s| s.parse().ok())
|
||||
.map(|ttl| if ttl == 0 { None } else { Some(ttl) })
|
||||
.unwrap_or(None);
|
||||
|
||||
Self::new(connection_string.to_string(), all_options, ttl)
|
||||
}
|
||||
|
||||
async fn create_table(&self, table_name: &str, schema: &Schema) -> Result<(), Status> {
|
||||
// Create table in DuckDB
|
||||
let sql = Self::schema_to_create_table_sql(table_name, schema);
|
||||
self.execute(&sql).await?;
|
||||
|
||||
// Register table in manager
|
||||
self.table_manager.create_table(table_name.to_string(), schema.clone()).await?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn insert_into_table(&self, table_name: &str, batch: RecordBatch) -> Result<(), Status> {
|
||||
let conn = self.conn.lock().await;
|
||||
let mut stmt = conn.prepare(&format!("INSERT INTO {} VALUES ({})",
|
||||
table_name,
|
||||
(0..batch.num_columns()).map(|_| "?").collect::<Vec<_>>().join(", ")
|
||||
)).map_err(|e| Status::internal(e.to_string()))?;
|
||||
|
||||
for row_idx in 0..batch.num_rows() {
|
||||
let mut param_values: Vec<Box<dyn ToSql>> = Vec::new();
|
||||
for col_idx in 0..batch.num_columns() {
|
||||
let col = batch.column(col_idx);
|
||||
match col.data_type() {
|
||||
DataType::Int64 => {
|
||||
let array = col.as_any().downcast_ref::<Int64Array>().unwrap();
|
||||
param_values.push(Box::new(array.value(row_idx)));
|
||||
}
|
||||
DataType::Float64 => {
|
||||
let array = col.as_any().downcast_ref::<Float64Array>().unwrap();
|
||||
param_values.push(Box::new(array.value(row_idx)));
|
||||
}
|
||||
DataType::Utf8 => {
|
||||
let array = col.as_any().downcast_ref::<StringArray>().unwrap();
|
||||
param_values.push(Box::new(array.value(row_idx).to_string()));
|
||||
}
|
||||
_ => return Err(Status::internal("Unsupported column type")),
|
||||
}
|
||||
}
|
||||
|
||||
let param_refs: Vec<&dyn ToSql> = param_values.iter().map(|p| p.as_ref()).collect();
|
||||
stmt.execute(param_refs.as_slice()).map_err(|e| Status::internal(e.to_string()))?;
|
||||
}
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn query_table(&self, table_name: &str, projection: Option<Vec<String>>) -> Result<RecordBatch, Status> {
|
||||
let schema = self.table_manager.get_table_schema(table_name).await?;
|
||||
|
||||
let mut builders: Vec<Box<dyn ArrayBuilder>> = schema.fields().iter()
|
||||
.map(|field| Self::create_array_builder(field))
|
||||
.collect();
|
||||
|
||||
let projection = projection.unwrap_or_else(|| {
|
||||
schema.fields().iter().map(|f| f.name().clone()).collect()
|
||||
});
|
||||
|
||||
let sql = format!(
|
||||
"SELECT {} FROM {}",
|
||||
projection.join(", "),
|
||||
table_name
|
||||
);
|
||||
|
||||
let conn = self.conn.lock().await;
|
||||
let mut stmt = conn.prepare(&sql)
|
||||
.map_err(|e| Status::internal(e.to_string()))?;
|
||||
|
||||
let mut rows = stmt.query(params![])
|
||||
.map_err(|e| Status::internal(e.to_string()))?;
|
||||
|
||||
while let Some(row) = rows.next().map_err(|e| Status::internal(e.to_string()))? {
|
||||
for (i, field) in schema.fields().iter().enumerate() {
|
||||
match field.data_type() {
|
||||
DataType::Int64 => {
|
||||
let builder = builders[i].as_any_mut().downcast_mut::<Int64Builder>().unwrap();
|
||||
match row.get::<usize, i64>(i) {
|
||||
Ok(value) => builder.append_value(value),
|
||||
Err(_) => builder.append_null(),
|
||||
}
|
||||
}
|
||||
DataType::Float64 => {
|
||||
let builder = builders[i].as_any_mut().downcast_mut::<Float64Builder>().unwrap();
|
||||
match row.get::<usize, f64>(i) {
|
||||
Ok(value) => builder.append_value(value),
|
||||
Err(_) => builder.append_null(),
|
||||
}
|
||||
}
|
||||
DataType::Utf8 => {
|
||||
let builder = builders[i].as_any_mut().downcast_mut::<StringBuilder>().unwrap();
|
||||
match row.get::<usize, String>(i) {
|
||||
Ok(value) => builder.append_value(value),
|
||||
Err(_) => builder.append_null(),
|
||||
}
|
||||
}
|
||||
_ => return Err(Status::internal("Unsupported column type")),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
let arrays: Vec<ArrayRef> = builders.into_iter()
|
||||
.map(|mut builder| Arc::new(builder.finish()) as ArrayRef)
|
||||
.collect();
|
||||
|
||||
Ok(RecordBatch::try_new(Arc::new(schema), arrays)
|
||||
.map_err(|e| Status::internal(format!("Failed to create record batch: {}", e)))?)
|
||||
}
|
||||
|
||||
async fn create_aggregation_view(&self, view: &AggregationView) -> Result<(), Status> {
|
||||
let columns: Vec<&str> = view.aggregate_columns.iter()
|
||||
.map(|s| s.as_str())
|
||||
.collect();
|
||||
|
||||
let sql = build_aggregate_query(
|
||||
&view.source_table,
|
||||
view.function,
|
||||
&view.group_by,
|
||||
&columns,
|
||||
None,
|
||||
None
|
||||
);
|
||||
|
||||
let view_name = format!("agg_view_{}", view.source_table);
|
||||
let conn = self.conn.lock().await;
|
||||
conn.execute(&format!("CREATE VIEW {} AS {}", view_name, sql), params![])
|
||||
.map_err(|e| Status::internal(format!("Failed to create view: {}", e)))?;
|
||||
|
||||
// Register view in manager
|
||||
self.table_manager.create_aggregation_view(
|
||||
view_name,
|
||||
view.source_table.clone(),
|
||||
view.function.clone(),
|
||||
view.group_by.clone(),
|
||||
view.window.clone(),
|
||||
view.aggregate_columns.clone(),
|
||||
).await?;
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn query_aggregation_view(&self, view_name: &str) -> Result<RecordBatch, Status> {
|
||||
self.query_table(view_name, None).await
|
||||
}
|
||||
|
||||
async fn drop_table(&self, table_name: &str) -> Result<(), Status> {
|
||||
let conn = self.conn.lock().await;
|
||||
conn.execute(&format!("DROP TABLE IF EXISTS {}", table_name), params![])
|
||||
.map_err(|e| Status::internal(format!("Failed to drop table: {}", e)))?;
|
||||
|
||||
self.table_manager.drop_table(table_name).await?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn drop_aggregation_view(&self, view_name: &str) -> Result<(), Status> {
|
||||
let conn = self.conn.lock().await;
|
||||
conn.execute(&format!("DROP VIEW IF EXISTS {}", view_name), params![])
|
||||
.map_err(|e| Status::internal(format!("Failed to drop view: {}", e)))?;
|
||||
|
||||
self.table_manager.drop_aggregation_view(view_name).await?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn table_manager(&self) -> &TableManager {
|
||||
&self.table_manager
|
||||
}
|
||||
}
|
||||
|
||||
impl DuckDbBackend {
|
||||
/// Executes a SQL query.
|
||||
async fn execute(&self, query: &str) -> Result<(), Status> {
|
||||
let conn = self.conn.lock().await;
|
||||
conn.execute(query, params![])
|
||||
.map_err(|e| Status::internal(e.to_string()))?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Converts an Arrow schema to a DuckDB CREATE TABLE statement
|
||||
fn schema_to_create_table_sql(table_name: &str, schema: &Schema) -> String {
|
||||
let mut sql = format!("CREATE TABLE IF NOT EXISTS \"{}\" (", table_name);
|
||||
let mut first = true;
|
||||
|
||||
for field in schema.fields() {
|
||||
if !first {
|
||||
sql.push_str(", ");
|
||||
}
|
||||
first = false;
|
||||
|
||||
sql.push_str(&format!("\"{}\" {}", field.name(), Self::arrow_type_to_duckdb_type(field.data_type())));
|
||||
}
|
||||
|
||||
sql.push_str(")");
|
||||
sql
|
||||
}
|
||||
|
||||
/// Converts an Arrow data type to a DuckDB type string
|
||||
fn arrow_type_to_duckdb_type(data_type: &DataType) -> &'static str {
|
||||
match data_type {
|
||||
DataType::Boolean => "BOOLEAN",
|
||||
DataType::Int8 => "TINYINT",
|
||||
DataType::Int16 => "SMALLINT",
|
||||
DataType::Int32 => "INTEGER",
|
||||
DataType::Int64 => "BIGINT",
|
||||
DataType::UInt8 => "TINYINT",
|
||||
DataType::UInt16 => "SMALLINT",
|
||||
DataType::UInt32 => "INTEGER",
|
||||
DataType::UInt64 => "BIGINT",
|
||||
DataType::Float32 => "REAL",
|
||||
DataType::Float64 => "DOUBLE",
|
||||
DataType::Utf8 => "VARCHAR",
|
||||
DataType::Binary => "BLOB",
|
||||
DataType::Date32 => "DATE",
|
||||
DataType::Date64 => "DATE",
|
||||
DataType::Time32(_) => "TIME",
|
||||
DataType::Time64(_) => "TIME",
|
||||
DataType::Timestamp(_, _) => "TIMESTAMP",
|
||||
_ => "VARCHAR", // Default to VARCHAR for unsupported types
|
||||
}
|
||||
}
|
||||
|
||||
fn create_array_builder(field: &Field) -> Box<dyn ArrayBuilder> {
|
||||
match field.data_type() {
|
||||
DataType::Int64 => Box::new(Int64Builder::new()),
|
||||
DataType::Float64 => Box::new(Float64Builder::new()),
|
||||
DataType::Utf8 => Box::new(StringBuilder::new()),
|
||||
_ => panic!("Unsupported column type"),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
+170
@@ -0,0 +1,170 @@
|
||||
//! Storage backends for metric data persistence and caching.
|
||||
//!
|
||||
//! This module provides multiple storage backend implementations:
|
||||
//! - `duckdb`: High-performance embedded database for caching and local storage
|
||||
//! - `adbc`: Arrow Database Connectivity for external database integration
|
||||
//! - `cached`: Two-tier storage with configurable caching layer
|
||||
//!
|
||||
//! Each backend implements the `StorageBackend` trait, providing a consistent
|
||||
//! interface for metric storage and retrieval operations.
|
||||
|
||||
pub mod adbc;
|
||||
pub mod duckdb;
|
||||
pub mod cache;
|
||||
pub mod table_manager;
|
||||
|
||||
use crate::config::Credentials;
|
||||
use crate::metrics::MetricRecord;
|
||||
use crate::aggregation::{AggregateFunction, GroupBy, AggregateResult, TimeWindow};
|
||||
use crate::storage::table_manager::{TableManager, AggregationView};
|
||||
use async_trait::async_trait;
|
||||
use std::collections::HashMap;
|
||||
use tonic::Status;
|
||||
use arrow_schema::Schema;
|
||||
use arrow_array::RecordBatch;
|
||||
|
||||
/// Batch-level aggregation state for efficient updates
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct BatchAggregation {
|
||||
/// The metric ID this aggregation belongs to
|
||||
pub metric_id: String,
|
||||
/// Start of the time window
|
||||
pub window_start: i64,
|
||||
/// End of the time window
|
||||
pub window_end: i64,
|
||||
/// Running sum within the window
|
||||
pub running_sum: f64,
|
||||
/// Running count within the window
|
||||
pub running_count: i64,
|
||||
/// Minimum value in the window
|
||||
pub min_value: f64,
|
||||
/// Maximum value in the window
|
||||
pub max_value: f64,
|
||||
}
|
||||
|
||||
/// Storage backend trait for metric data persistence.
|
||||
///
|
||||
/// This trait defines the interface that all storage backends must implement.
|
||||
/// It provides methods for:
|
||||
/// - Initialization and configuration
|
||||
/// - Metric data insertion
|
||||
/// - Metric data querying
|
||||
/// - SQL query preparation and execution
|
||||
/// - Aggregation of metrics
|
||||
/// - Table and view management
|
||||
#[async_trait]
|
||||
pub trait StorageBackend: Send + Sync + 'static {
|
||||
/// Initialize the storage backend.
|
||||
async fn init(&self) -> Result<(), Status>;
|
||||
|
||||
/// Insert metrics into storage.
|
||||
async fn insert_metrics(&self, metrics: Vec<MetricRecord>) -> Result<(), Status>;
|
||||
|
||||
/// Query metrics from storage.
|
||||
async fn query_metrics(&self, from_timestamp: i64) -> Result<Vec<MetricRecord>, Status>;
|
||||
|
||||
/// Prepare a SQL query and return a handle.
|
||||
/// The handle is backend-specific and opaque to the caller.
|
||||
async fn prepare_sql(&self, query: &str) -> Result<Vec<u8>, Status>;
|
||||
|
||||
/// Execute a prepared SQL query using its handle.
|
||||
/// The handle must have been obtained from prepare_sql.
|
||||
async fn query_sql(&self, statement_handle: &[u8]) -> Result<Vec<MetricRecord>, Status>;
|
||||
|
||||
/// Aggregate metrics using the specified function and grouping.
|
||||
async fn aggregate_metrics(
|
||||
&self,
|
||||
function: AggregateFunction,
|
||||
group_by: &GroupBy,
|
||||
from_timestamp: i64,
|
||||
to_timestamp: Option<i64>,
|
||||
) -> Result<Vec<AggregateResult>, Status>;
|
||||
|
||||
/// Create a new instance with the given options.
|
||||
/// The connection string and options are backend-specific.
|
||||
fn new_with_options(
|
||||
connection_string: &str,
|
||||
options: &HashMap<String, String>,
|
||||
credentials: Option<&Credentials>,
|
||||
) -> Result<Self, Status>
|
||||
where
|
||||
Self: Sized;
|
||||
|
||||
/// Create a new table with the given schema
|
||||
async fn create_table(&self, table_name: &str, schema: &Schema) -> Result<(), Status>;
|
||||
|
||||
/// Insert data into a table
|
||||
async fn insert_into_table(&self, table_name: &str, batch: RecordBatch) -> Result<(), Status>;
|
||||
|
||||
/// Query data from a table
|
||||
async fn query_table(&self, table_name: &str, projection: Option<Vec<String>>) -> Result<RecordBatch, Status>;
|
||||
|
||||
/// Create an aggregation view
|
||||
async fn create_aggregation_view(&self, view: &AggregationView) -> Result<(), Status>;
|
||||
|
||||
/// Query data from an aggregation view
|
||||
async fn query_aggregation_view(&self, view_name: &str) -> Result<RecordBatch, Status>;
|
||||
|
||||
/// Drop a table
|
||||
async fn drop_table(&self, table_name: &str) -> Result<(), Status>;
|
||||
|
||||
/// Drop an aggregation view
|
||||
async fn drop_aggregation_view(&self, view_name: &str) -> Result<(), Status>;
|
||||
|
||||
/// Get the table manager instance
|
||||
fn table_manager(&self) -> &TableManager;
|
||||
|
||||
/// Update batch-level aggregations.
|
||||
/// This is called during batch writes to maintain running aggregations.
|
||||
async fn update_batch_aggregations(
|
||||
&self,
|
||||
batch: &[MetricRecord],
|
||||
window: TimeWindow,
|
||||
) -> Result<Vec<BatchAggregation>, Status> {
|
||||
// Default implementation that processes the batch and updates aggregations
|
||||
let mut aggregations = HashMap::new();
|
||||
|
||||
for metric in batch {
|
||||
let (window_start, window_end) = window.window_bounds(metric.timestamp);
|
||||
let key = (metric.metric_id.clone(), window_start, window_end);
|
||||
|
||||
let agg = aggregations.entry(key).or_insert_with(|| BatchAggregation {
|
||||
metric_id: metric.metric_id.clone(),
|
||||
window_start,
|
||||
window_end,
|
||||
running_sum: 0.0,
|
||||
running_count: 0,
|
||||
min_value: f64::INFINITY,
|
||||
max_value: f64::NEG_INFINITY,
|
||||
});
|
||||
|
||||
// Update running aggregations
|
||||
agg.running_sum += metric.value_running_window_sum;
|
||||
agg.running_count += 1;
|
||||
agg.min_value = agg.min_value.min(metric.value_running_window_sum);
|
||||
agg.max_value = agg.max_value.max(metric.value_running_window_sum);
|
||||
}
|
||||
|
||||
Ok(aggregations.into_values().collect())
|
||||
}
|
||||
|
||||
/// Insert batch-level aggregations.
|
||||
/// This is called after update_batch_aggregations to persist the aggregations.
|
||||
async fn insert_batch_aggregations(
|
||||
&self,
|
||||
aggregations: Vec<BatchAggregation>,
|
||||
) -> Result<(), Status> {
|
||||
// Default implementation that stores aggregations in a separate table
|
||||
let mut batch = Vec::new();
|
||||
for agg in aggregations {
|
||||
batch.push(MetricRecord {
|
||||
metric_id: agg.metric_id,
|
||||
timestamp: agg.window_start,
|
||||
value_running_window_sum: agg.running_sum,
|
||||
value_running_window_avg: agg.running_sum / agg.running_count as f64,
|
||||
value_running_window_count: agg.running_count,
|
||||
});
|
||||
}
|
||||
self.insert_metrics(batch).await
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,140 @@
|
||||
use std::collections::HashMap;
|
||||
use std::sync::Arc;
|
||||
use tokio::sync::RwLock;
|
||||
use arrow_schema::Schema;
|
||||
use tonic::Status;
|
||||
use serde::{Serialize, Deserialize};
|
||||
use crate::aggregation::{TimeWindow, AggregateFunction, GroupBy};
|
||||
|
||||
/// Configuration for an aggregation view
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct AggregationView {
|
||||
pub source_table: String,
|
||||
pub function: AggregateFunction,
|
||||
pub group_by: GroupBy,
|
||||
pub window: TimeWindow,
|
||||
pub aggregate_columns: Vec<String>,
|
||||
}
|
||||
|
||||
#[derive(Debug)]
|
||||
pub struct TableManager {
|
||||
tables: Arc<RwLock<HashMap<String, Schema>>>,
|
||||
views: Arc<RwLock<HashMap<String, AggregationView>>>,
|
||||
}
|
||||
|
||||
impl Clone for TableManager {
|
||||
fn clone(&self) -> Self {
|
||||
Self {
|
||||
tables: self.tables.clone(),
|
||||
views: self.views.clone(),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl Default for TableManager {
|
||||
fn default() -> Self {
|
||||
Self::new()
|
||||
}
|
||||
}
|
||||
|
||||
impl TableManager {
|
||||
pub fn new() -> Self {
|
||||
Self {
|
||||
tables: Arc::new(RwLock::new(HashMap::new())),
|
||||
views: Arc::new(RwLock::new(HashMap::new())),
|
||||
}
|
||||
}
|
||||
|
||||
pub async fn create_table(&self, name: String, schema: Schema) -> Result<(), Status> {
|
||||
let mut tables = self.tables.write().await;
|
||||
if tables.contains_key(&name) {
|
||||
return Err(Status::already_exists(format!("Table {} already exists", name)));
|
||||
}
|
||||
tables.insert(name, schema);
|
||||
Ok(())
|
||||
}
|
||||
|
||||
pub async fn get_table_schema(&self, name: &str) -> Result<Schema, Status> {
|
||||
let tables = self.tables.read().await;
|
||||
tables.get(name)
|
||||
.cloned()
|
||||
.ok_or_else(|| Status::not_found(format!("Table {} not found", name)))
|
||||
}
|
||||
|
||||
pub async fn create_aggregation_view(
|
||||
&self,
|
||||
name: String,
|
||||
source_table: String,
|
||||
function: AggregateFunction,
|
||||
group_by: GroupBy,
|
||||
window: TimeWindow,
|
||||
aggregate_columns: Vec<String>,
|
||||
) -> Result<(), Status> {
|
||||
// Verify source table exists
|
||||
{
|
||||
let tables = self.tables.read().await;
|
||||
if !tables.contains_key(&source_table) {
|
||||
return Err(Status::not_found(format!("Source table {} not found", source_table)));
|
||||
}
|
||||
|
||||
// Verify aggregate columns exist in source table
|
||||
let schema = tables.get(&source_table).unwrap();
|
||||
for col in &aggregate_columns {
|
||||
if !schema.fields().iter().any(|f| f.name() == col) {
|
||||
return Err(Status::invalid_argument(format!(
|
||||
"Column {} not found in source table {}",
|
||||
col, source_table
|
||||
)));
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
let view = AggregationView {
|
||||
source_table,
|
||||
function,
|
||||
group_by,
|
||||
window,
|
||||
aggregate_columns,
|
||||
};
|
||||
|
||||
let mut views = self.views.write().await;
|
||||
if views.contains_key(&name) {
|
||||
return Err(Status::already_exists(format!("View {} already exists", name)));
|
||||
}
|
||||
views.insert(name, view);
|
||||
Ok(())
|
||||
}
|
||||
|
||||
pub async fn get_aggregation_view(&self, name: &str) -> Result<AggregationView, Status> {
|
||||
let views = self.views.read().await;
|
||||
views.get(name)
|
||||
.cloned()
|
||||
.ok_or_else(|| Status::not_found(format!("View {} not found", name)))
|
||||
}
|
||||
|
||||
pub async fn list_tables(&self) -> Vec<String> {
|
||||
let tables = self.tables.read().await;
|
||||
tables.keys().cloned().collect()
|
||||
}
|
||||
|
||||
pub async fn list_aggregation_views(&self) -> Vec<String> {
|
||||
let views = self.views.read().await;
|
||||
views.keys().cloned().collect()
|
||||
}
|
||||
|
||||
pub async fn drop_table(&self, name: &str) -> Result<(), Status> {
|
||||
let mut tables = self.tables.write().await;
|
||||
if tables.remove(name).is_none() {
|
||||
return Err(Status::not_found(format!("Table {} not found", name)));
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
pub async fn drop_aggregation_view(&self, name: &str) -> Result<(), Status> {
|
||||
let mut views = self.views.write().await;
|
||||
if views.remove(name).is_none() {
|
||||
return Err(Status::not_found(format!("View {} not found", name)));
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user