Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
110 changes: 110 additions & 0 deletions apps/backend-services/src/document/document-db.service.spec.ts
Original file line number Diff line number Diff line change
Expand Up @@ -48,6 +48,7 @@ const makeDocument = (overrides: Partial<DocumentData> = {}): DocumentData => ({
group_id: "group-1",
created_at: new Date("2024-01-01"),
updated_at: new Date("2024-01-01"),
purged_at: null,
...overrides,
});

Expand Down Expand Up @@ -311,6 +312,115 @@ describe("DocumentDbService", () => {
});
});

describe("findPurgeableEphemeralDocuments", () => {
it("matches any ephemeral target, filters status/unpurged, selects config", async () => {
mockPrismaDocument.findMany.mockResolvedValue([]);

await service.findPurgeableEphemeralDocuments(
[DocumentStatus.complete, DocumentStatus.failed],
50,
);

expect(mockPrismaDocument.findMany).toHaveBeenCalledWith({
where: {
status: { in: [DocumentStatus.complete, DocumentStatus.failed] },
purged_at: null,
workflowVersion: {
is: {
OR: [
{ config: { path: ["metadata", "ephemeral"], equals: true } },
{
config: {
path: ["metadata", "ephemeral", "files"],
equals: true,
},
},
{
config: {
path: ["metadata", "ephemeral", "temporalRecord"],
equals: true,
},
},
],
},
},
},
select: {
id: true,
group_id: true,
workflow_execution_id: true,
workflowVersion: { select: { config: true } },
},
orderBy: { updated_at: "asc" },
take: 50,
});
});

it("extracts the ephemeral policy from each workflow config", async () => {
mockPrismaDocument.findMany.mockResolvedValue([
{
id: "d1",
group_id: "g1",
workflow_execution_id: "wf-1",
workflowVersion: { config: { metadata: { ephemeral: true } } },
},
{
id: "d2",
group_id: "g1",
workflow_execution_id: "wf-2",
workflowVersion: {
config: { metadata: { ephemeral: { files: true } } },
},
},
{
id: "d3",
group_id: "g1",
workflow_execution_id: null,
workflowVersion: { config: { metadata: {} } },
},
]);

const result = await service.findPurgeableEphemeralDocuments(
[DocumentStatus.complete],
50,
);

expect(result).toEqual([
{
id: "d1",
group_id: "g1",
workflow_execution_id: "wf-1",
ephemeral: true,
},
{
id: "d2",
group_id: "g1",
workflow_execution_id: "wf-2",
ephemeral: { files: true, temporalRecord: false },
},
{
id: "d3",
group_id: "g1",
workflow_execution_id: null,
ephemeral: false,
},
]);
});
});

describe("markDocumentPurged", () => {
it("stamps purged_at on the document", async () => {
mockPrismaDocument.update.mockResolvedValue({});

await service.markDocumentPurged("d1");

expect(mockPrismaDocument.update).toHaveBeenCalledWith({
where: { id: "d1" },
data: { purged_at: expect.any(Date) },
});
});
});

describe("findOcrResult", () => {
it("should return the OCR result when found", async () => {
const ocrResult: Partial<OcrResult> = {
Expand Down
110 changes: 109 additions & 1 deletion apps/backend-services/src/document/document-db.service.ts
Original file line number Diff line number Diff line change
@@ -1,3 +1,4 @@
import type { EphemeralConfig } from "@ai-di/graph-workflow";
import { getErrorMessage } from "@ai-di/shared-logging";
import {
DocumentStatus,
Expand All @@ -16,6 +17,43 @@ import {
import { PrismaService } from "../database/prisma.service";
import type { DocumentData } from "./document-db.types";

/** A terminal document to purge, with its workflow's ephemeral policy. */
export interface PurgeableEphemeralDocument {
id: string;
group_id: string;
workflow_execution_id: string | null;
ephemeral: EphemeralConfig;
}

/**
* Extracts `metadata.ephemeral` from a workflow version config (stored as JSON),
* normalizing the object form to explicit booleans. Returns `false` when absent
* or malformed.
*/
function extractEphemeralConfig(
config: Prisma.JsonValue | null | undefined,
): EphemeralConfig {
if (!config || typeof config !== "object" || Array.isArray(config)) {
return false;
}
const metadata = (config as Record<string, unknown>).metadata;
if (!metadata || typeof metadata !== "object" || Array.isArray(metadata)) {
return false;
}
const ephemeral = (metadata as Record<string, unknown>).ephemeral;
if (ephemeral === true) {
return true;
}
if (ephemeral && typeof ephemeral === "object" && !Array.isArray(ephemeral)) {
const obj = ephemeral as Record<string, unknown>;
return {
files: obj.files === true,
temporalRecord: obj.temporalRecord === true,
};
}
return false;
}

@Injectable()
export class DocumentDbService {
constructor(
Expand All @@ -34,7 +72,7 @@ export class DocumentDbService {
* @returns The created document record.
*/
async createDocument(
data: Omit<DocumentData, "created_at" | "updated_at">,
data: Omit<DocumentData, "created_at" | "updated_at" | "purged_at">,
tx?: Prisma.TransactionClient,
): Promise<DocumentData> {
const client = tx ?? this.prisma;
Expand Down Expand Up @@ -361,6 +399,76 @@ export class DocumentDbService {
}
}

/**
* Finds documents eligible for ephemeral cleanup: those processed by a
* workflow whose config declares an ephemeral policy that deletes at least
* one target (`metadata.ephemeral` is `true`, or an object with `files` or
* `temporalRecord` set to `true`), in a terminal status, and not yet purged.
* Ephemerality is configured on the workflow — there is no global or
* per-group setting. Each result carries its workflow's ephemeral policy so
* the janitor knows which targets to delete.
*
* @param statuses - Terminal statuses considered safe to purge.
* @param limit - Maximum number of documents to return in one batch.
* @returns Records (id, group, Temporal execution id, ephemeral policy).
*/
async findPurgeableEphemeralDocuments(
statuses: DocumentStatus[],
limit: number,
): Promise<PurgeableEphemeralDocument[]> {
const rows = await this.prisma.document.findMany({
where: {
status: { in: statuses },
purged_at: null,
workflowVersion: {
is: {
OR: [
{ config: { path: ["metadata", "ephemeral"], equals: true } },
{
config: {
path: ["metadata", "ephemeral", "files"],
equals: true,
},
},
{
config: {
path: ["metadata", "ephemeral", "temporalRecord"],
equals: true,
},
},
],
},
},
},
select: {
id: true,
group_id: true,
workflow_execution_id: true,
workflowVersion: { select: { config: true } },
},
orderBy: { updated_at: "asc" },
take: limit,
});
return rows.map((row) => ({
id: row.id,
group_id: row.group_id,
workflow_execution_id: row.workflow_execution_id,
ephemeral: extractEphemeralConfig(row.workflowVersion?.config),
}));
}

/**
* Stamps a document as purged so the cleanup janitor will not reprocess it.
*
* @param id - The document ID to mark purged.
*/
async markDocumentPurged(id: string): Promise<void> {
await this.prisma.document.update({
where: { id },
data: { purged_at: new Date() },
});
}

/**
* Retrieves the most recent OCR result for a document.
*
Expand Down
2 changes: 2 additions & 0 deletions apps/backend-services/src/document/document.module.ts
Original file line number Diff line number Diff line change
Expand Up @@ -5,13 +5,15 @@ import { UploadNormalizationLimiter } from "../upload/upload-normalization-limit
import { DocumentController } from "./document.controller";
import { DocumentService } from "./document.service";
import { DocumentDbService } from "./document-db.service";
import { EphemeralDocumentCleanupService } from "./ephemeral-document-cleanup.service";
import { PdfNormalizationService } from "./pdf-normalization.service";

@Module({
imports: [BlobStorageModule, TemporalModule],
providers: [
DocumentDbService,
DocumentService,
EphemeralDocumentCleanupService,
PdfNormalizationService,
UploadNormalizationLimiter,
],
Expand Down
12 changes: 9 additions & 3 deletions apps/backend-services/src/document/document.service.ts
Original file line number Diff line number Diff line change
Expand Up @@ -88,7 +88,7 @@ export class DocumentService {
* @returns The created document record.
*/
async createDocument(
data: Omit<DocumentData, "created_at" | "updated_at">,
data: Omit<DocumentData, "created_at" | "updated_at" | "purged_at">,
tx?: Prisma.TransactionClient,
): Promise<DocumentData> {
this.logger.debug(`DocumentService.createDocument: ${data.id}`);
Expand Down Expand Up @@ -195,7 +195,10 @@ export class DocumentService {
);
}

const failedDoc: Omit<DocumentData, "created_at" | "updated_at"> = {
const failedDoc: Omit<
DocumentData,
"created_at" | "updated_at" | "purged_at"
> = {
id: documentId,
title,
original_filename: originalFilename,
Expand Down Expand Up @@ -224,7 +227,10 @@ export class DocumentService {
};
}

const documentData: Omit<DocumentData, "created_at" | "updated_at"> = {
const documentData: Omit<
DocumentData,
"created_at" | "updated_at" | "purged_at"
> = {
id: documentId,
title,
original_filename: originalFilename,
Expand Down
Loading
Loading