dataforge/db/

mod.rs

//! 数据库相关模块
//! 
//! 包含数据库填充、Schema解析、类型映射等功能

#[cfg(feature = "database")]
use sqlx::{Database, Pool};
use serde_json::Value;
use crate::error::Result;
use crate::generation::DataForge;

// 重新导出schema模块的类型
pub mod schema;
pub use schema::{TableSchema, FieldSchema, DataType, SchemaParser};

/// 数据库适配器 Trait
pub trait DatabaseAdapter {
    /// 自动推断表结构
    fn infer_schema(&self) -> Result<TableSchema>;
    
    /// 批量数据插入（支持事务回滚）
    fn bulk_insert(&self, table: &str, data: Vec<Value>) -> Result<usize>;
    
    /// 类型转换系统
    fn type_mapping(&self, rust_type: &str) -> String;
}

/// 现代化的数据库填充器（支持泛型数据库类型）
#[cfg(feature = "database")]
pub struct GenericDatabaseForge<DB: Database> {
    #[allow(dead_code)]
    pool: Pool<DB>,
    tables: Vec<TableConfig>,
    schema_parser: SchemaParser,
}

/// 简化的数据库填充器（不依赖sqlx）
pub struct SimpleDatabaseForge {
    #[allow(dead_code)]
    forge: DataForge,
    connection_string: String,
    batch_size: usize,
    schema_parser: SchemaParser,
}

/// 表配置
pub struct TableConfig {
    pub name: String,
    pub count: usize,
    pub fields: Vec<FieldConfig>,
}

/// 字段配置
pub struct FieldConfig {
    pub name: String,
    pub generator: Box<dyn Fn() -> Value + Send + Sync>,
}

#[cfg(feature = "database")]
impl<DB: Database> GenericDatabaseForge<DB> {
    /// 创建新的数据库填充器
    pub fn new(pool: Pool<DB>) -> Self {
        Self {
            pool,
            tables: Vec::new(),
            schema_parser: SchemaParser::new(),
        }
    }

    /// 添加表配置
    pub fn table<F>(mut self, name: &str, count: usize, config: F) -> Self
    where
        F: FnOnce(&mut TableBuilder) -> &mut TableBuilder,
    {
        let mut builder = TableBuilder::new(name.to_string());
        config(&mut builder);
        self.tables.push(builder.build(count));
        self
    }

    /// 执行数据填充
    pub async fn fill(&self) -> Result<()> {
        for table in &self.tables {
            self.fill_table(table).await?;
        }
        Ok(())
    }

    /// 填充单个表
    async fn fill_table(&self, _table: &TableConfig) -> Result<()> {
        // 这里需要根据具体的数据库类型实现。由于这是一个通用实现，我们只提供接口
        todo!("Database-specific implementation needed")
    }

    /// 获取Schema解析器
    pub fn schema_parser(&self) -> &SchemaParser {
        &self.schema_parser
    }

    /// 获取可变的Schema解析器
    pub fn schema_parser_mut(&mut self) -> &mut SchemaParser {
        &mut self.schema_parser
    }
}

impl SimpleDatabaseForge {
    /// 创建新的简化数据库填充器
    pub fn new(connection_string: &str) -> Self {
        Self {
            forge: DataForge::default(),
            connection_string: connection_string.to_string(),
            batch_size: 1000,
            schema_parser: SchemaParser::new(),
        }
    }

    /// 设置批次大小
    pub fn batch_size(mut self, size: usize) -> Self {
        self.batch_size = size;
        self
    }

    /// 配置表填充
    pub fn table<F>(self, table_name: &str, count: usize, config: F) -> TableFillConfig
    where
        F: FnOnce(&mut TableBuilder) -> &mut TableBuilder,
    {
        let mut builder = TableBuilder::new(table_name.to_string());
        config(&mut builder);
        
        TableFillConfig {
            forge: self,
            table_name: table_name.to_string(),
            count,
            fields: builder.build_simple(),
        }
    }

    /// 获取Schema解析器
    pub fn schema_parser(&self) -> &SchemaParser {
        &self.schema_parser
    }

    /// 获取可变的Schema解析器
    pub fn schema_parser_mut(&mut self) -> &mut SchemaParser {
        &mut self.schema_parser
    }

    /// 获取连接字符串
    pub fn connection_string(&self) -> &str {
        &self.connection_string
    }

    /// 获取当前批次大小
    pub fn get_batch_size(&self) -> usize {
        self.batch_size
    }
}

/// 表填充配置
pub struct TableFillConfig {
    #[allow(dead_code)]
    forge: SimpleDatabaseForge,
    table_name: String,
    count: usize,
    fields: Vec<(String, Box<dyn Fn() -> Value + Send + Sync>)>,
}

impl TableFillConfig {
    /// 执行填充
    #[cfg(feature = "database")]
    pub async fn fill(self) -> Result<usize> {
        // 这里应该实现实际的数据库连接和插入逻辑
        // 由于这是示例，我们只返回模拟结果
        println!("模拟填充表 {} 共 {} 条记录", self.table_name, self.count);
        
        // 生成数据
        for i in 0..self.count.min(5) { // 只打印前5条作为示例
            let mut record = serde_json::Map::new();
            for (field_name, generator) in &self.fields {
                record.insert(field_name.clone(), generator());
            }
            println!("记录 {}: {}", i + 1, serde_json::to_string_pretty(&record)?);
        }
        
        Ok(self.count)
    }

    /// 同步版本的填充
    pub fn fill_sync(self) -> Result<usize> {
        println!("填充表 {} 共 {} 条记录", self.table_name, self.count);
        
        // 生成数据
        for i in 0..self.count.min(5) { // 只打印前5条作为示例
            let mut record = serde_json::Map::new();
            for (field_name, generator) in &self.fields {
                record.insert(field_name.clone(), generator());
            }
            println!("记录 {}: {}", i + 1, serde_json::to_string_pretty(&record)?);
        }
        
        Ok(self.count)
    }

    /// 获取表名
    pub fn table_name(&self) -> &str {
        &self.table_name
    }

    /// 获取记录数量
    pub fn count(&self) -> usize {
        self.count
    }

    /// 获取字段配置
    pub fn fields(&self) -> &Vec<(String, Box<dyn Fn() -> Value + Send + Sync>)> {
        &self.fields
    }
}

/// 表构建器
pub struct TableBuilder {
    name: String,
    fields: Vec<FieldConfig>,
}

impl TableBuilder {
    fn new(name: String) -> Self {
        Self {
            name,
            fields: Vec::new(),
        }
    }

    /// 添加字段
    pub fn field<F>(&mut self, name: &str, generator: F) -> &mut Self
    where
        F: Fn() -> Value + Send + Sync + 'static,
    {
        self.fields.push(FieldConfig {
            name: name.to_string(),
            generator: Box::new(generator),
        });
        self
    }

    /// 构建表配置
    fn build(self, count: usize) -> TableConfig {
        TableConfig {
            name: self.name,
            count,
            fields: self.fields,
        }
    }

    /// 构建简单表配置（用于SimpleDatabaseForge）
    fn build_simple(self) -> Vec<(String, Box<dyn Fn() -> Value + Send + Sync>)> {
        self.fields.into_iter().map(|field| (field.name, field.generator)).collect()
    }
}

// 类型别名，保持向后兼容
// 默认的DatabaseForge指向简化版本
pub type DatabaseForge = SimpleDatabaseForge;

/// MySQL 数据库填充器（泛型版本）
#[cfg(feature = "database")]
pub type MySqlForge = GenericDatabaseForge<sqlx::MySql>;

/// PostgreSQL 数据库填充器（泛型版本）
#[cfg(feature = "database")]
pub type PostgresForge = GenericDatabaseForge<sqlx::Postgres>;

/// SQLite 数据库填充器（泛型版本）
#[cfg(feature = "database")]
pub type SqliteForge = GenericDatabaseForge<sqlx::Sqlite>;

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_simple_database_forge() {
        let forge = SimpleDatabaseForge::new("sqlite://test.db");
        assert_eq!(forge.connection_string, "sqlite://test.db");
        assert_eq!(forge.batch_size, 1000);
    }

    #[test]
    fn test_table_builder() {
        let mut builder = TableBuilder::new("users".to_string());
        builder
            .field("id", || Value::String("test".to_string()))
            .field("name", || Value::String("test_name".to_string()));

        let table = builder.build(100);
        assert_eq!(table.name, "users");
        assert_eq!(table.count, 100);
        assert_eq!(table.fields.len(), 2);
    }

    #[test]
    fn test_schema_parser_integration() {
        let forge = SimpleDatabaseForge::new("test://db");
        let parser = forge.schema_parser();
        
        // 测试Schema解析器是否正确初始化
        // 通过尝试推断一个简单的类型来验证解析器工作正常
        let json_schema = serde_json::json!({
            "title": "Test",
            "properties": {
                "id": {"type": "integer"}
            }
        });
        
        let result = parser.infer_from_json(&json_schema);
        assert!(result.is_ok());
    }
}