use super::common::*;
use crate::context::ParsedFile;
use pecto_core::model::*;
pub fn extract(file: &ParsedFile) -> Option<Capability> {
let full_text = &file.source;
if !full_text.contains("@Entity")
&& !full_text.contains("@Column")
&& !full_text.contains("new Schema(")
&& !full_text.contains("mongoose.Schema")
{
return None;
}
let mut entities = Vec::new();
if full_text.contains("@Entity") {
extract_typeorm_entities(full_text, &mut entities);
}
if full_text.contains("Schema(") {
extract_mongoose_schemas(full_text, &mut entities);
}
if entities.is_empty() {
return None;
}
let file_stem = file
.path
.rsplit('/')
.next()
.unwrap_or(&file.path)
.split('.')
.next()
.unwrap_or("unknown");
let capability_name = format!("{}-entity", to_kebab_case(file_stem));
let mut capability = Capability::new(capability_name, file.path.clone());
capability.entities = entities;
Some(capability)
}
fn extract_typeorm_entities(source: &str, entities: &mut Vec<Entity>) {
let mut remaining = source;
while let Some(entity_pos) = remaining.find("@Entity(") {
remaining = &remaining[entity_pos..];
let (class_name, bases) = remaining
.find("class ")
.map(|pos| {
let after = &remaining[pos + 6..];
let line_end = after.find('{').unwrap_or(after.len());
let class_line = &after[..line_end];
let name = class_line
.split([' ', '{', '\n'])
.next()
.unwrap_or("Unknown")
.trim()
.to_string();
let bases = if let Some(ext_pos) = class_line.find("extends ") {
let after_ext = &class_line[ext_pos + 8..];
let base = after_ext
.split([' ', '{', '\n', ','])
.next()
.unwrap_or("")
.trim()
.to_string();
if base.is_empty() {
Vec::new()
} else {
vec![base]
}
} else {
Vec::new()
};
(name, bases)
})
.unwrap_or_else(|| ("Unknown".to_string(), Vec::new()));
let table_name = remaining
.find("@Entity(")
.and_then(|pos| {
let after = &remaining[pos + 8..];
let arg = after.split(')').next()?;
if arg.contains('"') || arg.contains('\'') {
Some(clean_string_literal(arg.trim()))
} else {
None
}
})
.unwrap_or_else(|| class_name.to_lowercase());
let mut fields = Vec::new();
if let Some(class_start) = remaining.find('{') {
let class_body = &remaining[class_start..];
let mut depth = 0;
let mut end = class_body.len();
for (i, c) in class_body.chars().enumerate() {
match c {
'{' => depth += 1,
'}' => {
depth -= 1;
if depth == 0 {
end = i;
break;
}
}
_ => {}
}
}
let body = &class_body[1..end];
extract_typeorm_fields(body, &mut fields);
}
entities.push(Entity {
name: class_name,
table: table_name,
fields,
bases,
});
remaining = &remaining[1..];
if let Some(next) = remaining.find("class ") {
remaining = &remaining[next..];
} else {
break;
}
}
}
fn extract_typeorm_fields(body: &str, fields: &mut Vec<EntityField>) {
let decorators = [
"@PrimaryGeneratedColumn",
"@PrimaryColumn",
"@Column",
"@ManyToOne",
"@OneToMany",
"@ManyToMany",
"@OneToOne",
"@JoinColumn",
];
for line in body.lines() {
let trimmed = line.trim();
let has_decorator = decorators.iter().any(|d| trimmed.starts_with(d));
if !has_decorator {
continue;
}
let mut constraints = Vec::new();
if trimmed.starts_with("@PrimaryGeneratedColumn") {
constraints.push("@PrimaryGeneratedColumn".to_string());
} else if trimmed.starts_with("@PrimaryColumn") {
constraints.push("@PrimaryColumn".to_string());
} else if trimmed.starts_with("@Column") {
constraints.push("@Column".to_string());
if trimmed.contains("nullable: false") || trimmed.contains("nullable:false") {
constraints.push("required".to_string());
}
if trimmed.contains("unique: true") || trimmed.contains("unique:true") {
constraints.push("unique".to_string());
}
} else if trimmed.starts_with("@ManyToOne") {
constraints.push("@ManyToOne".to_string());
} else if trimmed.starts_with("@OneToMany") {
constraints.push("@OneToMany".to_string());
} else if trimmed.starts_with("@ManyToMany") {
constraints.push("@ManyToMany".to_string());
} else if trimmed.starts_with("@OneToOne") {
constraints.push("@OneToOne".to_string());
}
}
let lines: Vec<&str> = body.lines().collect();
let mut i = 0;
while i < lines.len() {
let trimmed = lines[i].trim();
if decorators.iter().any(|d| trimmed.starts_with(d)) {
let mut constraints = Vec::new();
if trimmed.contains("PrimaryGeneratedColumn") || trimmed.contains("PrimaryColumn") {
constraints.push("primary_key".to_string());
}
if trimmed.contains("ManyToOne") {
constraints.push("@ManyToOne".to_string());
}
if trimmed.contains("OneToMany") {
constraints.push("@OneToMany".to_string());
}
if trimmed.contains("ManyToMany") {
constraints.push("@ManyToMany".to_string());
}
if trimmed.contains("nullable: false") {
constraints.push("required".to_string());
}
if trimmed.contains("unique: true") {
constraints.push("unique".to_string());
}
if i + 1 < lines.len() {
let next = lines[i + 1].trim();
if next.contains(':') && !next.starts_with('@') && !next.starts_with("//") {
let parts: Vec<&str> = next.splitn(2, ':').collect();
if parts.len() == 2 {
let name = parts[0]
.trim()
.trim_start_matches("readonly ")
.trim()
.to_string();
let field_type = parts[1].trim().trim_end_matches(';').trim().to_string();
if !name.is_empty() && !name.starts_with("//") {
fields.push(EntityField {
name,
field_type,
constraints,
});
}
}
}
}
}
i += 1;
}
}
fn extract_mongoose_schemas(source: &str, entities: &mut Vec<Entity>) {
let mut remaining = source;
while let Some(pos) = remaining.find("Schema(") {
let before = &remaining[..pos];
let schema_name = before
.rsplit([' ', '\t', '='])
.find(|s| !s.is_empty() && *s != "new" && *s != "mongoose.")
.map(|s| {
s.trim()
.replace("Schema", "")
.replace("const ", "")
.replace("let ", "")
})
.unwrap_or_else(|| "Unknown".to_string());
let name = if schema_name.is_empty() || schema_name == "new" {
"Unknown".to_string()
} else {
schema_name
};
entities.push(Entity {
name: name.clone(),
table: name.to_lowercase(),
fields: Vec::new(), bases: Vec::new(),
});
remaining = &remaining[pos + 7..];
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::context::ParsedFile;
fn parse_file(source: &str, path: &str) -> ParsedFile {
ParsedFile::parse(source.to_string(), path.to_string()).unwrap()
}
#[test]
fn test_typeorm_entity() {
let source = r#"
import { Entity, PrimaryGeneratedColumn, Column, ManyToOne } from 'typeorm';
@Entity('users')
export class User {
@PrimaryGeneratedColumn()
id: number;
@Column({ nullable: false, unique: true })
email: string;
@Column()
name: string;
@ManyToOne(() => Organization)
organization: Organization;
}
"#;
let file = parse_file(source, "entities/user.entity.ts");
let capability = extract(&file).unwrap();
let entity = &capability.entities[0];
assert_eq!(entity.name, "User");
assert_eq!(entity.table, "users");
assert!(
entity.fields.len() >= 3,
"Should find fields, found {}",
entity.fields.len()
);
}
#[test]
fn test_no_entity() {
let source = r#"
export class UserService {
findAll() { return []; }
}
"#;
let file = parse_file(source, "user.service.ts");
assert!(extract(&file).is_none());
}
}