sqlcx_core/generator/python/
mod.rs

1pub mod asyncpg;
2pub mod common;
3pub mod mysql_connector;
4pub mod psycopg;
5pub mod pydantic;
6pub mod sqlite3_driver;
7
8use crate::config::TargetConfig;
9use crate::error::{Result, SqlcxError};
10use crate::generator::{DriverGenerator, GeneratedFile, LanguagePlugin, SchemaGenerator};
11use crate::ir::SqlcxIR;
12
13use self::pydantic::PydanticGenerator;
14
15pub struct PythonPlugin {
16    pub schema_name: String,
17    pub driver_name: String,
18}
19
20impl PythonPlugin {
21    pub fn new(schema: &str, driver: &str) -> Result<Self> {
22        resolve_schema(schema)?;
23        resolve_driver(driver)?;
24        Ok(Self {
25            schema_name: schema.to_string(),
26            driver_name: driver.to_string(),
27        })
28    }
29}
30
31fn resolve_schema(name: &str) -> Result<Box<dyn SchemaGenerator>> {
32    match name {
33        "pydantic" => Ok(Box::new(PydanticGenerator)),
34        _ => Err(SqlcxError::UnknownSchema(name.to_string())),
35    }
36}
37
38fn resolve_driver(name: &str) -> Result<Option<Box<dyn DriverGenerator>>> {
39    match name {
40        "none" => Ok(None),
41        "psycopg" => Ok(Some(Box::new(psycopg::PsycopgGenerator))),
42        "asyncpg" => Ok(Some(Box::new(asyncpg::AsyncpgGenerator))),
43        "sqlite3" => Ok(Some(Box::new(sqlite3_driver::Sqlite3Generator))),
44        "mysql-connector" => Ok(Some(Box::new(mysql_connector::MysqlConnectorGenerator))),
45        _ => Err(SqlcxError::UnknownDriver(name.to_string())),
46    }
47}
48
49impl LanguagePlugin for PythonPlugin {
50    fn generate(&self, ir: &SqlcxIR, config: &TargetConfig) -> Result<Vec<GeneratedFile>> {
51        let schema_gen = resolve_schema(&self.schema_name)?;
52        let overrides = &config.overrides;
53
54        let mut files = vec![schema_gen.generate(ir, overrides)?];
55
56        if let Some(driver_gen) = resolve_driver(&self.driver_name)? {
57            files.extend(driver_gen.generate(ir)?);
58        }
59
60        Ok(files)
61    }
62}
63
64#[cfg(test)]
65mod tests {
66    use super::*;
67    use crate::generator::LanguagePlugin;
68    use crate::parser::DatabaseParser;
69    use crate::parser::postgres::PostgresParser;
70    use std::collections::HashMap;
71
72    fn parse_fixture_ir() -> SqlcxIR {
73        let schema_sql = include_str!("../../../../../tests/fixtures/schema.sql");
74        let queries_sql = include_str!("../../../../../tests/fixtures/queries/users.sql");
75        let parser = PostgresParser::new();
76        let (tables, enums) = parser.parse_schema(schema_sql).unwrap();
77        let queries = parser
78            .parse_queries(queries_sql, &tables, &enums, "queries/users.sql")
79            .unwrap();
80        SqlcxIR {
81            tables,
82            queries,
83            enums,
84        }
85    }
86
87    #[test]
88    fn generates_one_file_with_no_driver() {
89        let ir = parse_fixture_ir();
90        let plugin = PythonPlugin::new("pydantic", "none").unwrap();
91        let config = TargetConfig {
92            language: "python".to_string(),
93            out: "./src/db".to_string(),
94            schema: "pydantic".to_string(),
95            driver: "none".to_string(),
96            overrides: HashMap::new(),
97        };
98        let files = plugin.generate(&ir, &config).unwrap();
99        assert_eq!(files.len(), 1);
100        assert!(files.iter().any(|f| f.path == "models.py"));
101    }
102
103    #[test]
104    fn generates_models_and_queries_with_psycopg() {
105        let ir = parse_fixture_ir();
106        let plugin = PythonPlugin::new("pydantic", "psycopg").unwrap();
107        let config = TargetConfig {
108            language: "python".to_string(),
109            out: "./src/db".to_string(),
110            schema: "pydantic".to_string(),
111            driver: "psycopg".to_string(),
112            overrides: HashMap::new(),
113        };
114        let files = plugin.generate(&ir, &config).unwrap();
115        assert_eq!(files.len(), 2);
116        assert!(files.iter().any(|f| f.path == "models.py"));
117        assert!(files.iter().any(|f| f.path.ends_with("_queries.py")));
118        assert!(!files.iter().any(|f| f.path == "client.py"));
119    }
120
121    #[test]
122    fn generates_models_and_queries_with_asyncpg() {
123        let ir = parse_fixture_ir();
124        let plugin = PythonPlugin::new("pydantic", "asyncpg").unwrap();
125        let config = TargetConfig {
126            language: "python".to_string(),
127            out: "./src/db".to_string(),
128            schema: "pydantic".to_string(),
129            driver: "asyncpg".to_string(),
130            overrides: HashMap::new(),
131        };
132        let files = plugin.generate(&ir, &config).unwrap();
133        assert_eq!(files.len(), 2);
134        assert!(files.iter().any(|f| f.path == "models.py"));
135        assert!(files.iter().any(|f| f.path.ends_with("_queries.py")));
136        assert!(!files.iter().any(|f| f.path == "client.py"));
137    }
138}
sqlcx_core/generator/python/mod.rs

sqlcx_core/generator/python/
mod.rs
