Coverage for tests/query/test_update_execution.py: 100%

1"""Update query construction, compilation, and execution acceptance tests.""" 

2 

3from __future__ import annotations 

4 

5from collections.abc import Callable 

6from typing import cast 

7 

8from snektest import assert_eq, assert_is_none, assert_raises, test 

9 

10from snekql import ( 

11 MISSING, 

12 Database, 

13 Fetched, 

14 Integer, 

15 Model, 

16 Pending, 

17 QueryConstructionError, 

18 Text, 

19 insert, 

20 select, 

21 update, 

22) 

23from snekql.query import compile_write_sql 

24from tests.helpers import NULL_LOGGER 

25 

26 

27@test(mark="fast") 

28def update_compilation_accepts_multiple_quoted_assignments() -> None: 

29 """Update compiles multiple assignments in caller-provided order.""" 

30 

31 class Order[S = Pending](Model[S, "Order[Fetched]"]): 

32 """Table model with identifiers requiring SQLite quoting.""" 

33 

34 __tablename__ = "select" 

35 status: Order.Col[str] = Text(nullable=False) 

36 where: Order.Col[str] = Text(nullable=False) 

37 

38 query = ( 

39 update(Order) 

40 .set( 

41 Order.where.to("new"), 

42 Order.status.to("done"), 

43 ) 

44 .where(Order.where.ne("old")) 

45 ) 

46 

47 sql, params = compile_write_sql(query) 

48 

49 expected_sql = 'UPDATE "select" SET "where" = ?, "status" = ? WHERE ("where" != ?)' 

50 assert_eq(sql, expected_sql) 

51 assert_eq(params, ("new", "done", "old")) 

52 

53 

54@test(mark="fast") 

55def update_assignments_must_belong_to_target_model() -> None: 

56 """Update set() rejects assignments built from another table model.""" 

57 

58 class User[S = Pending](Model[S, "User[Fetched]"]): 

59 """Target table model for ownership checks.""" 

60 

61 email: User.Col[str] = Text(nullable=False) 

62 

63 class AuditLog[S = Pending](Model[S, "AuditLog[Fetched]"]): 

64 """Unrelated table model for ownership checks.""" 

65 

66 message: AuditLog.Col[str] = Text(nullable=False) 

67 

68 set_assignments = cast("Callable[..., object]", update(User).set) 

69 

70 with assert_raises(QueryConstructionError): 

71 _ = set_assignments(AuditLog.message.to("wrong table")) 

72 

73 

74@test(mark="fast") 

75def update_rejects_generated_and_primary_key_assignments() -> None: 

76 """Generated and primary key columns are not update-assignable.""" 

77 

78 class User[S = Pending](Model[S, "User[Fetched]"]): 

79 """Table model with both protected and updateable columns.""" 

80 

81 account_id: User.Col[int] = Integer(primary_key=True) 

82 email: User.Col[str] = Text(nullable=False) 

83 revision: User.GenCol[int] = Integer(default=MISSING) 

84 

85 with assert_raises(QueryConstructionError): 

86 _ = update(User).set(User.revision.to(2)) 

87 

88 with assert_raises(QueryConstructionError): 

89 _ = update(User).set(User.account_id.to(2)) 

90 

91 

92@test(mark="medium") 

93async def update_execute_returns_none_and_supports_explicit_all() -> None: 

94 """tx.execute(update(...).all()) returns None after full-table update.""" 

95 

96 class User[S = Pending](Model[S, "User[Fetched]"]): 

97 """Table model updated through the async runtime.""" 

98 

99 id: User.GenCol[int] = Integer(primary_key=True, default=MISSING) 

100 email: User.Col[str] = Text(nullable=False) 

101 status: User.Col[str] = Text(nullable=False, default="active") 

102 

103 database = await Database.initialize( 

104 logger=NULL_LOGGER, database=":memory:", models=[User] 

105 ) 

106 try: 

107 async with database.transaction() as tx: 

108 await tx.execute(insert(User(email="a@example.com"))) 

109 await tx.execute(insert(User(email="b@example.com"))) 

110 

111 result = await tx.execute( 

112 update(User).set(User.status.to("disabled")).all(), 

113 ) 

114 statuses = await tx.fetch_all( 

115 select(User.status).all().order_by(User.id.asc()), 

116 ) 

117 finally: 

118 await database.close() 

119 

120 assert_is_none(result) 

121 assert_eq(statuses, ["disabled", "disabled"])