diff --git a/server/src/services/company-artifacts.ts b/server/src/services/company-artifacts.ts index 4551ac27acf..aa06d102aed 100644 --- a/server/src/services/company-artifacts.ts +++ b/server/src/services/company-artifacts.ts @@ -1,5 +1,5 @@ import { buffer } from "node:stream/consumers"; -import { and, desc, eq, isNotNull, isNull, notInArray, or, sql, type SQL } from "drizzle-orm"; +import { and, desc, eq, inArray, isNotNull, isNull, notInArray, or, sql, type SQL } from "drizzle-orm"; import { alias } from "drizzle-orm/pg-core"; import type { Db } from "@paperclipai/db"; import { @@ -32,6 +32,7 @@ import type { StorageService } from "../storage/types.js"; const TEXT_PREVIEW_BYTES = 4096; const PREVIEW_TEXT_MAX_LENGTH = 280; const GROUP_PREVIEW_ARTIFACT_LIMIT = 3; +const GROUPED_ARTIFACT_FETCH_LIMIT = COMPANY_ARTIFACTS_MAX_LIMIT * 10; type ArtifactCursor = { updatedAt: string; @@ -191,19 +192,33 @@ function pageByCursor( return { page, nextCursor }; } -async function loadIssueGroupingRows(db: Db, companyId: string) { - const rows = await db - .select({ - id: issues.id, - parentId: issues.parentId, - identifier: issues.identifier, - title: issues.title, - updatedAt: issues.updatedAt, - }) - .from(issues) - .where(eq(issues.companyId, companyId)); +async function loadIssueGroupingRows(db: Db, companyId: string, seedIssueIds: Iterable) { + const rowsById = new Map(); + let pending = [...new Set(seedIssueIds)]; - return new Map(rows.map((row): [string, IssueGroupingRow] => [row.id, row])); + while (pending.length > 0) { + const rows = await db + .select({ + id: issues.id, + parentId: issues.parentId, + identifier: issues.identifier, + title: issues.title, + updatedAt: issues.updatedAt, + }) + .from(issues) + .where(and(eq(issues.companyId, companyId), inArray(issues.id, pending))); + + const nextPending = new Set(); + for (const row of rows) { + rowsById.set(row.id, row); + if (row.parentId && !rowsById.has(row.parentId)) { + nextPending.add(row.parentId); + } + } + pending = [...nextPending]; + } + + return rowsById; } function getIssueSummary(issue: IssueGroupingRow) { @@ -312,6 +327,7 @@ export function companyArtifactsService(db: Db, storage?: StorageService) { if (!company) throw notFound("Company not found"); const fetchLimit = Math.min(query.limit + 1, COMPANY_ARTIFACTS_MAX_LIMIT + 1); + const sourceFetchLimit = groupBy ? GROUPED_ARTIFACT_FETCH_LIMIT : fetchLimit; const q = query.q ? `%${escapeLikePattern(query.q)}%` : null; const artifacts: CompanyArtifact[] = []; const workProductAttachmentIds = new Set(); @@ -328,6 +344,7 @@ export function companyArtifactsService(db: Db, storage?: StorageService) { ]; const documentCursor = groupBy ? undefined : cursorCondition(sql`${documents.updatedAt}`, documentArtifactId, cursor); if (documentCursor) documentConditions.push(documentCursor); + if (groupBy === "task" && query.groupIssueId) documentConditions.push(eq(issues.id, query.groupIssueId)); if (query.projectId) documentConditions.push(eq(issues.projectId, query.projectId)); if (q) { documentConditions.push(sql`( @@ -392,7 +409,7 @@ export function companyArtifactsService(db: Db, storage?: StorageService) { ) .where(and(...documentConditions)) .orderBy(desc(documents.updatedAt), desc(documentArtifactId)); - const documentRows = groupBy ? await documentRowsQuery : await documentRowsQuery.limit(fetchLimit); + const documentRows = await documentRowsQuery.limit(sourceFetchLimit); for (const row of documentRows) { const identifier = row.issueIdentifier ?? row.issueId; @@ -432,6 +449,11 @@ export function companyArtifactsService(db: Db, storage?: StorageService) { : cursorCondition(sql`${issueWorkProducts.updatedAt}`, workProductArtifactId, cursor); const workProductKind = contentTypeKindCondition(workProductContentType, query.kind); if (workProductCursor) workProductConditions.push(workProductCursor); + if (groupBy === "task" && query.groupIssueId) { + const selectedIssueCondition = eq(issues.id, query.groupIssueId); + workProductBaseConditions.push(selectedIssueCondition); + workProductConditions.push(selectedIssueCondition); + } if (workProductKind) { workProductBaseConditions.push(workProductKind); workProductConditions.push(workProductKind); @@ -499,7 +521,7 @@ export function companyArtifactsService(db: Db, storage?: StorageService) { ) .where(and(...workProductConditions)) .orderBy(desc(issueWorkProducts.updatedAt), desc(workProductArtifactId)); - const workProductRows = groupBy ? await workProductRowsQuery : await workProductRowsQuery.limit(fetchLimit); + const workProductRows = await workProductRowsQuery.limit(sourceFetchLimit); const workProductAttachmentRows = await db .select({ @@ -513,7 +535,8 @@ export function companyArtifactsService(db: Db, storage?: StorageService) { eq(issues.companyId, issueWorkProducts.companyId), ), ) - .where(and(...workProductBaseConditions, sql`${issueWorkProducts.metadata}->>'attachmentId' IS NOT NULL`)); + .where(and(...workProductBaseConditions, sql`${issueWorkProducts.metadata}->>'attachmentId' IS NOT NULL`)) + .limit(sourceFetchLimit); for (const row of workProductAttachmentRows) { if (row.attachmentId) { @@ -561,6 +584,7 @@ export function companyArtifactsService(db: Db, storage?: StorageService) { : cursorCondition(sql`${issueAttachments.updatedAt}`, attachmentArtifactId, cursor); const attachmentKind = contentTypeKindCondition(sql`${assets.contentType}`, query.kind); if (attachmentCursor) attachmentConditions.push(attachmentCursor); + if (groupBy === "task" && query.groupIssueId) attachmentConditions.push(eq(issues.id, query.groupIssueId)); if (attachmentKind) attachmentConditions.push(attachmentKind); if (query.projectId) attachmentConditions.push(eq(issues.projectId, query.projectId)); if (q) { @@ -620,7 +644,7 @@ export function companyArtifactsService(db: Db, storage?: StorageService) { ) .where(and(...attachmentConditions)) .orderBy(desc(issueAttachments.updatedAt), desc(attachmentArtifactId)); - const attachmentRows = groupBy ? await attachmentRowsQuery : await attachmentRowsQuery.limit(fetchLimit); + const attachmentRows = await attachmentRowsQuery.limit(sourceFetchLimit); const attachmentArtifacts = await Promise.all(attachmentRows.map(async (row): Promise => { if (workProductAttachmentIds.has(row.attachmentId)) return null; @@ -666,7 +690,9 @@ export function companyArtifactsService(db: Db, storage?: StorageService) { return { artifacts: page, nextCursor }; } - const issueRows = await loadIssueGroupingRows(db, companyId); + const issueSeedIds = new Set(artifacts.map((artifact) => artifact.issue.id)); + if (query.groupIssueId) issueSeedIds.add(query.groupIssueId); + const issueRows = await loadIssueGroupingRows(db, companyId, issueSeedIds); const groups = buildArtifactGroups({ artifacts: sorted, companyPrefix: company.issuePrefix,