Skip to content

Commit 3900ae4

Browse files
committed
log-server: fix cancelation of completed runs
1 parent 445d0ae commit 3900ae4

File tree

4 files changed

+157
-10
lines changed

4 files changed

+157
-10
lines changed

.changeset/curvy-hats-shave.md

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,5 @@
1+
---
2+
'@lightmill/log-server': patch
3+
---
4+
5+
Fix cancelation of completed runs. Requires database migration.

packages/log-server/__tests__/store.test.ts

Lines changed: 19 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -753,28 +753,40 @@ describe('SQLiteStore#setRunStatus', () => {
753753
it('refuses to update a completed run', async ({ expect, store, e1run1 }) => {
754754
await store.setRunStatus(e1run1, 'completed');
755755
await expect(
756-
store.setRunStatus(e1run1, 'completed'),
756+
store.setRunStatus(e1run1, 'idle'),
757757
).rejects.toThrowErrorMatchingInlineSnapshot(
758-
`[StoreError: Cannot update status of run 1 because the run is completed or canceled]`,
758+
`[StoreError: Cannot change status of run 1 to idle because the run is completed and can only be canceled.]`,
759759
);
760760
await expect(
761-
store.setRunStatus(e1run1, 'canceled'),
761+
store.setRunStatus(e1run1, 'running'),
762762
).rejects.toThrowErrorMatchingInlineSnapshot(
763-
`[StoreError: Cannot update status of run 1 because the run is completed or canceled]`,
763+
`[StoreError: Cannot change status of run 1 to running because the run is completed and can only be canceled.]`,
764764
);
765765
});
766766

767+
it('cancels a completed run', async ({ expect, store, e1run1 }) => {
768+
await store.setRunStatus(e1run1, 'completed');
769+
await expect(
770+
store.setRunStatus(e1run1, 'canceled'),
771+
).resolves.toBeUndefined();
772+
});
773+
767774
it('refuses to update a canceled run', async ({ expect, store, e1run1 }) => {
768775
await store.setRunStatus(e1run1, 'canceled');
769776
await expect(
770-
store.setRunStatus(e1run1, 'completed'),
777+
store.setRunStatus(e1run1, 'running'),
771778
).rejects.toThrowErrorMatchingInlineSnapshot(
772-
`[StoreError: Cannot update status of run 1 because the run is completed or canceled]`,
779+
`[StoreError: Cannot update status of run 1 because the run is canceled.]`,
773780
);
774781
await expect(
775782
store.setRunStatus(e1run1, 'canceled'),
776783
).rejects.toThrowErrorMatchingInlineSnapshot(
777-
`[StoreError: Cannot update status of run 1 because the run is completed or canceled]`,
784+
`[StoreError: Cannot update status of run 1 because the run is canceled.]`,
785+
);
786+
await expect(
787+
store.setRunStatus(e1run1, 'completed'),
788+
).rejects.toThrowErrorMatchingInlineSnapshot(
789+
`[StoreError: Cannot update status of run 1 because the run is canceled.]`,
778790
);
779791
});
780792

Lines changed: 121 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,121 @@
1+
import { type ColumnType, type GeneratedAlways, Kysely, sql } from 'kysely';
2+
import type { JsonObject } from 'type-fest';
3+
4+
export type DbRunId = number;
5+
export type DbExperimentId = number;
6+
export type DbLogId = number;
7+
export type DbLogSequenceId = number;
8+
9+
export type ExperimentTable = {
10+
experimentId: GeneratedAlways<DbExperimentId>;
11+
experimentName: ColumnType<string, string, never>;
12+
// We use ColumnType to indicate that the column cannot be updated.
13+
experimentCreatedAt: ColumnType<string, string, never>;
14+
};
15+
export const runStatuses = [
16+
'idle',
17+
'running',
18+
'completed',
19+
'canceled',
20+
'interrupted',
21+
] as const;
22+
export type RunStatus = (typeof runStatuses)[number];
23+
export type RunTable = {
24+
runId: GeneratedAlways<DbRunId>;
25+
experimentId: ColumnType<DbExperimentId, DbExperimentId, never>;
26+
runName?: ColumnType<string, string, never>;
27+
runStatus: RunStatus;
28+
runCreatedAt: ColumnType<string, string, never>;
29+
};
30+
export type LogSequenceTable = {
31+
sequenceId: GeneratedAlways<DbLogSequenceId>;
32+
runId: ColumnType<DbRunId, number, never>;
33+
sequenceNumber: ColumnType<number, number, never>;
34+
start: ColumnType<number, number, never>;
35+
};
36+
export type LogTable = {
37+
logId: GeneratedAlways<DbLogId>;
38+
sequenceId: ColumnType<DbLogSequenceId, DbLogSequenceId, never>;
39+
logNumber: ColumnType<number, number, never>;
40+
// Logs with no types are used to fill in missing log numbers.
41+
canceledBy?: ColumnType<number, never, never>;
42+
logType?: string;
43+
logValues?: JsonObject;
44+
};
45+
export type RunLogView = {
46+
experimentId: ColumnType<DbExperimentId, never, never>;
47+
experimentName: ColumnType<string, never, never>;
48+
runId: ColumnType<DbRunId, never, never>;
49+
runName: ColumnType<string, never, never>;
50+
runStatus: ColumnType<RunStatus, never, never>;
51+
logId: ColumnType<DbLogId, never, never>;
52+
sequenceId: ColumnType<DbLogSequenceId, never, never>;
53+
logNumber: ColumnType<number, never, never>;
54+
logType?: ColumnType<string, never, never>;
55+
logValues?: ColumnType<JsonObject, never, never>;
56+
};
57+
export type LogPropertyNameTable = {
58+
logId: ColumnType<DbLogId, number, never>;
59+
logPropertyName: ColumnType<string, string, never>;
60+
};
61+
export type Database = {
62+
experiment: ExperimentTable;
63+
run: RunTable;
64+
logSequence: LogSequenceTable;
65+
log: LogTable;
66+
runLogView: RunLogView;
67+
logPropertyName: LogPropertyNameTable;
68+
};
69+
70+
export async function up(db: Kysely<Database>) {
71+
await db.transaction().execute(async (trx) => {
72+
await sql`DROP TRIGGER prevent_run_status_update`.execute(trx);
73+
74+
await sql`
75+
CREATE TRIGGER prevent_canceled_run_status_update
76+
BEFORE UPDATE ON run
77+
WHEN (
78+
SELECT run_status FROM run WHERE run.run_id = OLD.run_id
79+
) = 'canceled'
80+
BEGIN
81+
SELECT RAISE(ABORT, 'Cannot update run status when the run is canceled');
82+
END;
83+
`.execute(trx);
84+
85+
await sql`
86+
CREATE TRIGGER prevent_completed_run_status_update
87+
BEFORE UPDATE ON run
88+
WHEN (
89+
(
90+
(SELECT run_status FROM run WHERE run.run_id = OLD.run_id) = 'completed'
91+
) AND (
92+
NEW.run_status <> 'canceled'
93+
)
94+
)
95+
BEGIN
96+
SELECT RAISE(ABORT, 'Completed runs can only be canceled');
97+
END;
98+
`.execute(trx);
99+
});
100+
}
101+
102+
export async function down(db: Kysely<Database>) {
103+
await db.transaction().execute(async (trx) => {
104+
await sql`
105+
DROP TRIGGER IF EXISTS prevent_canceled_run_status_update
106+
`.execute(trx);
107+
await sql`
108+
DROP TRIGGER IF EXISTS prevent_completed_run_status_update
109+
`.execute(trx);
110+
await sql`
111+
CREATE TRIGGER prevent_run_status_update
112+
BEFORE UPDATE ON run
113+
WHEN (
114+
SELECT run_status FROM run WHERE run.run_id = OLD.run_id
115+
) IN ('completed', 'canceled')
116+
BEGIN
117+
SELECT RAISE(ABORT, 'Cannot update run status when the run is completed or canceled');
118+
END;
119+
`.execute(trx);
120+
});
121+
}

packages/log-server/src/store.ts

Lines changed: 12 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -288,11 +288,20 @@ export class SQLiteStore {
288288
if (
289289
e instanceof SQLiteDB.SqliteError &&
290290
e.code === 'SQLITE_CONSTRAINT_TRIGGER' &&
291-
e.message ===
292-
'Cannot update run status when the run is completed or canceled'
291+
e.message === 'Completed runs can only be canceled'
293292
) {
294293
throw new StoreError(
295-
`Cannot update status of run ${runId} because the run is completed or canceled`,
294+
`Cannot change status of run ${runId} to ${status} because the run is completed and can only be canceled.`,
295+
StoreError.RUN_HAS_ENDED,
296+
{ cause: e },
297+
);
298+
} else if (
299+
e instanceof SQLiteDB.SqliteError &&
300+
e.code === 'SQLITE_CONSTRAINT_TRIGGER' &&
301+
e.message === 'Cannot update run status when the run is canceled'
302+
) {
303+
throw new StoreError(
304+
`Cannot update status of run ${runId} because the run is canceled.`,
296305
StoreError.RUN_HAS_ENDED,
297306
{ cause: e },
298307
);

0 commit comments

Comments
 (0)