use std::collections::HashMap; use std::path::PathBuf; use std::sync::atomic::{AtomicBool, Ordering}; use std::sync::{Arc, Mutex}; use crate::catalog::Catalog; use crate::config::{self, Config, Loaded}; use crate::heap::HeapTable; use crate::initsql::{self, InitColumn, InitTable}; use crate::log; use crate::protocol::{ reserved_table, Account, ColumnStatWire, IndexWire, Request, Response, Row, SchemaColumn, TableSchemaWire, USERS_TABLE, }; use crate::wal::{Wal, WalOp}; use crate::{log_debug, log_error, log_info, log_warn}; // 进程内存储运行时:Haskell 侧经 ffi.rs 调进来。 // request_line 的语义等同原来的「一行 JSON 请求 → 一行 JSON 响应」套接字协议。 struct Server { cfg: Config, tables: Mutex>>>, write_lock: Mutex<()>, catalog: Mutex, wal: Wal, recovery_required: AtomicBool, } type TableDescription = (Vec, u64, Vec, Vec); impl Server { /// 新建服务器状态 fn new(cfg: Config) -> std::io::Result { std::fs::create_dir_all(&cfg.data_dir)?; let catalog = Catalog::load(cfg.data_dir.join("catalog.json"))?; let wal = Wal::new(cfg.data_dir.join("wal.log")); Ok(Server { cfg, tables: Mutex::new(HashMap::new()), write_lock: Mutex::new(()), catalog: Mutex::new(catalog), wal, recovery_required: AtomicBool::new(false), }) } /// 表名 -> 数据文件 fn table_path(&self, table: &str) -> PathBuf { self.cfg.data_dir.join(format!("{}.db", table)) } /// (表名, 列名) 对应哪个索引文件 fn index_path(&self, table: &str, column: &str) -> PathBuf { if column == "id" { self.cfg.data_dir.join(format!("{}.idx", table)) } else { self.cfg .data_dir .join(format!("{}.{}.idx", table, column)) } } /// 数据字典文件 fn catalog_path(&self) -> PathBuf { self.cfg.data_dir.join("catalog.json") } /// 仅数据字典决定表是否存在 fn table_exists(&self, table: &str) -> std::io::Result { let catalog = self.catalog.lock().map_err(|_| std::io::Error::other("catalog lock poisoned"))?; Ok(catalog.describe(table).is_some()) } /// 一张表该开哪些索引 fn index_columns(&self, table: &str) -> Vec { let mut cols = vec!["id".to_string()]; let c = self.catalog.lock().unwrap(); if let Some(s) = c.describe(table) { for i in &s.indexes { if !cols.contains(&i.column) { cols.push(i.column.clone()); } } } cols } /// 拿表句柄(缓存复用) fn table_handle(&self, table: &str) -> std::io::Result>> { if let Some(h) = self.tables.lock().unwrap().get(table) { return Ok(h.clone()); } let columns = self.index_columns(table); let hidden = self.catalog.lock().map_err(|_| std::io::Error::other("catalog lock poisoned"))? .describe(table).map(|s| s.dropped_columns.iter().cloned().collect()).unwrap_or_default(); let paths: Vec<(String, PathBuf)> = columns .iter() .map(|c| (c.clone(), self.index_path(table, c))) .collect(); let mut map = self.tables.lock().unwrap(); if let Some(h) = map.get(table) { return Ok(h.clone()); } std::fs::create_dir_all(&self.cfg.data_dir)?; let mut t = HeapTable::open_with_indexes( self.table_path(table), &paths, self.cfg.page_size, self.cfg.btree_order, self.cfg.pool_size, )?; t.set_hidden_columns(hidden); let h = Arc::new(Mutex::new(t)); map.insert(table.to_string(), h.clone()); Ok(h) } /// 在表上跑闭包 fn with_table( &self, table: &str, f: impl FnOnce(&mut HeapTable) -> std::io::Result, ) -> std::io::Result { let h = self.table_handle(table)?; let mut guard = h.lock().unwrap(); f(&mut guard) } /// 表不存在就报错 fn with_existing_table( &self, table: &str, f: impl FnOnce(&mut HeapTable) -> std::io::Result, ) -> std::io::Result { if !self.table_exists(table)? { return Err(std::io::Error::new( std::io::ErrorKind::NotFound, format!("unknown table: {}", table), )); } self.with_table(table, f) } /// 把一条 WAL 操作落到磁盘 fn apply_op(&self, op: &WalOp) -> std::io::Result<()> { match op { WalOp::HideColumn { table, column } => { let mut c = self.catalog.lock().map_err(|_| std::io::Error::other("catalog lock poisoned"))?; c.remove_column(table, column)?; c.save(self.catalog_path())?; drop(c); if let Some(handle) = self.tables.lock().map_err(|_| std::io::Error::other("table lock poisoned"))?.get(table) { handle.lock().map_err(|_| std::io::Error::other("heap lock poisoned"))?.hide_column(column); } match std::fs::remove_file(self.index_path(table, column)) { Ok(()) => Ok(()), Err(e) if e.kind() == std::io::ErrorKind::NotFound => Ok(()), Err(e) => Err(e), } } WalOp::Compact { table, rows } => { self.with_table(table, |t| { t.replace_all(rows)?; t.flush() })?; let mut c = self.catalog.lock().unwrap(); c.mark_compacted(table)?; c.rebuild_stats(table, rows); c.save(self.catalog_path()) } WalOp::Accounts { accounts } => { let rows = account_rows(accounts)?; self.tables.lock().map_err(|_| std::io::Error::other("table lock poisoned"))?.remove(USERS_TABLE); for path in [self.table_path(USERS_TABLE), self.index_path(USERS_TABLE, "id")] { std::fs::File::create(path)?.sync_all()?; } self.with_table(USERS_TABLE, |t| { t.replace_all(&rows)?; t.flush() })?; let mut c = self.catalog.lock().map_err(|_| std::io::Error::other("catalog lock poisoned"))?; declare_system_schema(&mut c, &rows)?; c.save(self.catalog_path()) } WalOp::Insert { table, row } => { self.with_table(table, |t| t.insert_row(row))?; let mut c = self.catalog.lock().unwrap(); if c.record_insert(table, row) { c.save(self.catalog_path())?; } Ok(()) } WalOp::InsertBatch { table, rows } => { self.with_table(table, |t| t.insert_rows(rows))?; let mut c = self.catalog.lock().unwrap(); let mut changed = false; for r in rows { if c.record_insert(table, r) { changed = true; } } if changed { c.save(self.catalog_path())?; } Ok(()) } WalOp::DeleteKeys { table, keys } => { let deleted = self.with_table(table, |t| t.delete_by_keys(keys))?; let mut c = self.catalog.lock().unwrap(); c.record_delete(table, &deleted); Ok(()) } WalOp::ReplaceAll { table, rows } => { self.with_table(table, |t| t.replace_all(rows))?; let mut c = self.catalog.lock().unwrap(); c.rebuild_stats(table, rows); c.save(self.catalog_path()) } WalOp::DropColumn { table, column, rows } => { self.with_table(table, |t| { t.detach_index(column); t.replace_all(rows) })?; let mut c = self.catalog.lock().unwrap(); c.remove_column(table, column)?; c.rebuild_stats(table, rows); c.save(self.catalog_path())?; drop(c); let _ = std::fs::remove_file(self.index_path(table, column)); Ok(()) } WalOp::ReplaceSchema { table, columns, rows } => { let old_indexes: Vec = { let c = self.catalog.lock().map_err(|_| std::io::Error::other("catalog lock poisoned"))?; c.describe(table) .map(|s| s.indexes.iter().map(|i| i.column.clone()).collect()) .unwrap_or_default() }; self.with_table(table, |t| { for col in &old_indexes { t.detach_index(col); } t.set_hidden_columns(Default::default()); t.replace_all(rows) })?; let mut c = self.catalog.lock().unwrap(); c.set_columns(table, columns.clone())?; c.rebuild_stats(table, rows); c.save(self.catalog_path())?; drop(c); for col in old_indexes { let _ = std::fs::remove_file(self.index_path(table, &col)); } Ok(()) } } } /// 先写 WAL 再改数据 fn write_via_wal(&self, op: &WalOp) -> std::io::Result<()> { let _guard = self.write_lock.lock().unwrap(); if self.recovery_required.load(Ordering::Acquire) { return Err(std::io::Error::other("storage recovery required")); } { let c = self.catalog.lock().map_err(|_| std::io::Error::other("catalog lock poisoned"))?; let validate = |table: &str, rows: &[Row]| -> std::io::Result<()> { let Some(schema) = c.describe(table) else { return Ok(()); }; if let Some(name) = rows.iter().flat_map(|r| r.keys()).find(|k| schema.dropped_columns.contains(*k)) { return Err(std::io::Error::other(format!("unknown column: {name}"))); } Ok(()) }; match op { WalOp::Insert { table, row } => validate(table, std::slice::from_ref(row))?, WalOp::InsertBatch { table, rows } | WalOp::ReplaceAll { table, rows } => validate(table, rows)?, WalOp::HideColumn { table, column } if !c.describe(table).is_some_and(|s| s.columns.iter().any(|col| &col.name == column)) => { return Err(std::io::Error::other(format!("unknown column: {column}"))); } _ => {} } } std::fs::create_dir_all(&self.cfg.data_dir)?; let wal = &self.wal; wal.truncate()?; wal.append(op)?; let result = self.apply_op(op); if result.is_err() && matches!(op, WalOp::HideColumn { .. }) { self.recovery_required.store(true, Ordering::Release); return result; } if result.is_ok() { self.flush_op_tables(op)?; } wal.clear()?; result } /// 把涉及的表写回磁盘 fn flush_op_tables(&self, op: &WalOp) -> std::io::Result<()> { let table = match op { WalOp::HideColumn { .. } => return Ok(()), WalOp::Accounts { .. } => USERS_TABLE, WalOp::Insert { table, .. } | WalOp::InsertBatch { table, .. } | WalOp::DeleteKeys { table, .. } | WalOp::ReplaceAll { table, .. } | WalOp::DropColumn { table, .. } | WalOp::Compact { table, .. } | WalOp::ReplaceSchema { table, .. } => table, }; self.with_table(table, |t| t.flush()) } /// 启动时重放残留 WAL fn recover_wal(&self) -> std::io::Result<()> { std::fs::create_dir_all(&self.cfg.data_dir)?; let wal = &self.wal; match wal.read()? { None => Ok(()), Some(op) => { log_warn!(core, "replaying incomplete WAL entry"); self.apply_op(&op)?; self.flush_op_tables(&op)?; wal.clear()?; log_info!(core, "WAL replay complete"); Ok(()) } } } /// 启动时重建列统计 fn rebuild_stats(&self) -> std::io::Result<()> { let names: Vec = self.catalog.lock().unwrap().table_names(); let mut touched = false; for table in names { if !self.table_path(&table).exists() { continue; } let rows = self.with_table(&table, |t| t.scan())?; let mut c = self.catalog.lock().unwrap(); c.rebuild_stats(&table, &rows); touched = true; } if touched { let c = self.catalog.lock().unwrap(); c.save(self.catalog_path())?; } Ok(()) } /// 老数据迁移:类型名规范、缺列补 null fn migrate_data(&self) -> std::io::Result<()> { let names: Vec = self .catalog .lock() .map_err(|_| std::io::Error::other("catalog lock poisoned"))? .table_names(); for table in names { if !self.table_path(&table).exists() { continue; } let columns: Vec = match self .catalog .lock() .map_err(|_| std::io::Error::other("catalog lock poisoned"))? .describe(&table) { Some(schema) => schema.columns.clone(), None => continue, }; let rows = self.with_existing_table(&table, |t| t.scan())?; if rows.iter().all(|row| columns.iter().all(|c| row.contains_key(&c.name))) { continue; } let filled: Vec = rows .into_iter() .map(|mut row| { for column in &columns { row.entry(column.name.clone()).or_insert(serde_json::Value::Null); } row }) .collect(); log_info!(core, "migrating {}: {} rows", table, filled.len()); let op = WalOp::ReplaceAll { table, rows: filled }; self.write_via_wal(&op)?; } let mut catalog = self .catalog .lock() .map_err(|_| std::io::Error::other("catalog lock poisoned"))?; if catalog.normalize_types() { catalog.save(self.catalog_path())?; } Ok(()) } /// 给一列建索引 fn create_index(&self, table: &str, column: &str) -> std::io::Result<()> { if column == "id" { return Err(std::io::Error::new( std::io::ErrorKind::AlreadyExists, "id is the built-in index and cannot be created twice", )); } { let c = self.catalog.lock().unwrap(); let Some(schema) = c.describe(table) else { return Err(std::io::Error::new( std::io::ErrorKind::NotFound, format!("unknown table: {}", table), )); }; if !schema.columns.iter().any(|c| c.name == column) { return Err(std::io::Error::new( std::io::ErrorKind::NotFound, format!("unknown column: {}", column), )); } if schema.indexes.iter().any(|i| i.column == column) { return Err(std::io::Error::new( std::io::ErrorKind::AlreadyExists, format!("index on column \"{}\" already exists", column), )); } } let path = self.index_path(table, column); let (page_size, order, pool) = (self.cfg.page_size, self.cfg.btree_order, self.cfg.pool_size); self.with_table(table, |t| { t.build_index(column, &path, page_size, order, pool)?; t.flush() })?; let mut c = self.catalog.lock().unwrap(); c.add_index(table, column)?; c.save(self.catalog_path()) } /// 去掉一列的索引(连文件一起删) fn drop_index(&self, table: &str, column: &str) -> std::io::Result<()> { if column == "id" { return Err(std::io::Error::new( std::io::ErrorKind::InvalidInput, "id is the built-in index and cannot be dropped", )); } { let c = self.catalog.lock().unwrap(); if c.describe(table).is_none() { return Err(std::io::Error::new( std::io::ErrorKind::NotFound, format!("unknown table: {}", table), )); } } self.with_table(table, |t| { t.detach_index(column); Ok(()) })?; let mut c = self.catalog.lock().unwrap(); c.remove_index(table, column)?; c.save(self.catalog_path())?; drop(c); let _ = std::fs::remove_file(self.index_path(table, column)); Ok(()) } /// 删除列定义并隐藏旧字段,不改写堆页。 fn drop_column(&self, table: &str, column: &str) -> std::io::Result<()> { if column == "id" { return Err(std::io::Error::new( std::io::ErrorKind::InvalidInput, "the built-in id column cannot be dropped", )); } { let c = self.catalog.lock().unwrap(); let Some(schema) = c.describe(table) else { return Err(std::io::Error::new( std::io::ErrorKind::NotFound, format!("unknown table: {}", table), )); }; if !schema.columns.iter().any(|c| c.name == column) { return Err(std::io::Error::new( std::io::ErrorKind::NotFound, format!("unknown column: {}", column), )); } } let op = WalOp::HideColumn { table: table.to_string(), column: column.to_string(), }; self.write_via_wal(&op) } /// 整理一张表:按当前列定义重写堆页,回收隐藏列占的空间 fn compact_table(&self, table: &str) -> std::io::Result<()> { let rows = self.with_existing_table(table, |t| t.scan())?; let op = WalOp::Compact { table: table.to_string(), rows, }; self.write_via_wal(&op) } /// 启动时整理:把还留着隐藏列的表逐张重写 fn compact_tables(&self) -> std::io::Result<()> { let tables = { let c = self.catalog.lock().map_err(|_| std::io::Error::other("catalog lock poisoned"))?; c.dropped_tables() }; for table in tables { log_info!(core, "compacting table {}", table); self.compact_table(&table)?; } Ok(()) } /// 账号读写共享串行持久化边界 fn accounts_request(&self, request: Request) -> std::io::Result> { use std::io::{Error, ErrorKind}; let _guard = self.write_lock.lock().map_err(|_| Error::other("write lock poisoned"))?; if self.recovery_required.load(Ordering::Acquire) { return Err(Error::other("storage recovery required")); } let initialized = { let c = self.catalog.lock().map_err(|_| Error::other("catalog lock poisoned"))?; match c.describe(USERS_TABLE) { Some(schema) if schema.system => true, Some(_) => return Err(Error::new(ErrorKind::AlreadyExists, "reserved account table collision")), None => false, } }; if !initialized && (self.table_path(USERS_TABLE).exists() || self.index_path(USERS_TABLE, "id").exists()) { return Err(Error::new(ErrorKind::AlreadyExists, "reserved account file collision")); } let mut accounts: Vec = if initialized { if !self.table_path(USERS_TABLE).exists() { return Err(Error::other("missing account data")); } let mut loaded: Vec = self .with_existing_table(USERS_TABLE, |t| t.scan())? .iter() .map(decode_account_row) .collect::>()?; for account in loaded.iter_mut() { normalize_account(account); } loaded } else { Vec::new() }; match request { Request::AccountsList => return Ok(accounts), Request::AccountCreate { user, password_hash } => { let name = account_name(&user)?; if accounts.iter().any(|a| a.user == name) { return Err(Error::new(ErrorKind::AlreadyExists, "account already exists")); } let id = accounts.iter().map(|a| a.id).max().unwrap_or(0) .checked_add(1).ok_or_else(|| Error::other("account id exhausted"))?; accounts.push(new_account(id, &name, password_hash)?); } Request::AccountReset { user, password_hash } => { let name = account_name(&user)?; validate_account_hash(&password_hash)?; let account = accounts.iter_mut().find(|a| a.user == name) .ok_or_else(|| Error::new(ErrorKind::NotFound, "unknown account"))?; account.revision = account.revision.checked_add(1).ok_or_else(|| Error::other("account revision exhausted"))?; account.password_hash = password_hash; } Request::AccountLogin { user, at } => { let name = account_name(&user)?; let stamp = at.unwrap_or_else(now_stamp); let account = accounts.iter_mut().find(|a| a.user == name) .ok_or_else(|| Error::new(ErrorKind::NotFound, "unknown account"))?; account.last_login_at = Some(stamp); } Request::AccountDrop { user } => { let name = account_name(&user)?; let before = accounts.len(); accounts.retain(|a| a.user != name); if accounts.len() == before { return Err(Error::new(ErrorKind::NotFound, "unknown account")); } } _ => return Err(Error::other("unexpected account request")), } for row in account_rows(&accounts)? { HeapTable::encode_row(&row, self.cfg.page_size)?; } let op = WalOp::Accounts { accounts: accounts.clone() }; self.recovery_required.store(true, Ordering::Release); self.wal.truncate()?; self.wal.append(&op)?; self.apply_op(&op)?; self.wal.clear()?; self.recovery_required.store(false, Ordering::Release); Ok(accounts) } /// 一条请求翻成一条响应 fn handle_request(&self, req: Request) -> Response { let _ddl_guard = if matches!(req, Request::CreateTable { .. } | Request::DropTable { .. } | Request::CreateIndex { .. } | Request::DropIndex { .. }) { match self.write_lock.lock() { Ok(guard) => Some(guard), Err(_) => return Response::Error { message: "write lock poisoned".into() }, } } else { None }; if self.recovery_required.load(Ordering::Acquire) && !matches!(req, Request::Ping) { return Response::Error { message: "storage recovery required".into() }; } if !req.valid_columns() { return Response::Error { message: "invalid column name".into() }; } if let Some(table) = req.table() { if reserved_table(table) { return Response::Error { message: "reserved system table".into() }; } if table.is_empty() || !table.chars().all(|c| c.is_ascii_alphanumeric() || c == '_') { return Response::Error { message: "invalid table name".into() }; } } match req { request @ (Request::AccountsList | Request::AccountCreate { .. } | Request::AccountReset { .. } | Request::AccountLogin { .. } | Request::AccountDrop { .. }) => { match self.accounts_request(request) { Ok(accounts) => Response::Accounts { accounts }, Err(e) => Response::Error { message: e.to_string() }, } } Request::Ping => { log_debug!(request, "ping"); Response::Pong } Request::Scan { table, columns } => { log_debug!(request, "scan table={}", table); let result = (|| { if let Some(cols) = &columns { let c = self.catalog.lock().map_err(|_| std::io::Error::other("catalog lock poisoned"))?; let schema = c.describe(&table).ok_or_else(|| std::io::Error::other(format!("unknown table: {table}")))?; for col in cols { if !schema.columns.iter().any(|c| &c.name == col) { return Err(std::io::Error::other(format!("unknown column: {col}"))); } } } self.with_existing_table(&table, |t| t.scan_columns(columns.as_deref())) })(); match result { Ok(rows) => Response::Rows { rows }, Err(e) => { log_warn!(request, "scan {}: {}", table, e); Response::Error { message: format!("scan {}: {}", table, e), } } } } Request::Insert { table, row } => { log_debug!(request, "insert table={}", table); let tname = table.clone(); let op = WalOp::Insert { table, row }; match self.write_via_wal(&op) { Ok(()) => Response::Ok, Err(e) => { log_warn!(request, "insert {}: {}", tname, e); Response::Error { message: format!("insert {}: {}", tname, e), } } } } Request::InsertBatch { table, rows } => { log_debug!(request, "insert_batch table={} rows={}", table, rows.len()); let tname = table.clone(); let op = WalOp::InsertBatch { table, rows }; match self.write_via_wal(&op) { Ok(()) => Response::Ok, Err(e) => { log_warn!(request, "insert_batch {}: {}", tname, e); Response::Error { message: format!("insert_batch {}: {}", tname, e), } } } } Request::DeleteKeys { table, keys } => { log_debug!(request, "delete_keys table={} keys={}", table, keys.len()); let tname = table.clone(); let op = WalOp::DeleteKeys { table, keys }; match self.write_via_wal(&op) { Ok(()) => Response::Ok, Err(e) => { log_warn!(request, "delete_keys {}: {}", tname, e); Response::Error { message: format!("delete_keys {}: {}", tname, e), } } } } Request::LookupByIndex { table, column, key } => { log_debug!(request, "lookup_by_index table={} column={} key={}", table, column, key); let indexed = match self.with_existing_table(&table, |t| Ok(t.has_index(&column))) { Ok(v) => v, Err(e) => { log_warn!(request, "lookup_by_index {}: {}", table, e); return Response::Error { message: format!("lookup_by_index {}: {}", table, e), }; } }; if !indexed { return Response::NoIndex; } let found = match key.as_i64() { Some(k) => self.with_table(&table, |t| t.get_all_by_column_key(&column, k)), None => match key.as_str() { Some(s) => self.with_table(&table, |t| t.get_all_by_string_key(&column, s)), None => { return Response::Error { message: "lookup_by_index: key must be a number or a string".into(), }; } }, }; match found { Ok(rows) => Response::Rows { rows }, Err(e) => { log_warn!(request, "lookup_by_index {}: {}", table, e); Response::Error { message: format!("lookup_by_index {}: {}", table, e), } } } } Request::RangeByIndex { table, column, lo, lo_inclusive, hi, hi_inclusive } => { log_debug!(request, "range_by_index table={} column={}", table, column); let indexed = match self.with_existing_table(&table, |t| Ok(t.has_index(&column))) { Ok(v) => v, Err(e) => { log_warn!(request, "range_by_index {}: {}", table, e); return Response::Error { message: format!("range_by_index {}: {}", table, e), }; } }; if !indexed { return Response::NoIndex; } let lo = lo.as_ref().map(|v| (v, lo_inclusive)); let hi = hi.as_ref().map(|v| (v, hi_inclusive)); match self.with_table(&table, |t| t.scan_range(&column, lo, hi)) { Ok(Some(rows)) => Response::Rows { rows }, Ok(None) => Response::NoIndex, Err(e) => { log_warn!(request, "range_by_index {}: {}", table, e); Response::Error { message: format!("range_by_index {}: {}", table, e), } } } } Request::ReplaceAll { table, rows } => { log_debug!(request, "replace_all table={} rows={}", table, rows.len()); let tname = table.clone(); let op = WalOp::ReplaceAll { table, rows }; match self.write_via_wal(&op) { Ok(()) => Response::Ok, Err(e) => { log_warn!(request, "replace_all {}: {}", tname, e); Response::Error { message: format!("replace_all {}: {}", tname, e), } } } } Request::ListTables => { log_debug!(request, "list_tables"); match self.list_tables() { Ok(tables) => Response::Tables { tables }, Err(e) => { log_warn!(request, "list_tables: {}", e); Response::Error { message: format!("list_tables: {}", e), } } } } Request::ListCatalog => { log_debug!(request, "list_catalog"); let c = self.catalog.lock().unwrap(); let schemas: Vec = c .all_tables() .into_iter() .filter(|(name, _)| !reserved_table(name)) .map(|(name, ts)| TableSchemaWire { table: name.to_string(), columns: ts.columns.clone(), row_count: ts.row_count, indexes: index_wires(ts), stats: stat_wires(ts), }) .collect(); Response::Catalog { schemas } } Request::DescribeTable { table } => { log_debug!(request, "describe_table table={}", table); match self.describe_table(&table) { Ok((cols, count, indexes, stats)) => Response::Schema { table, columns: cols, row_count: count, indexes, stats, }, Err(e) => { log_warn!(request, "describe_table {}: {}", table, e); Response::Error { message: format!("describe_table {}: {}", table, e), } } } } Request::CreateTable { table, columns } => { log_debug!(request, "create_table table={} cols={}", table, columns.len()); { let c = self.catalog.lock().unwrap(); if c.describe(&table).is_some() { return Response::Error { message: format!("create_table {}: table already exists", table), }; } } if let Err(e) = self.with_table(&table, |_t| Ok(())) { return Response::Error { message: format!("create_table {}: {}", table, e), }; } let mut c = self.catalog.lock().unwrap(); if let Err(e) = c.create_table(&table, columns) { return Response::Error { message: format!("create_table {}: {}", table, e), }; } if let Err(e) = c.save(self.catalog_path()) { return Response::Error { message: format!("create_table {}: catalog save: {}", table, e), }; } Response::Ok } Request::DropTable { table } => { log_debug!(request, "drop_table table={}", table); let dropped_indexes: Vec = { let mut c = self.catalog.lock().unwrap(); let idx = c .describe(&table) .map(|s| s.indexes.iter().map(|i| i.column.clone()).collect()) .unwrap_or_default(); if let Err(e) = c.drop_table(&table) { return Response::Error { message: format!("drop_table {}: {}", table, e), }; } if let Err(e) = c.save(self.catalog_path()) { return Response::Error { message: format!("drop_table {}: catalog save: {}", table, e), }; } idx }; self.tables.lock().unwrap().remove(&table); let _ = std::fs::remove_file(self.table_path(&table)); let _ = std::fs::remove_file(self.index_path(&table, "id")); for col in dropped_indexes { let _ = std::fs::remove_file(self.index_path(&table, &col)); } Response::Ok } Request::CreateIndex { table, column } => { log_debug!(request, "create_index table={} column={}", table, column); match self.create_index(&table, &column) { Ok(()) => Response::Ok, Err(e) => { log_warn!(request, "create_index {}({}): {}", table, column, e); Response::Error { message: format!("create_index {}({}): {}", table, column, e), } } } } Request::DropIndex { table, column } => { log_debug!(request, "drop_index table={} column={}", table, column); match self.drop_index(&table, &column) { Ok(()) => Response::Ok, Err(e) => { log_warn!(request, "drop_index {}({}): {}", table, column, e); Response::Error { message: format!("drop_index {}({}): {}", table, column, e), } } } } Request::DropColumn { table, column } => { log_debug!(request, "drop_column table={} column={}", table, column); match self.drop_column(&table, &column) { Ok(()) => Response::Ok, Err(e) => { log_warn!(request, "drop_column {}({}): {}", table, column, e); Response::Error { message: format!("drop_column {}({}): {}", table, column, e), } } } } Request::Compact { table } => { log_debug!(request, "compact table={}", table); match self.compact_table(&table) { Ok(()) => Response::Ok, Err(e) => { log_warn!(request, "compact {}: {}", table, e); Response::Error { message: format!("compact {}: {}", table, e), } } } } Request::ReplaceSchema { table, columns, rows } => { log_debug!(request, "replace_schema table={} cols={}", table, columns.len()); let tname = table.clone(); let op = WalOp::ReplaceSchema { table, columns, rows }; match self.write_via_wal(&op) { Ok(()) => Response::Ok, Err(e) => { log_warn!(request, "replace_schema {}: {}", tname, e); Response::Error { message: format!("replace_schema {}: {}", tname, e), } } } } } } /// 只列出数据字典中的业务表 fn list_tables(&self) -> std::io::Result> { let c = self.catalog.lock().map_err(|_| std::io::Error::other("catalog lock poisoned"))?; Ok(c.table_names().into_iter().filter(|t| !reserved_table(t)).collect()) } /// 查一张表的 schema fn describe_table( &self, table: &str, ) -> std::io::Result { let c = self.catalog.lock().unwrap(); if let Some(schema) = c.describe(table) { return Ok(( schema.columns.clone(), schema.row_count, index_wires(schema), stat_wires(schema), )); } drop(c); Err(std::io::Error::new( std::io::ErrorKind::NotFound, format!("unknown table: {}", table), )) } } /// 账号列表转成行 fn account_rows(accounts: &[Account]) -> std::io::Result> { accounts.iter().map(|a| Ok(HashMap::from([ ("id".into(), serde_json::json!(a.id)), ("user".into(), serde_json::Value::String(a.user.clone())), ("password_hash".into(), serde_json::Value::String(a.password_hash.clone())), ("registered_at".into(), serde_json::Value::String(a.registered_at.clone())), ("last_login_at".into(), match &a.last_login_at { Some(stamp) => serde_json::Value::String(stamp.clone()), None => serde_json::Value::Null, }), ("revision".into(), serde_json::json!(a.revision)), ]))).collect() } /// 按 init.sql 把系统表写进字典 fn declare_system_schema(catalog: &mut Catalog, rows: &[Row]) -> std::io::Result<()> { for table in system_schema() { let stats: &[Row] = if table.name == USERS_TABLE { rows } else { &[] }; catalog.rebuild_stats(&table.name, stats); catalog.set_columns(&table.name, table.columns.iter().map(InitColumn::to_schema).collect())?; catalog.rebuild_stats(&table.name, stats); catalog.mark_system(&table.name)?; for index in &table.indexes { catalog.add_index(&table.name, index)?; } } Ok(()) } /// 读一行账号(兼容老格式) fn decode_account_row(row: &Row) -> std::io::Result { if let Some(encoded) = row.get("account").and_then(serde_json::Value::as_str) { return serde_json::from_str(encoded).map_err(|_| std::io::Error::other("invalid account record")); } let id = row.get("id").and_then(serde_json::Value::as_i64) .ok_or_else(|| std::io::Error::other("invalid account record"))?; let user = row.get("user").and_then(serde_json::Value::as_str) .ok_or_else(|| std::io::Error::other("invalid account record"))?; let password_hash = row.get("password_hash").and_then(serde_json::Value::as_str) .ok_or_else(|| std::io::Error::other("invalid account record"))?; Ok(Account { id, user: user.to_string(), password_hash: password_hash.to_string(), revision: row.get("revision").and_then(serde_json::Value::as_u64).unwrap_or(1), registered_at: row.get("registered_at").and_then(serde_json::Value::as_str).unwrap_or("").to_string(), last_login_at: row.get("last_login_at").and_then(serde_json::Value::as_str).map(str::to_string), }) } /// 老记录没有注册时间就补当前时间 fn normalize_account(account: &mut Account) { if account.registered_at.is_empty() { account.registered_at = now_stamp(); } } /// 当前 UTC 时间,写法与引擎的 timestamp 一致 fn now_stamp() -> String { let seconds = std::time::SystemTime::now() .duration_since(std::time::UNIX_EPOCH) .map(|d| d.as_secs() as i64) .unwrap_or(0); format_timestamp(seconds) } /// Unix 秒转时间戳文本 fn format_timestamp(seconds: i64) -> String { let days = seconds.div_euclid(86_400); let rest = seconds.rem_euclid(86_400); let (year, month, day) = civil_from_days(days); format!( "{:04}-{:02}-{:02} {:02}:{:02}:{:02}", year, month, day, rest / 3600, (rest % 3600) / 60, rest % 60 ) } /// 天数 → (年, 月, 日) fn civil_from_days(days: i64) -> (i64, u32, u32) { let z = days + 719_468; let era = if z >= 0 { z } else { z - 146_096 } / 146_097; let doe = z - era * 146_097; let yoe = (doe - doe / 1460 + doe / 36_524 - doe / 146_096) / 365; let year = yoe + era * 400; let doy = doe - (365 * yoe + yoe / 4 - yoe / 100); let mp = (5 * doy + 2) / 153; let day = (doy - (153 * mp + 2) / 5 + 1) as u32; let month = if mp < 10 { mp + 3 } else { mp - 9 } as u32; (if month <= 2 { year + 1 } else { year }, month, day) } /// 规范化并校验账号名 fn account_name(user: &str) -> std::io::Result { if user.is_empty() || user.len() > 64 || !user.bytes().all(|c| c.is_ascii_alphanumeric() || b"_.-".contains(&c)) { return Err(std::io::Error::other("invalid account name")); } Ok(user.to_ascii_lowercase()) } /// 校验口令哈希长度 fn validate_account_hash(hash: &str) -> std::io::Result<()> { if hash.is_empty() || hash.len() > 256 { return Err(std::io::Error::other("invalid account credential")); } Ok(()) } /// 造一条新账号记录 fn new_account(id: i64, user: &str, password_hash: String) -> std::io::Result { validate_account_hash(&password_hash)?; Ok(Account { id, user: account_name(user)?, password_hash, revision: 1, registered_at: now_stamp(), last_login_at: None, }) } /// 一张表的索引定义转线上格式 fn index_wires(schema: &crate::catalog::TableSchema) -> Vec { let mut out = vec![IndexWire { column: "id".to_string(), }]; for i in &schema.indexes { if i.column != "id" { out.push(IndexWire { column: i.column.clone(), }); } } out } /// 一张表的列统计 -> 线上格式 fn stat_wires(schema: &crate::catalog::TableSchema) -> Vec { schema .stats .iter() .map(|(name, s)| ColumnStatWire { name: name.clone(), distinct: s.distinct, capped: s.capped, }) .collect() } impl Drop for Server { /// 析构时保存数据字典 fn drop(&mut self) { if let Ok(c) = self.catalog.lock() { let _ = c.save(self.catalog_path()); } } } /// 系统库名:账号表住在这里,其中的表都是系统表。服务启动时只建它一个, /// 别的库一律要显式建(没有默认工作库)。 const SYSTEM_DATABASE: &str = "system"; /// 系统库的表结构来自这份脚本(编译期嵌进来) const SYSTEM_SCHEMA_SQL: &str = include_str!("../../../scripts/init.sql"); /// 系统库 schema:脚本缺表或多表都按它来 fn system_schema() -> &'static [InitTable] { static TABLES: std::sync::OnceLock> = std::sync::OnceLock::new(); TABLES.get_or_init(|| initsql::parse(SYSTEM_SCHEMA_SQL).unwrap_or_else(|e| panic!("scripts/init.sql: {e}"))) } /// 保留库:不能建、不能删。 fn reserved_database(name: &str) -> bool { name == SYSTEM_DATABASE } /// 每个库拥有独立的目录、字典、句柄和 WAL。 struct Databases { root: Config, system: Server, named: Mutex>, } impl Databases { /// 装载系统库与已有的具名库 fn new(cfg: Config) -> std::io::Result { // 系统库按需创建:目录、字典与系统表缺失时自动补齐。 // 其余库只在 data/databases// 下已存在时才装载,启动不会凭空建库。 let system = Self::open(&cfg, SYSTEM_DATABASE)?; let root = cfg.data_dir.join("databases"); std::fs::create_dir_all(&root)?; let canonical_root = std::fs::canonicalize(&root)?; if canonical_root.parent() != Some(std::fs::canonicalize(&cfg.data_dir)?.as_path()) { return Err(std::io::Error::other("database root escapes data directory")); } let mut named = HashMap::new(); for entry in std::fs::read_dir(root)? { let entry = entry?; let name = entry.file_name().to_string_lossy().to_string(); if entry.file_type()?.is_dir() && valid_database(&name) && !reserved_database(&name) && name == name.to_ascii_lowercase() && entry.path().join("catalog.json").is_file() { if std::fs::canonicalize(entry.path())?.parent() != Some(canonical_root.as_path()) { return Err(std::io::Error::other("database directory escapes database root")); } let mut config = cfg.clone(); config.data_dir = entry.path(); let server = Server::new(config)?; server.recover_wal()?; server.rebuild_stats()?; server.migrate_data()?; server.compact_tables()?; named.insert(name, server); } } Ok(Self { root: cfg, system, named: Mutex::new(named) }) } /// 按名字打开或新建一个库 fn open(cfg: &Config, name: &str) -> std::io::Result { let mut config = cfg.clone(); config.data_dir = cfg.data_dir.join(name); let server = Server::new(config)?; if !server.catalog_path().is_file() { server.catalog.lock().map_err(|_| std::io::Error::other("catalog lock poisoned"))?.save(server.catalog_path())?; } server.recover_wal()?; server.rebuild_stats()?; server.migrate_data()?; server.compact_tables()?; Ok(server) } /// 库级分发:先处理库管理,再转给选中的库 fn request(&self, mut value: serde_json::Value) -> std::io::Result { use std::io::Error; let method = value.get("method").and_then(|v| v.as_str()).unwrap_or("").to_string(); let name = match value.get("database") { None => String::new(), Some(serde_json::Value::String(s)) => s.to_ascii_lowercase(), _ => return Err(Error::other("invalid database name")), }; if !name.is_empty() && !valid_database(&name) { return Err(Error::other("invalid database name")); } let mut named = self.named.lock().map_err(|_| Error::other("database lock poisoned"))?; match method.as_str() { "all_catalogs" => { let mut schemas = Vec::new(); for (database, server) in std::iter::once((SYSTEM_DATABASE, &self.system)) .chain(named.iter().map(|(name, server)| (name.as_str(), server))) { match server.handle_request(Request::ListCatalog) { Response::Catalog { schemas: entries } => { for mut entry in entries { entry.table = format!("{database}.{}", entry.table); schemas.push(entry); } } other => return Ok(other), } } return Ok(Response::Catalog { schemas }); } "list_databases" => { let mut databases: Vec = named.keys().cloned().collect(); databases.sort(); databases.push(SYSTEM_DATABASE.into()); return Ok(Response::Tables { tables: databases }); } "create_database" => { if !valid_database(&name) { return Err(Error::other("invalid database name")); } if reserved_database(&name) || named.contains_key(&name) { return Err(Error::other("database already exists")); } let mut cfg = self.root.clone(); cfg.data_dir = cfg.data_dir.join("databases").join(&name); std::fs::create_dir(&cfg.data_dir)?; let server = Server::new(cfg)?; server.catalog.lock().map_err(|_| Error::other("catalog lock poisoned"))?.save(server.catalog_path())?; named.insert(name, server); return Ok(Response::Ok); } "drop_database" => { if name == SYSTEM_DATABASE { return Err(Error::other("cannot drop the system database")); } let server = named.get(&name).ok_or_else(|| Error::other("unknown database"))?; let root = std::fs::canonicalize(self.root.data_dir.join("databases"))?; let target = std::fs::canonicalize(&server.cfg.data_dir)?; if target.parent() != Some(root.as_path()) || !target.join("catalog.json").is_file() { return Err(Error::other("invalid database directory")); } let config = server.cfg.clone(); drop(named.remove(&name)); let stamp = std::time::SystemTime::now().duration_since(std::time::UNIX_EPOCH).map_err(Error::other)?.as_nanos(); let tombstone = root.join(format!(".dropped-{name}-{stamp}")); if let Err(error) = std::fs::rename(&target, &tombstone) { named.insert(name, Server::new(config)?); return Err(error); } std::fs::remove_dir_all(&tombstone).map_err(|e| Error::other(format!("database dropped; directory cleanup failed: {e}")))?; return Ok(Response::Ok); } "use_database" => { return if reserved_database(&name) || named.contains_key(&name) { Ok(Response::Ok) } else { Err(Error::other("unknown database")) }; } _ => {} } let mut selected = name; let qualified = value.get("table").and_then(|v| v.as_str()).map(str::to_owned); if let Some((database, table)) = qualified.as_deref().and_then(|t| t.split_once('.')) { if !valid_database(database) { return Err(Error::other("invalid database name")); } selected = database.to_ascii_lowercase(); value["table"] = serde_json::Value::String(table.to_string()); } let request: Request = serde_json::from_value(value).map_err(Error::other)?; let global = matches!(request, Request::AccountsList | Request::AccountCreate { .. } | Request::AccountReset { .. } | Request::AccountLogin { .. } | Request::AccountDrop { .. } | Request::Ping); let server = if global { &self.system } else if selected.is_empty() { return Err(Error::other("no database selected")); } else if selected == SYSTEM_DATABASE { &self.system } else { named.get(&selected).ok_or_else(|| Error::other("unknown database"))? }; Ok(server.handle_request(request)) } } /// 校验库名是否合法 fn valid_database(name: &str) -> bool { let lower = name.to_ascii_lowercase(); let device = matches!(lower.as_str(), "con" | "prn" | "aux" | "nul") || (lower.len() == 4 && (lower.starts_with("com") || lower.starts_with("lpt")) && lower.as_bytes()[3].is_ascii_digit()); !name.is_empty() && name.len() <= 64 && !device && !reserved_table(name) && name.bytes().all(|c| c.is_ascii_alphanumeric() || c == b'_') && name.as_bytes().first().is_some_and(|c| c.is_ascii_alphabetic() || *c == b'_') } /// 把配置来源打进日志 fn log_config(loaded: &Loaded) { match &loaded.config_path { Some(p) => log_info!(core, "config file: {}", p.display()), None => log_info!(core, "config file: none"), } for (name, value, origin) in &loaded.origins { log_info!(core, " {:<18} {} [{}]", name, value, origin.describe()); } } /// 进程内存储句柄:一个数据目录一份状态,跨线程共享(内部各自加锁)。 pub struct Storage { databases: Databases, } impl Storage { /// 打开数据目录并初始化存储 pub fn open(config_path: Option<&str>) -> Result { let loaded = config::load(config_path)?; log::init(); log::set_level(loaded.config.log_level); log_info!(core, "chusql-core-storage {} starting", env!("CARGO_PKG_VERSION")); log_config(&loaded); let databases = Databases::new(loaded.config).map_err(|e| e.to_string())?; Ok(Storage { databases }) } /// 一行请求 JSON 换一行响应 JSON pub fn request_line(&self, line: &str) -> String { let response = match serde_json::from_str::(line) { Ok(value) => match self.databases.request(value) { Ok(response) => response, Err(e) => Response::Error { message: e.to_string() }, }, Err(e) => { log_warn!(core, "bad request line: {}", e); Response::Error { message: format!("parse error: {}", e), } } }; match serde_json::to_string(&response) { Ok(text) => text, Err(e) => { log_error!(core, "response encode failed: {}", e); "{\"status\":\"error\",\"message\":\"response encode failed\"}".to_string() } } } }