use std::collections::HashMap;
use std::fmt;
use std::sync::Arc;
use dashmap::DashMap;
use serde::{Deserialize, Serialize};
use crate::error::SchemaError;
use crate::types::FieldValue;
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)]
pub enum ColumnRole {
Timestamp,
Tag,
Field,
}
impl fmt::Display for ColumnRole {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
match self {
Self::Timestamp => write!(f, "timestamp"),
Self::Tag => write!(f, "tag"),
Self::Field => write!(f, "field"),
}
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)]
pub enum ColumnType {
Timestamp,
String,
F64,
I64,
U64,
Bool,
}
impl ColumnType {
#[inline]
#[must_use]
pub fn from_field_value(value: &FieldValue) -> Self {
match value {
FieldValue::F64(_) => Self::F64,
FieldValue::I64(_) => Self::I64,
FieldValue::U64(_) => Self::U64,
FieldValue::Bool(_) => Self::Bool,
FieldValue::String(_) => Self::String,
}
}
}
impl fmt::Display for ColumnType {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
match self {
Self::Timestamp => write!(f, "timestamp"),
Self::String => write!(f, "string"),
Self::F64 => write!(f, "f64"),
Self::I64 => write!(f, "i64"),
Self::U64 => write!(f, "u64"),
Self::Bool => write!(f, "bool"),
}
}
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct ColumnDef {
pub name: String,
pub column_type: ColumnType,
pub role: ColumnRole,
}
#[derive(Debug, Clone, Serialize)]
pub struct MeasurementSchema {
measurement: String,
columns: Vec<ColumnDef>,
#[serde(skip)]
column_index: HashMap<String, usize>,
}
impl PartialEq for MeasurementSchema {
fn eq(&self, other: &Self) -> bool {
self.measurement == other.measurement && self.columns == other.columns
}
}
impl Eq for MeasurementSchema {}
impl<'de> serde::Deserialize<'de> for MeasurementSchema {
fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
where
D: serde::Deserializer<'de>,
{
#[derive(Deserialize)]
struct Helper {
measurement: String,
columns: Vec<ColumnDef>,
}
let h = Helper::deserialize(deserializer)?;
let column_index = h
.columns
.iter()
.enumerate()
.map(|(i, c)| (c.name.clone(), i))
.collect();
Ok(Self {
measurement: h.measurement,
columns: h.columns,
column_index,
})
}
}
impl MeasurementSchema {
#[must_use]
pub fn new(measurement: impl Into<String>) -> Self {
let ts_col = ColumnDef {
name: crate::TIME_COLUMN.to_string(),
column_type: ColumnType::Timestamp,
role: ColumnRole::Timestamp,
};
let mut column_index = HashMap::new();
column_index.insert(ts_col.name.clone(), 0);
Self {
measurement: measurement.into(),
columns: vec![ts_col],
column_index,
}
}
#[inline]
#[must_use]
pub fn measurement(&self) -> &str {
&self.measurement
}
#[inline]
#[must_use]
pub fn columns(&self) -> &[ColumnDef] {
&self.columns
}
#[must_use]
pub fn add_tag(&mut self, name: &str) -> bool {
if let Some(existing) = self.columns.iter().find(|c| c.name == name) {
if existing.role != ColumnRole::Tag {
tracing::error!(
measurement = %self.measurement,
column = name,
existing_role = ?existing.role,
"tag name conflicts with existing column role — rejecting"
);
}
return false;
}
let idx = self.columns.len();
self.columns.push(ColumnDef {
name: name.to_string(),
column_type: ColumnType::String,
role: ColumnRole::Tag,
});
self.column_index.insert(name.to_string(), idx);
true
}
pub fn add_field(&mut self, name: &str, field_value: &FieldValue) -> Result<bool, SchemaError> {
let new_type = ColumnType::from_field_value(field_value);
if let Some(existing) = self.columns.iter().find(|c| c.name == name) {
if existing.role != ColumnRole::Field {
return Err(SchemaError::TypeConflict {
measurement: self.measurement.clone(),
field: name.to_string(),
expected: format!("{:?}", existing.role),
got: "Field".to_string(),
});
}
if existing.column_type != new_type {
return Err(SchemaError::TypeConflict {
measurement: self.measurement.clone(),
field: name.to_string(),
expected: existing.column_type.to_string(),
got: new_type.to_string(),
});
}
return Ok(false);
}
let idx = self.columns.len();
self.columns.push(ColumnDef {
name: name.to_string(),
column_type: new_type,
role: ColumnRole::Field,
});
self.column_index.insert(name.to_string(), idx);
Ok(true)
}
#[inline]
#[must_use]
pub fn column(&self, name: &str) -> Option<&ColumnDef> {
self.column_index
.get(name)
.and_then(|&i| self.columns.get(i))
}
pub fn push_column(&mut self, col: ColumnDef) {
if self.column_index.contains_key(&col.name) {
return; }
let idx = self.columns.len();
self.column_index.insert(col.name.clone(), idx);
self.columns.push(col);
}
#[must_use]
pub fn tag_count(&self) -> usize {
self.columns
.iter()
.filter(|c| c.role == ColumnRole::Tag)
.count()
}
#[must_use]
pub fn field_count(&self) -> usize {
self.columns
.iter()
.filter(|c| c.role == ColumnRole::Field)
.count()
}
#[must_use]
pub fn tag_names(&self) -> Vec<&str> {
self.columns
.iter()
.filter(|c| c.role == ColumnRole::Tag)
.map(|c| c.name.as_str())
.collect()
}
#[must_use]
pub fn field_names(&self) -> Vec<&str> {
self.columns
.iter()
.filter(|c| c.role == ColumnRole::Field)
.map(|c| c.name.as_str())
.collect()
}
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub enum SchemaAction {
CreateMeasurement(MeasurementSchema),
AddColumn {
measurement: String,
column: ColumnDef,
},
}
#[derive(Debug, Clone)]
pub struct SchemaRegistry {
schemas: Arc<DashMap<String, Arc<MeasurementSchema>>>,
changes: Arc<std::sync::Mutex<()>>,
}
impl SchemaRegistry {
#[must_use]
pub fn new() -> Self {
Self {
schemas: Arc::new(DashMap::new()),
changes: Arc::new(std::sync::Mutex::new(())),
}
}
pub fn register_batch(
&self,
points: &[&crate::types::Point],
) -> Result<Vec<SchemaAction>, SchemaError> {
let mut all_known = true;
for point in points {
match self.schemas.get(point.series_key().measurement()) {
Some(schema) => {
if !Self::point_fits(&schema, point)? {
all_known = false;
}
}
None => all_known = false,
}
}
if all_known {
return Ok(Vec::new());
}
let _guard = self
.changes
.lock()
.unwrap_or_else(std::sync::PoisonError::into_inner);
let mut pending: std::collections::BTreeMap<String, MeasurementSchema> =
std::collections::BTreeMap::new();
let mut actions = Vec::new();
for point in points {
let name = point.series_key().measurement();
let schema = match pending.get_mut(name) {
Some(s) => s,
None => {
let fresh = match self.schemas.get(name) {
Some(existing) => (**existing).clone(),
None => {
let s = MeasurementSchema::new(name);
actions.push(SchemaAction::CreateMeasurement(s.clone()));
s
}
};
pending.entry(name.to_string()).or_insert(fresh)
}
};
for tag_key in point.series_key().tag_keys() {
if schema.add_tag(tag_key) {
actions.push(SchemaAction::AddColumn {
measurement: name.to_string(),
column: ColumnDef {
name: tag_key.to_string(),
column_type: ColumnType::String,
role: ColumnRole::Tag,
},
});
}
}
for (field_name, field_value) in point.fields() {
if schema.add_field(field_name.as_ref(), field_value)? {
actions.push(SchemaAction::AddColumn {
measurement: name.to_string(),
column: ColumnDef {
name: field_name.to_string(),
column_type: ColumnType::from_field_value(field_value),
role: ColumnRole::Field,
},
});
}
}
}
for (name, schema) in pending {
self.schemas.insert(name, Arc::new(schema));
}
Ok(actions)
}
fn point_fits(
schema: &MeasurementSchema,
point: &crate::types::Point,
) -> Result<bool, SchemaError> {
for tag_key in point.series_key().tag_keys() {
if schema.column(tag_key).is_none() {
return Ok(false);
}
}
for (field_name, field_value) in point.fields() {
match schema.column(field_name.as_ref()) {
Some(existing) => {
let new_type = ColumnType::from_field_value(field_value);
if existing.column_type != new_type {
return Err(SchemaError::TypeConflict {
measurement: schema.measurement().to_string(),
field: field_name.to_string(),
expected: existing.column_type.to_string(),
got: new_type.to_string(),
});
}
}
None => return Ok(false),
}
}
Ok(true)
}
pub fn register_point(
&self,
point: &crate::types::Point,
) -> Result<Vec<SchemaAction>, SchemaError> {
self.register_batch(&[point])
}
#[must_use]
pub fn lookup(&self, measurement: &str) -> Option<Arc<MeasurementSchema>> {
self.schemas.get(measurement).map(|r| r.value().clone())
}
pub fn register_schema(&self, name: &str, schema: MeasurementSchema) {
self.schemas
.entry(name.to_string())
.or_insert(Arc::new(schema));
}
pub fn register_measurement(&self, ms: MeasurementSchema) {
self.schemas
.entry(ms.measurement().to_string())
.or_insert(Arc::new(ms));
}
pub fn apply_add_column(&self, measurement: &str, column: ColumnDef) {
if let Some(mut schema_arc) = self.schemas.get_mut(measurement) {
if let Some(existing) = schema_arc
.value()
.columns()
.iter()
.find(|c| c.name == column.name)
{
if existing.role != column.role {
tracing::warn!(
measurement,
column = %column.name,
existing_role = ?existing.role,
new_role = ?column.role,
"WAL replay column role conflict — skipping"
);
}
return;
}
let schema = Arc::make_mut(schema_arc.value_mut());
match column.role {
ColumnRole::Tag => {
let _ = schema.add_tag(&column.name);
}
ColumnRole::Field => {
schema.push_column(column);
}
ColumnRole::Timestamp => {
}
}
}
}
#[must_use]
pub fn measurement_count(&self) -> usize {
self.schemas.len()
}
#[must_use]
pub fn measurement_names(&self) -> Vec<String> {
self.schemas.iter().map(|r| r.key().clone()).collect()
}
#[must_use]
pub fn remove(&self, measurement: &str) -> Option<Arc<MeasurementSchema>> {
self.schemas.remove(measurement).map(|(_, v)| v)
}
}
impl Default for SchemaRegistry {
fn default() -> Self {
Self::new()
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::types::{FieldValue, Point, SeriesKey};
use std::collections::BTreeMap;
fn make_point(
measurement: &str,
tags: &[(&str, &str)],
fields: &[(&str, FieldValue)],
ts: i64,
) -> Point {
let tag_map: BTreeMap<String, String> = tags
.iter()
.map(|(k, v)| ((*k).to_string(), (*v).to_string()))
.collect();
let field_map: BTreeMap<String, FieldValue> = fields
.iter()
.map(|(k, v)| ((*k).to_string(), v.clone()))
.collect();
let key = SeriesKey::new(measurement, tag_map).unwrap();
Point::new(key, field_map, ts).unwrap()
}
#[test]
fn measurement_schema_new() {
let schema = MeasurementSchema::new("cpu");
assert_eq!(schema.measurement(), "cpu");
assert_eq!(schema.columns().len(), 1);
assert_eq!(schema.columns()[0].name, crate::TIME_COLUMN);
assert_eq!(schema.columns()[0].role, ColumnRole::Timestamp);
}
#[test]
fn measurement_schema_add_tag() {
let mut schema = MeasurementSchema::new("cpu");
assert!(schema.add_tag("host"));
assert!(!schema.add_tag("host")); assert_eq!(schema.tag_count(), 1);
assert_eq!(schema.columns()[1].name, "host");
assert_eq!(schema.columns()[1].role, ColumnRole::Tag);
}
#[test]
fn measurement_schema_add_field() {
let mut schema = MeasurementSchema::new("cpu");
let added = schema.add_field("value", &FieldValue::F64(1.0)).unwrap();
assert!(added);
let added = schema.add_field("value", &FieldValue::F64(2.0)).unwrap();
assert!(!added); assert_eq!(schema.field_count(), 1);
}
#[test]
fn measurement_schema_type_conflict() {
let mut schema = MeasurementSchema::new("cpu");
schema.add_field("value", &FieldValue::F64(1.0)).unwrap();
let err = schema.add_field("value", &FieldValue::I64(1)).unwrap_err();
assert!(matches!(err, SchemaError::TypeConflict { .. }));
}
#[test]
fn schema_registry_auto_create() {
let registry = SchemaRegistry::new();
let point = make_point(
"cpu",
&[("host", "srv1")],
&[("value", FieldValue::F64(72.5))],
1_000_000_000,
);
let actions = registry.register_point(&point).unwrap();
assert_eq!(actions.len(), 3);
assert!(matches!(&actions[0], SchemaAction::CreateMeasurement(_)));
let schema = registry.lookup("cpu").unwrap();
assert_eq!(schema.measurement(), "cpu");
assert_eq!(schema.tag_count(), 1);
assert_eq!(schema.field_count(), 1);
}
#[test]
fn schema_registry_additive_evolution() {
let registry = SchemaRegistry::new();
let p1 = make_point(
"cpu",
&[("host", "srv1")],
&[("value", FieldValue::F64(72.5))],
1,
);
registry.register_point(&p1).unwrap();
let p2 = make_point(
"cpu",
&[("host", "srv1"), ("dc", "eu")],
&[
("value", FieldValue::F64(73.0)),
("cores", FieldValue::I64(8)),
],
2,
);
let actions = registry.register_point(&p2).unwrap();
assert_eq!(actions.len(), 2);
let schema = registry.lookup("cpu").unwrap();
assert_eq!(schema.tag_count(), 2);
assert_eq!(schema.field_count(), 2);
}
#[test]
fn schema_registry_type_conflict_rejected() {
let registry = SchemaRegistry::new();
let p1 = make_point("cpu", &[], &[("value", FieldValue::F64(72.5))], 1);
registry.register_point(&p1).unwrap();
let p2 = make_point("cpu", &[], &[("value", FieldValue::I64(72))], 2);
let err = registry.register_point(&p2).unwrap_err();
assert!(matches!(err, SchemaError::TypeConflict { .. }));
}
#[test]
fn schema_registry_multiple_measurements() {
let registry = SchemaRegistry::new();
let p1 = make_point("cpu", &[], &[("value", FieldValue::F64(72.5))], 1);
let p2 = make_point("mem", &[], &[("used", FieldValue::U64(1024))], 1);
registry.register_point(&p1).unwrap();
registry.register_point(&p2).unwrap();
assert_eq!(registry.measurement_count(), 2);
assert!(registry.lookup("cpu").is_some());
assert!(registry.lookup("mem").is_some());
assert!(registry.lookup("disk").is_none());
}
#[test]
fn schema_registry_remove() {
let registry = SchemaRegistry::new();
let p = make_point("cpu", &[], &[("value", FieldValue::F64(1.0))], 1);
registry.register_point(&p).unwrap();
assert!(registry.lookup("cpu").is_some());
let removed = registry.remove("cpu");
assert!(removed.is_some());
assert!(registry.lookup("cpu").is_none());
assert_eq!(registry.measurement_count(), 0);
}
#[test]
fn schema_registry_measurement_names() {
let registry = SchemaRegistry::new();
let p1 = make_point("cpu", &[], &[("v", FieldValue::F64(1.0))], 1);
let p2 = make_point("mem", &[], &[("v", FieldValue::U64(1))], 1);
registry.register_point(&p1).unwrap();
registry.register_point(&p2).unwrap();
let mut names = registry.measurement_names();
names.sort();
assert_eq!(names, vec!["cpu", "mem"]);
}
#[test]
fn schema_serde_roundtrip() {
let mut schema = MeasurementSchema::new("cpu");
let _ = schema.add_tag("host");
schema.add_field("value", &FieldValue::F64(1.0)).unwrap();
let json = serde_json::to_string(&schema).unwrap();
let back: MeasurementSchema = serde_json::from_str(&json).unwrap();
assert_eq!(schema, back);
}
#[test]
fn column_type_from_field_value() {
assert_eq!(
ColumnType::from_field_value(&FieldValue::F64(1.0)),
ColumnType::F64
);
assert_eq!(
ColumnType::from_field_value(&FieldValue::I64(1)),
ColumnType::I64
);
assert_eq!(
ColumnType::from_field_value(&FieldValue::U64(1)),
ColumnType::U64
);
assert_eq!(
ColumnType::from_field_value(&FieldValue::Bool(true)),
ColumnType::Bool
);
assert_eq!(
ColumnType::from_field_value(&FieldValue::String("x".into())),
ColumnType::String
);
}
#[test]
fn column_type_display() {
assert_eq!(ColumnType::Timestamp.to_string(), "timestamp");
assert_eq!(ColumnType::String.to_string(), "string");
assert_eq!(ColumnType::F64.to_string(), "f64");
assert_eq!(ColumnType::I64.to_string(), "i64");
assert_eq!(ColumnType::U64.to_string(), "u64");
assert_eq!(ColumnType::Bool.to_string(), "bool");
}
#[test]
fn column_role_display() {
assert_eq!(ColumnRole::Timestamp.to_string(), "timestamp");
assert_eq!(ColumnRole::Tag.to_string(), "tag");
assert_eq!(ColumnRole::Field.to_string(), "field");
}
#[test]
fn measurement_schema_column_not_found() {
let schema = MeasurementSchema::new("cpu");
assert!(schema.column("nonexistent").is_none());
}
#[test]
fn schema_registry_remove_nonexistent() {
let registry = SchemaRegistry::new();
let removed = registry.remove("nope");
assert!(removed.is_none());
}
#[test]
fn schema_action_serde_roundtrip() {
let action = SchemaAction::CreateMeasurement(MeasurementSchema::new("cpu"));
let json = serde_json::to_string(&action).unwrap();
let back: SchemaAction = serde_json::from_str(&json).unwrap();
assert_eq!(action, back);
let action = SchemaAction::AddColumn {
measurement: "cpu".into(),
column: ColumnDef {
name: "value".into(),
column_type: ColumnType::F64,
role: ColumnRole::Field,
},
};
let json = serde_json::to_string(&action).unwrap();
let back: SchemaAction = serde_json::from_str(&json).unwrap();
assert_eq!(action, back);
}
#[test]
fn measurement_schema_accessors() {
let mut schema = MeasurementSchema::new("cpu");
assert!(schema.add_tag("host"));
assert_eq!(schema.measurement(), "cpu");
assert_eq!(schema.columns().len(), 2);
}
}