Skip to content

Commit 98e17b2

Browse files
egobraindideex
egobrain
and
dideex
authored
Row locks (#11)
Co-authored-by: Nick Tushin <[email protected]>
1 parent 680bcf7 commit 98e17b2

File tree

4 files changed

+171
-17
lines changed

4 files changed

+171
-17
lines changed

‎include/query.hrl

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -6,12 +6,12 @@
66
data = []:: q:data(),
77
select = #{} :: q:select(),
88
set = #{} :: q:set() | #query{},
9-
tables = [] :: [{real, iolist(), reference()} | q:table()],
9+
tables = [] :: [q:real_table() | q:table()],
1010
joins = [] :: [{q:join_type(), qast:ast_node(), qast:ast_node()}],
1111
group_by = [] :: [qast:ast_node()],
1212
order_by = [] :: q:order(),
1313
on_conflict = #{} :: #{q:conflict_target() => q:conflict_action()},
1414
limit :: non_neg_integer() | undefined,
1515
offset :: non_neg_integer() | undefined,
16-
for_update = false :: boolean()
16+
lock :: {q:row_lock_level(), [q:real_table()]} | undefined
1717
}).

‎src/q.erl

Lines changed: 43 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -6,7 +6,8 @@
66
-export([
77
pipe/2,
88

9-
get/2
9+
get/2,
10+
lookup_tables/2
1011
]).
1112

1213
-export([
@@ -24,6 +25,8 @@
2425
order_by/1, order_by/2,
2526
limit/1, limit/2,
2627
offset/1, offset/2,
28+
29+
lock/1, lock/2, lock/3,
2730
for_update/0, for_update/1,
2831

2932
distinct/0, distinct/1,
@@ -44,22 +47,28 @@
4447
-type order() :: [{qast:ast_node(), asc | desc}].
4548
-type distinct() :: all | [atom()].
4649
-type join_type() :: inner | left | right | full | {left, outer} | {right, outer} | {full, outer}.
50+
-type row_lock_level() :: for_update | for_no_key_update | for_share | for_key_share.
4751
-type qfun() :: fun((query()) -> query()).
4852
-type conflict_target() :: any | [atom()].
4953
-type conflict_action() :: nothing | #{atom() => qast:ast_node()}.
5054

55+
%% internal
56+
-type real_table() :: {real, iolist(), reference()}.
57+
5158
-export_type([query/0]).
5259

5360
-export_type([
5461
model/0,
5562
table/0,
63+
real_table/0,
5664
schema/0,
5765
data/0,
5866
select/0,
5967
set/0,
6068
order/0,
6169
distinct/0,
6270
join_type/0,
71+
row_lock_level/0,
6372
qfun/0,
6473
conflict_target/0,
6574
conflict_action/0
@@ -339,12 +348,44 @@ offset(Value) -> fun(Q) -> offset(Value, Q) end.
339348
offset(Value, Q) ->
340349
Q#query{offset=Value}.
341350

351+
-spec lock(row_lock_level()) -> qfun().
352+
lock(RowLockLevel) ->
353+
lock(RowLockLevel, fun(RealTables) -> RealTables end).
354+
355+
-spec lock(row_lock_level(), fun(([RealTable]) -> [RealTable])) -> qfun() when
356+
RealTable :: real_table().
357+
lock(RowLockLevel, Fun) ->
358+
fun(Q) -> lock(RowLockLevel, Fun, Q) end.
359+
360+
-spec lock(row_lock_level(), fun(([RealTable]) -> [RealTable]), query()) -> query() when
361+
RealTable :: real_table().
362+
lock(RowLockLevel, Fun, #query{tables = AllTables} = Q) ->
363+
RealTables = [T || {real, _Table, _TRef} = T <- AllTables],
364+
Q#query{lock = {RowLockLevel, Fun(RealTables)}}.
365+
366+
-spec lookup_tables(model() | [model()], [RealTable]) -> [RealTable] when
367+
RealTable :: real_table().
368+
%% @THROWS {unknown_table, model()}
369+
lookup_tables(Models, Tables) when is_list(Models) ->
370+
lists:flatmap(
371+
fun(M) ->
372+
TableName = equery_utils:wrap(maps:get(table, get_schema(M))),
373+
RealTables = [T || {real, Table, _TRef} = T <- Tables, Table =:= TableName],
374+
case RealTables of
375+
[] -> error({unknown_table, M});
376+
_ -> RealTables
377+
end
378+
end,
379+
Models);
380+
lookup_tables(Model, Tables) ->
381+
lookup_tables([Model], Tables).
382+
342383
-spec for_update() -> qfun().
343384
for_update() -> fun(Q) -> for_update(Q) end.
344385

345386
-spec for_update(Q) -> Q when Q :: query().
346387
for_update(Q) ->
347-
Q#query{for_update=true}.
388+
lock(for_update, fun(T) -> T end, Q).
348389

349390
-spec data(fun((data()) -> data())) -> qfun().
350391
data(Fun) -> fun(Q) -> data(Fun, Q) end.

‎src/qsql.erl

Lines changed: 17 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -23,7 +23,7 @@ select(#query{
2323
order_by=OrderBy,
2424
limit=Limit,
2525
offset=Offset,
26-
for_update=ForUpdate
26+
lock=Lock
2727
}) ->
2828
{Fields, Opts} = fields_and_opts(Schema, RFields),
2929
qast:exp([
@@ -38,7 +38,7 @@ select(#query{
3838
order_by_exp(OrderBy),
3939
limit_exp(Limit),
4040
offset_exp(Offset),
41-
for_update(ForUpdate)
41+
lock(Lock)
4242
], Opts).
4343

4444
-spec insert(q:query()) -> qast:ast_node().
@@ -286,7 +286,18 @@ conflict_action_exp(Set) when is_map(Set) ->
286286
], qast:raw(","))
287287
]).
288288

289-
for_update(true) ->
290-
qast:raw(" for update");
291-
for_update(false) ->
292-
qast:raw("").
289+
lock(undefined) ->
290+
qast:raw("");
291+
lock({RowLockLevel, Tables}) ->
292+
Aliases = lists:map(fun({real, _Table, TRef}) -> qast:alias(TRef) end, Tables),
293+
qast:exp([
294+
qast:raw(" "),
295+
case RowLockLevel of
296+
for_update -> qast:raw("for update");
297+
for_no_key_update -> qast:raw("for no key update");
298+
for_share -> qast:raw("for share");
299+
for_key_share -> qast:raw("for key share")
300+
end,
301+
qast:raw(" of "),
302+
qast:join(Aliases, qast:raw(","))
303+
]).

‎test/q_tests.erl

Lines changed: 109 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -23,7 +23,6 @@
2323
-define(USER_FIELDS_WITHOUT(L), maps:without(L, ?USER_FIELDS)).
2424
-define(USER_FIELDS_LIST, ?MAPS_TO_LIST(?USER_FIELDS)).
2525
-define(USER_FIELDS_LIST(L), ?MAPS_TO_LIST(?USER_FIELDS(L))).
26-
-define(USER_FIELDS_LIST_WITHOUT(L), ?MAPS_TO_LIST(maps:without(L, ?USER_FIELDS))).
2726

2827
-define(COMMENT_SCHEMA, #{
2928
fields => #{
@@ -40,10 +39,6 @@
4039
-define(TREE_FIELDS, maps:get(fields, tree_m:schema())).
4140
-define(TREE_FIELDS_LIST, ?MAPS_TO_LIST(?TREE_FIELDS)).
4241

43-
-define(COMMENT_FIELDS, maps:get(fields, ?COMMENT_SCHEMA)).
44-
-define(COMMENT_FIELDS_LIST, ?MAPS_TO_LIST(?COMMENT_FIELDS)).
45-
-define(COMMENT_FIELDS_LIST_WITHOUT(L), ?MAPS_TO_LIST(maps:without(L, ?COMMENT_FIELDS))).
46-
4742
schema() -> ?USER_SCHEMA.
4843

4944
%% =============================================================================
@@ -52,7 +47,7 @@ schema() -> ?USER_SCHEMA.
5247

5348
schema_test() ->
5449
?assertEqual(?USER_SCHEMA, q:get(schema, q:from(?USER_SCHEMA))),
55-
?assertEqual(?USER_SCHEMA#{model => ?MODULE}, q:get(schema, q:from(?MODULE))).
50+
?assertEqual(maps:put(model, ?MODULE, ?USER_SCHEMA), q:get(schema, q:from(?MODULE))).
5651

5752
data_test() ->
5853
?assertEqual(
@@ -94,11 +89,118 @@ q_test() ->
9489
"on (\"__alias-0\".\"id\" = \"__alias-1\".\"author\") "
9590
"where ((\"__alias-0\".\"name\" = $1) and (\"__alias-1\".\"text\" = $2)) "
9691
"order by \"__alias-0\".\"name\" ASC,\"__alias-0\".\"id\" DESC "
97-
"for update">>,
92+
"for update of \"__alias-0\"">>,
9893
Sql),
9994
?assertEqual([<<"test1">>, <<"test2">>], Args),
10095
?assertEqual({model, undefined, ?USER_FIELDS_LIST}, Feilds).
10196

97+
q_lock_with_several_tables_test() ->
98+
{Sql, Args, Feilds} = to_sql(
99+
qsql:select(q:pipe(q:from(?USER_SCHEMA), [
100+
q:where(
101+
fun([#{name := Name}]) ->
102+
pg_sql:'=:='(Name, <<"test1">>)
103+
end),
104+
q:using(?COMMENT_SCHEMA),
105+
q:where(
106+
fun([#{id := UserId}, #{author := AuthorId, text := Name}]) ->
107+
pg_sql:'andalso'(
108+
pg_sql:'=:='(UserId, AuthorId),
109+
pg_sql:'=:='(Name, <<"test2">>))
110+
end),
111+
q:order_by(
112+
fun([#{name := Name, id := Id}|_]) ->
113+
[{Name, asc}, {Id, desc}]
114+
end),
115+
q:lock(for_update, fun(Tables) -> q:lookup_tables(?COMMENT_SCHEMA, Tables) end)
116+
]))),
117+
?assertEqual(
118+
<<"select "
119+
"\"__alias-0\".\"id\" as \"id\","
120+
"\"__alias-0\".\"name\" as \"name\","
121+
"\"__alias-0\".\"password\" as \"password\","
122+
"\"__alias-0\".\"salt\" as \"salt\" "
123+
"from \"users\" as \"__alias-0\","
124+
"\"comments\" as \"__alias-1\" "
125+
"where ((\"__alias-0\".\"name\" = $1) and ((\"__alias-0\".\"id\" = \"__alias-1\".\"author\") and (\"__alias-1\".\"text\" = $2))) "
126+
"order by \"__alias-0\".\"name\" ASC,\"__alias-0\".\"id\" DESC "
127+
"for update of \"__alias-1\"">>,
128+
Sql),
129+
?assertEqual([<<"test1">>, <<"test2">>], Args),
130+
?assertEqual({model, undefined, ?USER_FIELDS_LIST}, Feilds).
131+
132+
q_lock_lookup_error_test() ->
133+
?assertException(
134+
error,
135+
{unknown_table, #{table := <<"comments">>}},
136+
qsql:select(q:pipe(q:from(?USER_SCHEMA), [
137+
q:where(
138+
fun([#{name := Name}]) ->
139+
pg_sql:'=:='(Name, <<"test1">>)
140+
end),
141+
q:lock(for_update, fun(Tables) -> q:lookup_tables(?COMMENT_SCHEMA, Tables) end)
142+
]))).
143+
144+
q_lock_for_no_key_update_test() ->
145+
{Sql, _Args, _Feilds} = to_sql(
146+
qsql:select(q:pipe(q:from(?USER_SCHEMA), [
147+
q:where(
148+
fun([#{name := Name}]) ->
149+
pg_sql:'=:='(Name, <<"test1">>)
150+
end),
151+
q:lock(for_no_key_update)
152+
]))),
153+
?assertEqual(
154+
<<"select "
155+
"\"__alias-0\".\"id\" as \"id\","
156+
"\"__alias-0\".\"name\" as \"name\","
157+
"\"__alias-0\".\"password\" as \"password\","
158+
"\"__alias-0\".\"salt\" as \"salt\" "
159+
"from \"users\" as \"__alias-0\" "
160+
"where (\"__alias-0\".\"name\" = $1) "
161+
"for no key update of \"__alias-0\"">>,
162+
Sql).
163+
164+
q_lock_for_share_test() ->
165+
{Sql, _Args, _Feilds} = to_sql(
166+
qsql:select(q:pipe(q:from(?USER_SCHEMA), [
167+
q:where(
168+
fun([#{name := Name}]) ->
169+
pg_sql:'=:='(Name, <<"test1">>)
170+
end),
171+
q:lock(for_share)
172+
]))),
173+
?assertEqual(
174+
<<"select "
175+
"\"__alias-0\".\"id\" as \"id\","
176+
"\"__alias-0\".\"name\" as \"name\","
177+
"\"__alias-0\".\"password\" as \"password\","
178+
"\"__alias-0\".\"salt\" as \"salt\" "
179+
"from \"users\" as \"__alias-0\" "
180+
"where (\"__alias-0\".\"name\" = $1) "
181+
"for share of \"__alias-0\"">>,
182+
Sql).
183+
184+
q_lock_for_key_share_test() ->
185+
{Sql, _Args, _Feilds} = to_sql(
186+
qsql:select(q:pipe(q:from(?USER_SCHEMA), [
187+
q:where(
188+
fun([#{name := Name}]) ->
189+
pg_sql:'=:='(Name, <<"test1">>)
190+
end),
191+
q:lock(for_key_share)
192+
]))),
193+
?assertEqual(
194+
<<"select "
195+
"\"__alias-0\".\"id\" as \"id\","
196+
"\"__alias-0\".\"name\" as \"name\","
197+
"\"__alias-0\".\"password\" as \"password\","
198+
"\"__alias-0\".\"salt\" as \"salt\" "
199+
"from \"users\" as \"__alias-0\" "
200+
"where (\"__alias-0\".\"name\" = $1) "
201+
"for key share of \"__alias-0\"">>,
202+
Sql).
203+
102204
q_from_query_test() ->
103205
BaseQuery = q:pipe(q:from(?USER_SCHEMA), [
104206
q:select(fun([#{id := Id, name := Name}]) -> #{num => Id*2, name => Name} end)

0 commit comments

Comments
 (0)