diff --git a/pkg/codegen/modelgenerate/generate_model.go b/pkg/codegen/modelgenerate/generate_model.go index 4118760..6f21a7b 100644 --- a/pkg/codegen/modelgenerate/generate_model.go +++ b/pkg/codegen/modelgenerate/generate_model.go @@ -21,6 +21,14 @@ func GenerateIDType(table schema.Table) *ast.GenDecl { } } +func fkFieldName(col schema.Column) string { + if col.IsNonCodeTableForeignKey() { + return textutils.SnakeToCamel(strings.TrimSuffix(col.Name, "_id")) + "ID" + } else { + return textutils.SnakeToCamel(col.Name) + } +} + // GenerateModelAST produces an AST for a struct type corresponding to the model. // TODO: generate the right field types here based on column types. func GenerateModelAST(table schema.Table) *ast.GenDecl { @@ -37,9 +45,9 @@ func GenerateModelAST(table schema.Table) *ast.GenDecl { Tag: &ast.BasicLit{Kind: token.STRING, Value: "`db:\"rowid\" json:\"id\"`"}, }) default: - if col.IsForeignKey && strings.HasSuffix(col.Name, "_id") { + if col.IsNonCodeTableForeignKey() { fields = append(fields, &ast.Field{ - Names: []*ast.Ident{ast.NewIdent(textutils.SnakeToCamel(strings.TrimSuffix(col.Name, "_id")) + "ID")}, + Names: []*ast.Ident{ast.NewIdent(fkFieldName(col))}, Type: ast.NewIdent(schema.TypenameFromTablename(col.ForeignKeyTargetTable) + "ID"), Tag: &ast.BasicLit{Kind: token.STRING, Value: fmt.Sprintf("`db:\"%s\" json:\"%s\"`", col.Name, col.Name)}, }) @@ -191,11 +199,13 @@ func GenerateSaveItemFunc(tbl schema.Table) *ast.FuncDecl { return funcDecl } +func getByIDFuncName(tblname string) string { + return "Get" + schema.TypenameFromTablename(tblname) + "ByID" +} + // GenerateGetItemByIDFunc produces an AST for the `GetXyzByID()` function. // E.g., a table with `table.TypeName = "foods"` will produce a "GetFoodByID()" function. func GenerateGetItemByIDFunc(tbl schema.Table) *ast.FuncDecl { - funcName := "Get" + tbl.GoTypeName + "ByID" - recv := &ast.FieldList{List: []*ast.Field{{Names: []*ast.Ident{ast.NewIdent("db")}, Type: ast.NewIdent("DB")}}} arg := &ast.FieldList{List: []*ast.Field{{Names: []*ast.Ident{ast.NewIdent("id")}, Type: ast.NewIdent(tbl.TypeIDName)}}} result := &ast.FieldList{List: []*ast.Field{{Names: []*ast.Ident{ast.NewIdent("ret")}, Type: ast.NewIdent(tbl.GoTypeName)}, {Names: []*ast.Ident{ast.NewIdent("err")}, Type: ast.NewIdent("error")}}} @@ -228,7 +238,7 @@ func GenerateGetItemByIDFunc(tbl schema.Table) *ast.FuncDecl { funcDecl := &ast.FuncDecl{ Recv: recv, - Name: ast.NewIdent(funcName), + Name: ast.NewIdent(getByIDFuncName(tbl.TableName)), Type: &ast.FuncType{Params: arg, Results: result}, Body: funcBody, } diff --git a/pkg/schema/table.go b/pkg/schema/table.go index 3763e76..2d9e741 100644 --- a/pkg/schema/table.go +++ b/pkg/schema/table.go @@ -1,5 +1,7 @@ package schema +import "strings" + // Column represents a single column in a table. type Column struct { TableName string `db:"table_name"` @@ -20,6 +22,10 @@ func (c Column) IsNullableForeignKey() bool { return !c.IsNotNull && !c.IsPrimaryKey && c.IsForeignKey } +func (c Column) IsNonCodeTableForeignKey() bool { + return c.IsForeignKey && strings.HasSuffix(c.Name, "_id") +} + // Table is a single SQLite table. type Table struct { TableName string `db:"name"`