orm: support different foreign key types, not just an integer id (#19337)

vlib/orm/orm.v

@@ -453,6 +453,7 @@ pub fn orm_table_gen(table string, q string, defaults bool, def_unique_len int,
 	mut unique_fields := []string{}
 	mut unique := map[string][]string{}
 	mut primary := ''
+	mut primary_typ := 0
 
 	for field in fields {
 		if field.is_arr {
@@ -468,7 +469,7 @@ pub fn orm_table_gen(table string, q string, defaults bool, def_unique_len int,
 		mut field_name := sql_field_name(field)
 		mut ctyp := sql_from_v(sql_field_type(field)) or {
 			field_name = '${field_name}_id'
-			sql_from_v(7)!
+			sql_from_v(primary_typ)!
 		}
 		for attr in field.attrs {
 			match attr.name {
@@ -480,6 +481,7 @@ pub fn orm_table_gen(table string, q string, defaults bool, def_unique_len int,
 				}
 				'primary' {
 					primary = field.name
+					primary_typ = field.typ
 				}
 				'unique' {
 					if attr.arg != '' {

vlib/v/checker/orm.v

@@ -9,6 +9,7 @@ import v.util
 const (
 	v_orm_prefix              = 'V ORM'
 	fkey_attr_name            = 'fkey'
+	pkey_attr_name            = 'primary'
 	connection_interface_name = 'orm.Connection'
 )
 
@@ -47,7 +48,28 @@ fn (mut c Checker) sql_expr(mut node ast.SqlExpr) ast.Type {
 	non_primitive_fields := c.get_orm_non_primitive_fields(fields)
 	mut sub_structs := map[int]ast.SqlExpr{}
 
+	mut has_pkey_attr := false
+	mut pkey_field := ast.StructField{}
+	for field in fields {
+		for attr in field.attrs {
+			if attr.name == checker.pkey_attr_name {
+				if has_pkey_attr {
+					c.orm_error('a struct can only have one primary key', field.pos)
+				}
+				has_pkey_attr = true
+				pkey_field = field
+			}
+		}
+	}
+
 	for field in non_primitive_fields {
+		if c.table.sym(field.typ).kind == .array && !has_pkey_attr {
+			c.orm_error('a struct that has a field that holds an array must have a primary key',
+				field.pos)
+		}
+
+		c.check_orm_struct_field_attributes(field)
+
 		typ := c.get_type_of_field_with_related_table(field)
 
 		mut subquery_expr := ast.SqlExpr{
@@ -98,6 +120,27 @@ fn (mut c Checker) sql_expr(mut node ast.SqlExpr) ast.Type {
 			or_block: ast.OrExpr{}
 		}
 
+		if c.table.sym(field.typ).kind == .array {
+			mut where_expr := subquery_expr.where_expr
+			if mut where_expr is ast.InfixExpr {
+				where_expr.left_type = pkey_field.typ
+				where_expr.right_type = pkey_field.typ
+
+				mut left := where_expr.left
+				if mut left is ast.Ident {
+					left.name = pkey_field.name
+				}
+
+				mut right := where_expr.right
+				if mut right is ast.Ident {
+					mut right_info := right.info
+					if mut right_info is ast.IdentVar {
+						right_info.typ = pkey_field.typ
+					}
+				}
+			}
+		}
+
 		sub_structs[int(typ)] = subquery_expr
 	}
 

vlib/v/checker/tests/orm_fkey_has_pkey.out

@@ -0,0 +1,7 @@
+vlib/v/checker/tests/orm_fkey_has_pkey.vv:5:2: error: V ORM: a struct that has a field that holds an array must have a primary key
+    3 | struct Person {
+    4 |     id    int
+    5 |     child []Child [fkey: 'person_id']
+      |     ~~~~~~~~~~~~~
+    6 | }
+    7 |

vlib/v/checker/tests/orm_fkey_has_pkey.vv

@@ -0,0 +1,18 @@
+import db.sqlite
+
+struct Person {
+	id    int
+	child []Child [fkey: 'person_id']
+}
+
+struct Child {
+	id        int [primary; sql: serial]
+	person_id int
+}
+
+fn main() {
+	db := sqlite.connect(':memory:')!
+	_ := sql db {
+		select from Person
+	}!
+}

vlib/v/checker/tests/orm_multiple_pkeys.out

@@ -0,0 +1,7 @@
+vlib/v/checker/tests/orm_multiple_pkeys.vv:5:2: error: V ORM: a struct can only have one primary key
+    3 | struct Person {
+    4 |     id   int    [primary]
+    5 |     name string [primary]
+      |     ~~~~~~~~~~~
+    6 | }
+    7 |

vlib/v/checker/tests/orm_multiple_pkeys.vv

@@ -0,0 +1,13 @@
+import db.sqlite
+
+struct Person {
+	id   int    [primary]
+	name string [primary]
+}
+
+fn main() {
+	db := sqlite.connect(':memory:')!
+	_ := sql db {
+		select from Person
+	}!
+}

vlib/v/gen/c/orm.v

@@ -302,7 +302,15 @@ fn (mut g Gen) write_orm_insert_with_last_ids(node ast.SqlStmtLine, connection_v
 	}
 
 	fields := node.fields.filter(g.table.sym(it.typ).kind != .array)
-	primary_field_name := g.get_orm_struct_primary_field_name(fields) or { '' }
+	primary_field := g.get_orm_struct_primary_field(fields) or { ast.StructField{} }
+
+	mut is_serial := false
+	for attr in primary_field.attrs {
+		if attr.kind == .plain && attr.name == 'sql' && attr.arg.to_lower() == 'serial' {
+			is_serial = true
+		}
+	}
+	is_serial = is_serial && primary_field.typ == ast.int_type
 
 	for sub in subs {
 		g.sql_stmt_line(sub, connection_var_name, or_expr)
@@ -374,7 +382,7 @@ fn (mut g Gen) write_orm_insert_with_last_ids(node ast.SqlStmtLine, connection_v
 	g.indent--
 	g.writeln('),')
 	g.writeln('.types = __new_array_with_default_noscan(0, 0, sizeof(int), 0),')
-	g.writeln('.primary_column_name = _SLIT("${primary_field_name}"),')
+	g.writeln('.primary_column_name = _SLIT("${primary_field.name}"),')
 	g.writeln('.kinds = __new_array_with_default_noscan(0, 0, sizeof(orm__OperationKind), 0),')
 	g.writeln('.is_and = __new_array_with_default_noscan(0, 0, sizeof(bool), 0),')
 	g.indent--
@@ -384,7 +392,19 @@ fn (mut g Gen) write_orm_insert_with_last_ids(node ast.SqlStmtLine, connection_v
 
 	if arrs.len > 0 {
 		mut id_name := g.new_tmp_var()
+		if is_serial {
+			// use last_insert_id if current struct has `int [primary; sql: serial]`
 			g.writeln('orm__Primitive ${id_name} = orm__int_to_primitive(orm__Connection_name_table[${connection_var_name}._typ]._method_last_id(${connection_var_name}._object));')
+		} else {
+			// else use the primary key value
+			mut sym := g.table.sym(primary_field.typ)
+			mut typ := sym.cname
+			if typ == 'time__Time' {
+				typ = 'time'
+			}
+			g.writeln('orm__Primitive ${id_name} = orm__${typ}_to_primitive(${node.object_var_name}${member_access_type}${c_name(primary_field.name)});')
+		}
+
 		for i, mut arr in arrs {
 			c_field_name := c_name(field_names[i])
 			idx := g.new_tmp_var()
@@ -656,7 +676,10 @@ fn (mut g Gen) write_orm_where_expr(expr ast.Expr, mut fields []string, mut pare
 		}
 		ast.Ident {
 			if g.sql_side == .left {
-				fields << g.get_orm_column_name_from_struct_field(g.get_orm_current_table_field(expr.name))
+				field := g.get_orm_current_table_field(expr.name) or {
+					verror('field "${expr.name}" does not exist on "${g.sql_table_name}"')
+				}
+				fields << g.get_orm_column_name_from_struct_field(field)
 			} else {
 				data << expr
 			}
@@ -686,7 +709,7 @@ fn (mut g Gen) write_orm_where_expr(expr ast.Expr, mut fields []string, mut pare
 // write_orm_select writes C code that calls ORM functions for selecting rows.
 fn (mut g Gen) write_orm_select(node ast.SqlExpr, connection_var_name string, left_expr_string string, or_expr ast.OrExpr) {
 	mut fields := []ast.StructField{}
-	mut primary_field_name := g.get_orm_struct_primary_field_name(node.fields) or { '' }
+	mut primary_field := g.get_orm_struct_primary_field(node.fields) or { ast.StructField{} }
 
 	for field in node.fields {
 		mut skip := false
@@ -732,8 +755,8 @@ fn (mut g Gen) write_orm_select(node ast.SqlExpr, connection_var_name string, le
 	g.writeln('.has_limit = ${node.has_limit},')
 	g.writeln('.has_offset = ${node.has_offset},')
 
-	if primary_field_name != '' {
+	if primary_field.name != '' {
-		g.writeln('.primary = _SLIT("${primary_field_name}"),')
+		g.writeln('.primary = _SLIT("${primary_field.name}"),')
 	}
 
 	select_fields := fields.filter(g.table.sym(it.typ).kind != .array)
@@ -927,11 +950,11 @@ fn (mut g Gen) write_orm_select(node ast.SqlExpr, connection_var_name string, le
 				where_expr.left = left_where_expr
 				where_expr.right = ast.SelectorExpr{
 					pos: right_where_expr.pos
-					field_name: primary_field_name
+					field_name: primary_field.name
 					is_mut: false
 					expr: right_where_expr
 					expr_type: (right_where_expr.info as ast.IdentVar).typ
-					typ: ast.int_type
+					typ: (right_where_expr.info as ast.IdentVar).typ
 					scope: 0
 				}
 				mut sql_expr_select_array := ast.SqlExpr{
@@ -1040,7 +1063,7 @@ fn (g &Gen) get_table_name_by_struct_type(typ ast.Type) string {
 }
 
 // get_orm_current_table_field returns the current processing table's struct field by name.
-fn (g &Gen) get_orm_current_table_field(name string) ast.StructField {
+fn (g &Gen) get_orm_current_table_field(name string) ?ast.StructField {
 	info := g.table.sym(g.table.type_idxs[g.sql_table_name]).struct_info()
 
 	for field in info.fields {
@@ -1049,7 +1072,7 @@ fn (g &Gen) get_orm_current_table_field(name string) ast.StructField {
 		}
 	}
 
-	return ast.StructField{}
+	return none
 }
 
 // get_orm_column_name_from_struct_field converts the struct field to a table column name.
@@ -1071,12 +1094,12 @@ fn (g &Gen) get_orm_column_name_from_struct_field(field ast.StructField) string
 	return name
 }
 
-// get_orm_struct_primary_field_name returns the table's primary column name.
+// get_orm_struct_primary_field returns the table's primary column field.
-fn (_ &Gen) get_orm_struct_primary_field_name(fields []ast.StructField) ?string {
+fn (_ &Gen) get_orm_struct_primary_field(fields []ast.StructField) ?ast.StructField {
 	for field in fields {
 		for attr in field.attrs {
 			if attr.name == 'primary' {
-				return field.name
+				return field
 			}
 		}
 	}

vlib/v/tests/orm_sub_array_struct_test.v

@@ -13,6 +13,18 @@ mut:
 	name      string
 }
 
+struct ParentString {
+	name     string        [primary]
+	children []ChildString [fkey: 'parent_name']
+}
+
+struct ChildString {
+mut:
+	id          int    [primary; sql: serial]
+	parent_name string
+	name        string
+}
+
 fn test_orm_array() {
 	mut db := sqlite.connect(':memory:') or { panic(err) }
 	sql db {
@@ -53,6 +65,45 @@ fn test_orm_array() {
 	assert parent.children[1].name == 'def'
 }
 
+fn test_orm_array_different_pkey_type() {
+	mut db := sqlite.connect(':memory:') or { panic(err) }
+	sql db {
+		create table ParentString
+		create table ChildString
+	}!
+
+	new_parent := ParentString{
+		name: 'test'
+		children: [
+			ChildString{
+				name: 'abc'
+			},
+			ChildString{
+				name: 'def'
+			},
+		]
+	}
+
+	sql db {
+		insert new_parent into ParentString
+	}!
+
+	parents := sql db {
+		select from ParentString where name == 'test'
+	}!
+
+	sql db {
+		drop table ParentString
+		drop table ChildString
+	}!
+
+	parent := parents.first()
+	assert parent.name == new_parent.name
+	assert parent.children.len == new_parent.children.len
+	assert parent.children[0].name == 'abc'
+	assert parent.children[1].name == 'def'
+}
+
 fn test_orm_relationship() {
 	mut db := sqlite.connect(':memory:') or { panic(err) }
 	sql db {
@@ -119,3 +170,70 @@ fn test_orm_relationship() {
 
 	assert children.len == 2
 }
+
+fn test_orm_relationship_different_pkey_type() {
+	mut db := sqlite.connect(':memory:') or { panic(err) }
+	sql db {
+		create table ParentString
+		create table ChildString
+	}!
+
+	mut child := ChildString{
+		name: 'abc'
+	}
+
+	new_parent := ParentString{
+		name: 'test'
+		children: []
+	}
+	sql db {
+		insert new_parent into ParentString
+	}!
+
+	mut parents := sql db {
+		select from ParentString where name == 'test'
+	}!
+
+	mut parent := parents.first()
+	child.parent_name = parent.name
+	child.name = 'atum'
+
+	sql db {
+		insert child into ChildString
+	}!
+
+	child.name = 'bacon'
+
+	sql db {
+		insert child into ChildString
+	}!
+
+	assert parent.name == new_parent.name
+	assert parent.children.len == 0
+
+	parents = sql db {
+		select from ParentString where name == 'test'
+	}!
+
+	parent = parents.first()
+	assert parent.name == new_parent.name
+	assert parent.children.len == 2
+	assert parent.children[0].name == 'atum'
+	assert parent.children[1].name == 'bacon'
+
+	mut children := sql db {
+		select from ChildString
+	}!
+
+	assert children.len == 2
+
+	sql db {
+		drop table ParentString
+	}!
+
+	children = sql db {
+		select from ChildString
+	}!
+
+	assert children.len == 2
+}