125 lines
3.8 KiB
Go
125 lines
3.8 KiB
Go
package generator
|
|
|
|
import (
|
|
"fmt"
|
|
"strings"
|
|
|
|
"github.com/entity-maker/entity-maker/internal/naming"
|
|
)
|
|
|
|
// GenerateModel generates the dataclass model
|
|
func GenerateModel(ctx *Context) (string, error) {
|
|
var b strings.Builder
|
|
|
|
// Imports
|
|
b.WriteString("from dataclasses import dataclass\n")
|
|
|
|
// Check what we need to import
|
|
needsDatetime := NeedsDatetimeImport(ctx.TableInfo.Columns)
|
|
needsDecimal := NeedsDecimalImport(ctx.TableInfo.Columns)
|
|
|
|
if needsDatetime {
|
|
b.WriteString("from datetime import datetime\n")
|
|
}
|
|
if needsDecimal {
|
|
b.WriteString("from decimal import Decimal\n")
|
|
}
|
|
|
|
b.WriteString("\n")
|
|
b.WriteString("from televend_core.databases.base_model import Base\n")
|
|
|
|
// Import related models for foreign keys
|
|
fkImports := make(map[string]string) // module_name -> entity_name
|
|
for _, fk := range ctx.TableInfo.ForeignKeys {
|
|
moduleName := GetRelationshipModuleName(fk.ForeignTableName)
|
|
entityName := GetRelationshipEntityName(fk.ForeignTableName)
|
|
fkImports[moduleName] = entityName
|
|
}
|
|
|
|
// Import enum types
|
|
if len(ctx.TableInfo.EnumTypes) > 0 {
|
|
b.WriteString(fmt.Sprintf("from televend_core.databases.televend_repositories.%s.enum import (\n",
|
|
ctx.ModuleName))
|
|
for _, enumType := range ctx.TableInfo.EnumTypes {
|
|
enumName := naming.ToPascalCase(enumType.TypeName)
|
|
b.WriteString(fmt.Sprintf(" %s,\n", enumName))
|
|
}
|
|
b.WriteString(")\n")
|
|
}
|
|
|
|
// Write foreign key imports
|
|
for moduleName, entityName := range fkImports {
|
|
b.WriteString(fmt.Sprintf("from televend_core.databases.televend_repositories.%s.model import %s\n",
|
|
moduleName, entityName))
|
|
}
|
|
|
|
b.WriteString("\n\n")
|
|
|
|
// Class definition
|
|
b.WriteString("@dataclass\n")
|
|
b.WriteString(fmt.Sprintf("class %s(Base):\n", ctx.EntityName))
|
|
|
|
// Get required and optional columns
|
|
requiredCols := GetRequiredColumns(ctx.TableInfo.Columns)
|
|
optionalCols := GetOptionalColumns(ctx.TableInfo.Columns)
|
|
|
|
// Required fields (non-nullable, no default, not auto-increment, not PK)
|
|
for _, col := range requiredCols {
|
|
fieldName := col.Name
|
|
pythonType := GetPythonTypeForColumn(col, ctx)
|
|
|
|
// Regular field
|
|
b.WriteString(fmt.Sprintf(" %s: %s\n", fieldName, pythonType))
|
|
|
|
// Add relationship field if this is a foreign key and column ends with _id
|
|
// (to avoid name clashes with FK columns that don't follow _id convention)
|
|
if fk := GetForeignKeyForColumn(col.Name, ctx.TableInfo.ForeignKeys); fk != nil {
|
|
if strings.HasSuffix(col.Name, "_id") {
|
|
relationName := GetRelationshipName(col.Name)
|
|
entityName := GetRelationshipEntityName(fk.ForeignTableName)
|
|
b.WriteString(fmt.Sprintf(" %s: %s\n", relationName, entityName))
|
|
}
|
|
}
|
|
}
|
|
|
|
// Empty line between required and optional
|
|
if len(requiredCols) > 0 && len(optionalCols) > 0 {
|
|
b.WriteString("\n")
|
|
}
|
|
|
|
// Optional fields
|
|
for _, col := range optionalCols {
|
|
// Skip primary key, we'll add it at the end
|
|
if col.IsPrimaryKey {
|
|
continue
|
|
}
|
|
|
|
fieldName := col.Name
|
|
pythonType := GetPythonTypeForColumn(col, ctx)
|
|
|
|
b.WriteString(fmt.Sprintf(" %s: %s | None = None\n", fieldName, pythonType))
|
|
|
|
// Add relationship field if this is a foreign key and column ends with _id
|
|
// (to avoid name clashes with FK columns that don't follow _id convention)
|
|
if fk := GetForeignKeyForColumn(col.Name, ctx.TableInfo.ForeignKeys); fk != nil {
|
|
if strings.HasSuffix(col.Name, "_id") {
|
|
relationName := GetRelationshipName(col.Name)
|
|
entityName := GetRelationshipEntityName(fk.ForeignTableName)
|
|
b.WriteString(fmt.Sprintf(" %s: %s | None = None\n", relationName, entityName))
|
|
}
|
|
}
|
|
}
|
|
|
|
// Add primary key at the end
|
|
for _, col := range ctx.TableInfo.Columns {
|
|
if col.IsPrimaryKey {
|
|
b.WriteString("\n")
|
|
pythonType := GetPythonTypeForColumn(col, ctx)
|
|
b.WriteString(fmt.Sprintf(" %s: %s | None = None\n", col.Name, pythonType))
|
|
break
|
|
}
|
|
}
|
|
|
|
return b.String(), nil
|
|
}
|