Adding api for uploading and retrieving screenshots, also adding improvements to get ancestors

apps/server/src/api/v2/chats/[id]/index.ts

@@ -1,10 +1,12 @@
 import { Hono } from 'hono';
 import GET from './GET';
+import SCREENSHOT from './screenshot';
 import SHARING from './sharing';
 
 const app = new Hono();
 
 app.route('/', GET);
 app.route('/sharing', SHARING);
+app.route('/screenshot', SCREENSHOT);
 
 export default app;

apps/server/src/api/v2/chats/[id]/screenshot/GET.ts

@@ -0,0 +1,57 @@
+import { checkPermission } from '@buster/access-controls';
+import { getChatById } from '@buster/database/queries';
+import { getAssetScreenshotSignedUrl } from '@buster/search';
+import { AssetIdParamsSchema, type GetScreenshotResponse } from '@buster/server-shared/screenshots';
+import { zValidator } from '@hono/zod-validator';
+import { Hono } from 'hono';
+import { HTTPException } from 'hono/http-exception';
+
+const app = new Hono().get('/', zValidator('param', AssetIdParamsSchema), async (c) => {
+  const { id } = c.req.valid('param');
+  const user = c.get('busterUser');
+
+  const chat = await getChatById(id);
+  if (!chat) {
+    throw new HTTPException(404, { message: 'Chat not found' });
+  }
+
+  const permission = await checkPermission({
+    userId: user.id,
+    assetId: id,
+    assetType: 'chat',
+    requiredRole: 'can_view',
+    workspaceSharing: chat.workspaceSharing,
+    organizationId: chat.organizationId,
+  });
+
+  if (!permission.hasAccess) {
+    throw new HTTPException(403, {
+      message: 'You do not have permission to view this chat',
+    });
+  }
+
+  let signedUrl = '';
+  let success = true;
+
+  try {
+    signedUrl = await getAssetScreenshotSignedUrl({
+      assetType: 'chat',
+      assetId: id,
+      organizationId: chat.organizationId,
+    });
+  } catch (error) {
+    console.error('Failed to generate chat screenshot URL', {
+      chatId: id,
+      error,
+    });
+    success = false;
+  }
+
+  const response: GetScreenshotResponse = {
+    success,
+    url: signedUrl,
+  };
+  return c.json(response);
+});
+
+export default app;

apps/server/src/api/v2/chats/[id]/screenshot/PUT.ts

@@ -0,0 +1,53 @@
+import { checkPermission } from '@buster/access-controls';
+import { getChatById } from '@buster/database/queries';
+import {
+  AssetIdParamsSchema,
+  PutScreenshotRequestSchema,
+  type PutScreenshotResponse,
+} from '@buster/server-shared/screenshots';
+import { zValidator } from '@hono/zod-validator';
+import { Hono } from 'hono';
+import { HTTPException } from 'hono/http-exception';
+import { uploadScreenshotHandler } from '../../../../../shared-helpers/upload-screenshot-handler';
+
+const app = new Hono().put(
+  '/',
+  zValidator('json', PutScreenshotRequestSchema),
+  zValidator('param', AssetIdParamsSchema),
+  async (c) => {
+    const assetId = c.req.valid('param').id;
+    const { base64Image } = c.req.valid('json');
+    const user = c.get('busterUser');
+
+    const chat = await getChatById(assetId);
+    if (!chat) {
+      throw new HTTPException(404, { message: 'Chat not found' });
+    }
+
+    const permission = await checkPermission({
+      userId: user.id,
+      assetId,
+      assetType: 'chat',
+      requiredRole: 'can_edit',
+      workspaceSharing: chat.workspaceSharing,
+      organizationId: chat.organizationId,
+    });
+
+    if (!permission.hasAccess) {
+      throw new HTTPException(403, {
+        message: 'You do not have permission to upload a screenshot for this chat',
+      });
+    }
+
+    const result: PutScreenshotResponse = await uploadScreenshotHandler({
+      assetType: 'chat',
+      assetId,
+      base64Image,
+      organizationId: chat.organizationId,
+    });
+
+    return c.json(result);
+  }
+);
+
+export default app;

apps/server/src/api/v2/chats/[id]/screenshot/index.ts

@@ -0,0 +1,7 @@
+import { Hono } from 'hono';
+import GET from './GET';
+import PUT from './PUT';
+
+const app = new Hono().route('/', GET).route('/', PUT);
+
+export default app;

apps/server/src/api/v2/dashboards/[id]/index.ts

@@ -1,11 +1,13 @@
 import { Hono } from 'hono';
 import '../../../../types/hono.types';
 import dashboardByIdRoutes from './GET';
+import SCREENSHOT from './screenshot';
 import SHARING from './sharing';
 
 const app = new Hono()
   //  /dashboards/:id GET
   .route('/', dashboardByIdRoutes)
-  .route('/sharing', SHARING);
+  .route('/sharing', SHARING)
+  .route('/screenshot', SCREENSHOT);
 
 export default app;

apps/server/src/api/v2/dashboards/[id]/screenshot/GET.ts

@@ -0,0 +1,57 @@
+import { checkPermission } from '@buster/access-controls';
+import { getDashboardById } from '@buster/database/queries';
+import { getAssetScreenshotSignedUrl } from '@buster/search';
+import { AssetIdParamsSchema, type GetScreenshotResponse } from '@buster/server-shared/screenshots';
+import { zValidator } from '@hono/zod-validator';
+import { Hono } from 'hono';
+import { HTTPException } from 'hono/http-exception';
+
+const app = new Hono().get('/', zValidator('param', AssetIdParamsSchema), async (c) => {
+  const dashboardId = c.req.valid('param').id;
+  const user = c.get('busterUser');
+
+  const dashboard = await getDashboardById({ dashboardId });
+  if (!dashboard) {
+    throw new HTTPException(404, { message: 'Dashboard not found' });
+  }
+
+  const permission = await checkPermission({
+    userId: user.id,
+    assetId: dashboardId,
+    assetType: 'dashboard_file',
+    requiredRole: 'can_view',
+    workspaceSharing: dashboard.workspaceSharing,
+    organizationId: dashboard.organizationId,
+  });
+
+  if (!permission.hasAccess) {
+    throw new HTTPException(403, {
+      message: 'You do not have permission to view this dashboard',
+    });
+  }
+
+  let signedUrl = '';
+  let success = true;
+
+  try {
+    signedUrl = await getAssetScreenshotSignedUrl({
+      assetType: 'dashboard_file',
+      assetId: dashboardId,
+      organizationId: dashboard.organizationId,
+    });
+  } catch (error) {
+    console.error('Failed to generate dashboard screenshot URL', {
+      dashboardId,
+      error,
+    });
+    success = false;
+  }
+
+  const response: GetScreenshotResponse = {
+    success,
+    url: signedUrl,
+  };
+  return c.json(response);
+});
+
+export default app;

apps/server/src/api/v2/dashboards/[id]/screenshot/PUT.ts

@@ -0,0 +1,53 @@
+import { checkPermission } from '@buster/access-controls';
+import { getDashboardById } from '@buster/database/queries';
+import {
+  AssetIdParamsSchema,
+  PutScreenshotRequestSchema,
+  type PutScreenshotResponse,
+} from '@buster/server-shared/screenshots';
+import { zValidator } from '@hono/zod-validator';
+import { Hono } from 'hono';
+import { HTTPException } from 'hono/http-exception';
+import { uploadScreenshotHandler } from '../../../../../shared-helpers/upload-screenshot-handler';
+
+const app = new Hono().put(
+  '/',
+  zValidator('json', PutScreenshotRequestSchema),
+  zValidator('param', AssetIdParamsSchema),
+  async (c) => {
+    const assetId = c.req.valid('param').id;
+    const { base64Image } = c.req.valid('json');
+    const user = c.get('busterUser');
+
+    const dashboard = await getDashboardById({ dashboardId: assetId });
+    if (!dashboard) {
+      throw new HTTPException(404, { message: 'Dashboard not found' });
+    }
+
+    const permission = await checkPermission({
+      userId: user.id,
+      assetId,
+      assetType: 'dashboard_file',
+      requiredRole: 'can_edit',
+      workspaceSharing: dashboard.workspaceSharing,
+      organizationId: dashboard.organizationId,
+    });
+
+    if (!permission.hasAccess) {
+      throw new HTTPException(403, {
+        message: 'You do not have permission to upload a screenshot for this dashboard',
+      });
+    }
+
+    const result: PutScreenshotResponse = await uploadScreenshotHandler({
+      assetType: 'dashboard_file',
+      assetId,
+      base64Image,
+      organizationId: dashboard.organizationId,
+    });
+
+    return c.json(result);
+  }
+);
+
+export default app;

apps/server/src/api/v2/dashboards/[id]/screenshot/index.ts

@@ -0,0 +1,7 @@
+import { Hono } from 'hono';
+import GET from './GET';
+import PUT from './PUT';
+
+const app = new Hono().route('/', GET).route('/', PUT);
+
+export default app;

apps/server/src/api/v2/metric_files/[id]/index.ts

@@ -3,6 +3,7 @@ import { standardErrorHandler } from '../../../../utils/response';
 import GET from './GET';
 import DATA from './data/GET';
 import DOWNLOAD from './download/GET';
+import SCREENSHOT from './screenshot';
 import SHARING from './sharing';
 
 const app = new Hono()
@@ -10,6 +11,7 @@ const app = new Hono()
   .route('/data', DATA)
   .route('/download', DOWNLOAD)
   .route('/sharing', SHARING)
+  .route('/screenshot', SCREENSHOT)
   .onError(standardErrorHandler);
 
 export default app;

apps/server/src/api/v2/metric_files/[id]/screenshot/GET.ts

@@ -0,0 +1,57 @@
+import { checkPermission } from '@buster/access-controls';
+import { getMetricFileById } from '@buster/database/queries';
+import { getAssetScreenshotSignedUrl } from '@buster/search';
+import { AssetIdParamsSchema, type GetScreenshotResponse } from '@buster/server-shared/screenshots';
+import { zValidator } from '@hono/zod-validator';
+import { Hono } from 'hono';
+import { HTTPException } from 'hono/http-exception';
+
+const app = new Hono().get('/', zValidator('param', AssetIdParamsSchema), async (c) => {
+  const metricId = c.req.valid('param').id;
+  const user = c.get('busterUser');
+
+  const metric = await getMetricFileById(metricId);
+  if (!metric) {
+    throw new HTTPException(404, { message: 'Metric not found' });
+  }
+
+  const permission = await checkPermission({
+    userId: user.id,
+    assetId: metricId,
+    assetType: 'metric_file',
+    requiredRole: 'can_view',
+    workspaceSharing: metric.workspaceSharing,
+    organizationId: metric.organizationId,
+  });
+
+  if (!permission.hasAccess) {
+    throw new HTTPException(403, {
+      message: 'You do not have permission to view this metric',
+    });
+  }
+
+  let signedUrl = '';
+  let success = true;
+
+  try {
+    signedUrl = await getAssetScreenshotSignedUrl({
+      assetType: 'metric_file',
+      assetId: metricId,
+      organizationId: metric.organizationId,
+    });
+  } catch (error) {
+    console.error('Failed to generate metric screenshot URL', {
+      metricId,
+      error,
+    });
+    success = false;
+  }
+
+  const response: GetScreenshotResponse = {
+    success,
+    url: signedUrl,
+  };
+  return c.json(response);
+});
+
+export default app;

apps/server/src/api/v2/metric_files/[id]/screenshot/PUT.ts

@@ -0,0 +1,57 @@
+import { checkPermission } from '@buster/access-controls';
+import { getMetricFileById } from '@buster/database/queries';
+import {
+  AssetIdParamsSchema,
+  PutScreenshotRequestSchema,
+  type PutScreenshotResponse,
+} from '@buster/server-shared/screenshots';
+import { zValidator } from '@hono/zod-validator';
+import { Hono } from 'hono';
+import { HTTPException } from 'hono/http-exception';
+import { uploadScreenshotHandler } from '../../../../../shared-helpers/upload-screenshot-handler';
+
+const app = new Hono().put(
+  '/',
+  zValidator('json', PutScreenshotRequestSchema),
+  zValidator('param', AssetIdParamsSchema),
+  async (c) => {
+    const assetId = c.req.valid('param').id;
+    const { base64Image } = c.req.valid('json');
+    const user = c.get('busterUser');
+
+    const metric = await getMetricFileById(assetId);
+    if (!metric) {
+      throw new HTTPException(404, { message: 'Metric not found' });
+    }
+
+    const permission = await checkPermission({
+      userId: user.id,
+      assetId,
+      assetType: 'metric_file',
+      requiredRole: 'can_edit',
+      workspaceSharing: metric.workspaceSharing,
+      organizationId: metric.organizationId,
+    });
+
+    if (!permission.hasAccess) {
+      throw new HTTPException(403, {
+        message: 'You do not have permission to upload a screenshot for this metric',
+      });
+    }
+
+    if (!assetId) {
+      throw new HTTPException(404, { message: 'Metric not found' });
+    }
+
+    const result: PutScreenshotResponse = await uploadScreenshotHandler({
+      assetType: 'metric_file',
+      assetId,
+      base64Image,
+      organizationId: metric.organizationId,
+    });
+
+    return c.json(result);
+  }
+);
+
+export default app;

apps/server/src/api/v2/metric_files/[id]/screenshot/index.ts

@@ -0,0 +1,7 @@
+import { Hono } from 'hono';
+import GET from './GET';
+import PUT from './PUT';
+
+const app = new Hono().route('/', GET).route('/', PUT);
+
+export default app;

apps/server/src/api/v2/reports/[id]/index.ts

@@ -1,8 +1,13 @@
 import { Hono } from 'hono';
 import GET from './GET';
 import PUT from './PUT';
+import SCREENSHOT from './screenshot';
 import SHARING from './sharing';
 
-const app = new Hono().route('/', GET).route('/', PUT).route('/sharing', SHARING);
+const app = new Hono()
+  .route('/', GET)
+  .route('/', PUT)
+  .route('/sharing', SHARING)
+  .route('/screenshot', SCREENSHOT);
 
 export default app;

apps/server/src/api/v2/reports/[id]/screenshot/GET.ts

@@ -0,0 +1,58 @@
+import { checkPermission } from '@buster/access-controls';
+import { getReportFileById } from '@buster/database/queries';
+import { getAssetScreenshotSignedUrl } from '@buster/search';
+import { AssetIdParamsSchema, type GetScreenshotResponse } from '@buster/server-shared/screenshots';
+import { zValidator } from '@hono/zod-validator';
+import { Hono } from 'hono';
+import { HTTPException } from 'hono/http-exception';
+
+const app = new Hono().get('/', zValidator('param', AssetIdParamsSchema), async (c) => {
+  const reportId = c.req.valid('param').id;
+  const user = c.get('busterUser');
+
+  const report = await getReportFileById({ reportId, userId: user.id });
+
+  if (!report) {
+    throw new HTTPException(404, { message: 'Report not found' });
+  }
+
+  const permission = await checkPermission({
+    userId: user.id,
+    assetId: reportId,
+    assetType: 'report_file',
+    requiredRole: 'can_view',
+    workspaceSharing: report.workspace_sharing,
+    organizationId: report.organization_id,
+  });
+
+  if (!permission.hasAccess) {
+    throw new HTTPException(403, {
+      message: 'You do not have permission to view this report',
+    });
+  }
+
+  let signedUrl = '';
+  let success = true;
+
+  try {
+    signedUrl = await getAssetScreenshotSignedUrl({
+      assetType: 'report_file',
+      assetId: reportId,
+      organizationId: report.organization_id,
+    });
+  } catch (error) {
+    console.error('Failed to generate report screenshot URL', {
+      reportId,
+      error,
+    });
+    success = false;
+  }
+
+  const response: GetScreenshotResponse = {
+    success,
+    url: signedUrl,
+  };
+  return c.json(response);
+});
+
+export default app;

apps/server/src/api/v2/reports/[id]/screenshot/PUT.ts

@@ -0,0 +1,57 @@
+import { checkPermission } from '@buster/access-controls';
+import { getReportFileById } from '@buster/database/queries';
+import {
+  AssetIdParamsSchema,
+  PutScreenshotRequestSchema,
+  type PutScreenshotResponse,
+} from '@buster/server-shared/screenshots';
+import { zValidator } from '@hono/zod-validator';
+import { Hono } from 'hono';
+import { HTTPException } from 'hono/http-exception';
+import { uploadScreenshotHandler } from '../../../../../shared-helpers/upload-screenshot-handler';
+
+const app = new Hono().put(
+  '/',
+  zValidator('json', PutScreenshotRequestSchema),
+  zValidator('param', AssetIdParamsSchema),
+  async (c) => {
+    const assetId = c.req.valid('param').id;
+    const { base64Image } = c.req.valid('json');
+    const user = c.get('busterUser');
+
+    const report = await getReportFileById({ reportId: assetId, userId: user.id });
+    if (!report) {
+      throw new HTTPException(404, { message: 'Report not found' });
+    }
+
+    if (!report) {
+      throw new HTTPException(404, { message: 'Report not found' });
+    }
+
+    const permission = await checkPermission({
+      userId: user.id,
+      assetId,
+      assetType: 'report_file',
+      requiredRole: 'can_edit',
+      workspaceSharing: report.workspace_sharing,
+      organizationId: report.organization_id,
+    });
+
+    if (!permission.hasAccess) {
+      throw new HTTPException(403, {
+        message: 'You do not have permission to upload a screenshot for this report',
+      });
+    }
+
+    const result: PutScreenshotResponse = await uploadScreenshotHandler({
+      assetType: 'report_file',
+      assetId,
+      base64Image,
+      organizationId: report.organization_id,
+    });
+
+    return c.json(result);
+  }
+);
+
+export default app;

apps/server/src/api/v2/reports/[id]/screenshot/index.ts

@@ -0,0 +1,7 @@
+import { Hono } from 'hono';
+import GET from './GET';
+import PUT from './PUT';
+
+const app = new Hono().route('/', GET).route('/', PUT);
+
+export default app;

apps/server/src/shared-helpers/upload-screenshot-handler.ts

@@ -0,0 +1,89 @@
+import { getProviderForOrganization } from '@buster/data-source';
+import type { AssetType } from '@buster/server-shared/assets';
+import { AssetTypeSchema } from '@buster/server-shared/assets';
+import {
+  PutScreenshotRequestSchema,
+  type PutScreenshotResponse,
+  PutScreenshotResponseSchema,
+} from '@buster/server-shared/screenshots';
+import z from 'zod';
+
+export const UploadScreenshotParamsSchema = PutScreenshotRequestSchema.extend({
+  assetType: AssetTypeSchema,
+  assetId: z.string().uuid('Asset ID must be a valid UUID'),
+  organizationId: z.string().uuid('Organization ID must be a valid UUID'),
+});
+
+export type UploadScreenshotParams = z.infer<typeof UploadScreenshotParamsSchema>;
+
+function getExtensionFromContentType(contentType: string): string {
+  switch (contentType) {
+    case 'image/jpeg':
+    case 'image/jpg':
+      return '.jpg';
+    case 'image/png':
+      return '.png';
+    case 'image/webp':
+      return '.webp';
+    default:
+      return '.png';
+  }
+}
+
+function parseBase64Image(base64Image: string): {
+  buffer: Buffer;
+  contentType: string;
+  extension: string;
+} {
+  const dataUriPattern = /^data:(?<mime>[^;]+);base64,(?<data>.+)$/;
+  const match = base64Image.match(dataUriPattern);
+
+  const contentType = match?.groups?.mime ?? 'image/png';
+  const base64Data = match?.groups?.data ?? base64Image;
+
+  const buffer = Buffer.from(base64Data, 'base64');
+
+  if (buffer.length === 0) {
+    throw new Error('Provided image data is empty');
+  }
+
+  return {
+    buffer,
+    contentType,
+    extension: getExtensionFromContentType(contentType),
+  };
+}
+
+function buildScreenshotKey(
+  assetType: AssetType,
+  assetId: string,
+  extension: string,
+  organizationId: string
+): string {
+  return `screenshots/${organizationId}/${assetType}-${assetId}${extension}`;
+}
+
+export async function uploadScreenshotHandler(
+  params: UploadScreenshotParams
+): Promise<PutScreenshotResponse> {
+  const { assetType, assetId, base64Image, organizationId } =
+    UploadScreenshotParamsSchema.parse(params);
+
+  const { buffer, contentType, extension } = parseBase64Image(base64Image);
+
+  const targetKey = buildScreenshotKey(assetType, assetId, extension, organizationId);
+
+  const provider = await getProviderForOrganization(organizationId);
+  const result = await provider.upload(targetKey, buffer, {
+    contentType,
+  });
+
+  if (!result.success) {
+    throw new Error(result.error ?? 'Failed to upload screenshot');
+  }
+
+  return PutScreenshotResponseSchema.parse({
+    success: true,
+    bucketKey: result.key,
+  });
+}

packages/data-source/src/storage/storage-factory.test.ts

@@ -156,7 +156,7 @@ describe('Storage Factory', () => {
       expect(createR2Provider).toHaveBeenCalledWith({
         provider: 'r2',
         accountId: 'test-account',
-        bucket: 'metric-exports',
+        bucket: 'development',
         accessKeyId: 'test-key',
         secretAccessKey: 'test-secret',
       });
@@ -198,7 +198,7 @@ describe('Storage Factory', () => {
       process.env.R2_ACCOUNT_ID = 'default-account';
       process.env.R2_ACCESS_KEY_ID = 'default-key';
       process.env.R2_SECRET_ACCESS_KEY = 'default-secret';
-      // Intentionally not setting R2_BUCKET so it uses the default 'metric-exports'
+      // Intentionally not setting R2_BUCKET so it uses the default 'development'
     });
 
     afterEach(() => {
@@ -296,7 +296,7 @@ describe('Storage Factory', () => {
       expect(createR2Provider).toHaveBeenCalledWith({
         provider: 'r2',
         accountId: 'default-account',
-        bucket: 'metric-exports',
+        bucket: 'development',
         accessKeyId: 'default-key',
         secretAccessKey: 'default-secret',
       });
@@ -314,7 +314,7 @@ describe('Storage Factory', () => {
       expect(createR2Provider).toHaveBeenCalledWith({
         provider: 'r2',
         accountId: 'default-account',
-        bucket: 'metric-exports',
+        bucket: 'development',
         accessKeyId: 'default-key',
         secretAccessKey: 'default-secret',
       });
@@ -329,7 +329,7 @@ describe('Storage Factory', () => {
       expect(createR2Provider).toHaveBeenCalledWith({
         provider: 'r2',
         accountId: 'default-account',
-        bucket: 'metric-exports',
+        bucket: 'development',
         accessKeyId: 'default-key',
         secretAccessKey: 'default-secret',
       });
@@ -349,7 +349,7 @@ describe('Storage Factory', () => {
       expect(createR2Provider).toHaveBeenCalledWith({
         provider: 'r2',
         accountId: 'default-account',
-        bucket: 'metric-exports',
+        bucket: 'development',
         accessKeyId: 'default-key',
         secretAccessKey: 'default-secret',
       });
@@ -374,7 +374,7 @@ describe('Storage Factory', () => {
       expect(createR2Provider).toHaveBeenCalledWith({
         provider: 'r2',
         accountId: 'default-account',
-        bucket: 'metric-exports',
+        bucket: 'development',
         accessKeyId: 'default-key',
         secretAccessKey: 'default-secret',
       });

packages/data-source/src/storage/storage-factory.ts

@@ -29,7 +29,7 @@ export function getDefaultProvider(): StorageProvider {
   const accountId = process.env.R2_ACCOUNT_ID;
   const accessKeyId = process.env.R2_ACCESS_KEY_ID;
   const secretAccessKey = process.env.R2_SECRET_ACCESS_KEY;
-  const bucket = process.env.R2_BUCKET || 'metric-exports';
+  const bucket = process.env.R2_BUCKET ?? 'development';
 
   if (!accountId || !accessKeyId || !secretAccessKey) {
     throw new Error('Default R2 storage credentials not configured');

packages/database/drizzle/0115_glamorous_penance.sql

@@ -0,0 +1,2 @@
+CREATE INDEX "idx_perm_active_identity_asset" ON "asset_permissions" USING btree ("identity_type","identity_id","asset_type","asset_id") WHERE "asset_permissions"."deleted_at" is null;--> statement-breakpoint
+CREATE INDEX "idx_as2_active_by_org" ON "asset_search_v2" USING btree ("organization_id") WHERE "asset_search_v2"."deleted_at" is null;

packages/database/drizzle/meta/_journal.json

@@ -799,6 +799,13 @@
       "when": 1759256209020,
       "tag": "0114_lovely_risque",
       "breakpoints": true
+    },
+    {
+      "idx": 115,
+      "version": "7",
+      "when": 1759336097951,
+      "tag": "0115_glamorous_penance",
+      "breakpoints": true
     }
   ]
 }

packages/database/src/queries/assets/asset-ancestors.ts

@@ -1,4 +1,4 @@
-import { and, eq, isNull, sql } from 'drizzle-orm';
+import { and, eq, inArray, isNull, sql } from 'drizzle-orm';
 import { db } from '../../connection';
 import {
   chats,
@@ -11,11 +11,30 @@ import {
   metricFilesToReportFiles,
   reportFiles,
 } from '../../schema';
-import type { Ancestor, AssetAncestors } from '../../schema-types';
+import type { Ancestor, AssetAncestors, AssetType } from '../../schema-types';
 
 // Type for database transaction
 type DatabaseTransaction = Parameters<Parameters<typeof db.transaction>[0]>[0];
 
+type BatchedAsset = {
+  assetId: string;
+  assetType: AssetType;
+};
+
+type BatchedAncestorRow = {
+  assetId: string;
+  ancestorId: string;
+  ancestorTitle: string;
+  ancestorType: 'chat' | 'collection' | 'dashboard_file' | 'report_file';
+};
+
+const ancestorTypeToKey: Record<BatchedAncestorRow['ancestorType'], keyof AssetAncestors> = {
+  chat: 'chats',
+  collection: 'collections',
+  dashboard_file: 'dashboards',
+  report_file: 'reports',
+};
+
 /**
  * Get chat ancestors as a subquery
  */
@@ -165,6 +184,160 @@ export async function getAssetAncestors(
   return ancestors;
 }
 
+export type GetAssetAncestorsForAssetsInput = {
+  assets: BatchedAsset[];
+  userId: string;
+  organizationId: string;
+};
+
+export async function getAssetAncestorsForAssets(
+  input: GetAssetAncestorsForAssetsInput
+): Promise<Record<string, AssetAncestors>> {
+  const { assets } = input;
+  if (assets.length === 0) {
+    return {};
+  }
+
+  const assetIds: string[] = [];
+  const metricAssetIds: string[] = [];
+
+  for (const asset of assets) {
+    assetIds.push(asset.assetId);
+    if (asset.assetType === 'metric_file') {
+      metricAssetIds.push(asset.assetId);
+    }
+  }
+
+  const chatRowsPromise: Promise<BatchedAncestorRow[]> = assetIds.length
+    ? db
+        .select({
+          assetId: messagesToFiles.fileId,
+          ancestorId: chats.id,
+          ancestorTitle: chats.title,
+          ancestorType: sql<'chat'>`'chat'`.as('ancestorType'),
+        })
+        .from(messagesToFiles)
+        .innerJoin(messages, eq(messages.id, messagesToFiles.messageId))
+        .innerJoin(chats, eq(chats.id, messages.chatId))
+        .where(
+          and(
+            inArray(messagesToFiles.fileId, assetIds),
+            isNull(messagesToFiles.deletedAt),
+            isNull(messages.deletedAt),
+            isNull(chats.deletedAt)
+          )
+        )
+    : // .then((rows) =>
+      //   rows.map((row) => ({
+      //     assetId: row.assetId,
+      //     ancestorId: row.ancestorId,
+      //     ancestorTitle: row.ancestorTitle,
+      //     ancestorType: 'chat',
+      //   }))
+      // )
+      Promise.resolve([]);
+
+  const collectionRowsPromise: Promise<BatchedAncestorRow[]> = assetIds.length
+    ? db
+        .select({
+          assetId: collectionsToAssets.assetId,
+          ancestorId: collections.id,
+          ancestorTitle: collections.name,
+          ancestorType: sql<'collection'>`'collection'`.as('ancestorType'),
+        })
+        .from(collectionsToAssets)
+        .innerJoin(collections, eq(collections.id, collectionsToAssets.collectionId))
+        .where(
+          and(
+            inArray(collectionsToAssets.assetId, assetIds),
+            isNull(collectionsToAssets.deletedAt),
+            isNull(collections.deletedAt)
+          )
+        )
+    : Promise.resolve([]);
+
+  const dashboardRowsPromise: Promise<BatchedAncestorRow[]> = metricAssetIds.length
+    ? db
+        .select({
+          assetId: metricFilesToDashboardFiles.metricFileId,
+          ancestorId: dashboardFiles.id,
+          ancestorTitle: dashboardFiles.name,
+          ancestorType: sql<'dashboard_file'>`'dashboard_file'`.as('ancestorType'),
+        })
+        .from(metricFilesToDashboardFiles)
+        .innerJoin(
+          dashboardFiles,
+          eq(dashboardFiles.id, metricFilesToDashboardFiles.dashboardFileId)
+        )
+        .where(
+          and(
+            inArray(metricFilesToDashboardFiles.metricFileId, metricAssetIds),
+            isNull(metricFilesToDashboardFiles.deletedAt),
+            isNull(dashboardFiles.deletedAt)
+          )
+        )
+    : Promise.resolve([]);
+
+  const reportRowsPromise: Promise<BatchedAncestorRow[]> = metricAssetIds.length
+    ? db
+        .select({
+          assetId: metricFilesToReportFiles.metricFileId,
+          ancestorId: reportFiles.id,
+          ancestorTitle: reportFiles.name,
+          ancestorType: sql<'report_file'>`'report_file'`.as('ancestorType'),
+        })
+        .from(metricFilesToReportFiles)
+        .innerJoin(reportFiles, eq(reportFiles.id, metricFilesToReportFiles.reportFileId))
+        .where(
+          and(
+            inArray(metricFilesToReportFiles.metricFileId, metricAssetIds),
+            isNull(metricFilesToReportFiles.deletedAt),
+            isNull(reportFiles.deletedAt)
+          )
+        )
+    : Promise.resolve([]);
+
+  const [chatRows, collectionRows, dashboardRows, reportRows] = await Promise.all([
+    chatRowsPromise,
+    collectionRowsPromise,
+    dashboardRowsPromise,
+    reportRowsPromise,
+  ]);
+
+  const ancestorsByAssetId: Record<string, AssetAncestors> = {};
+  for (const asset of assets) {
+    ancestorsByAssetId[asset.assetId] = {
+      chats: [],
+      collections: [],
+      dashboards: [],
+      reports: [],
+    };
+  }
+
+  const appendRows = (rows: BatchedAncestorRow[]) => {
+    for (const row of rows) {
+      const ancestorBucket = ancestorsByAssetId[row.assetId];
+      if (!ancestorBucket) {
+        continue;
+      }
+
+      const key = ancestorTypeToKey[row.ancestorType];
+      const ancestor: Ancestor = {
+        id: row.ancestorId,
+        title: row.ancestorTitle,
+      };
+      ancestorBucket[key].push(ancestor);
+    }
+  };
+
+  appendRows(chatRows);
+  appendRows(collectionRows);
+  appendRows(dashboardRows);
+  appendRows(reportRows);
+
+  return ancestorsByAssetId;
+}
+
 export async function getAssetChatAncestors(
   assetId: string,
   tx?: DatabaseTransaction

packages/database/src/queries/assets/index.ts

@@ -23,9 +23,12 @@ export {
   getMetricDashboardAncestors,
   getMetricReportAncestors,
   getAssetAncestors,
+  getAssetAncestorsForAssets,
   getAssetAncestorsWithTransaction,
 } from './asset-ancestors';
 
+export type { GetAssetAncestorsForAssetsInput } from './asset-ancestors';
+
 export {
   getAssetLatestVersion,
   GetAssetLatestVersionInputSchema,

packages/database/src/schema.ts

@@ -1483,6 +1483,9 @@ export const assetPermissions = pgTable(
     index('idx_perm_active_asset_identity')
       .on(table.assetId, table.assetType, table.identityId, table.identityType)
       .where(isNull(table.deletedAt)),
+    index('idx_perm_active_identity_asset')
+      .on(table.identityType, table.identityId, table.assetType, table.assetId)
+      .where(isNull(table.deletedAt)),
   ]
 );
 
@@ -1913,6 +1916,7 @@ export const assetSearchV2 = pgTable(
     index('idx_as2_active_by_asset')
       .on(table.assetId, table.assetType)
       .where(isNull(table.deletedAt)),
+    index('idx_as2_active_by_org').on(table.organizationId).where(isNull(table.deletedAt)),
   ]
 );
 

packages/search/src/asset-search/get-asset-screenshot.ts

@@ -0,0 +1,57 @@
+import { type StorageProvider, getProviderForOrganization } from '@buster/data-source';
+import { type AssetType, AssetTypeSchema } from '@buster/server-shared/assets';
+import { z } from 'zod';
+
+const GetAssetScreenshotParamsSchema = z.object({
+  assetType: AssetTypeSchema,
+  assetId: z.string().uuid(),
+  organizationId: z.string().uuid(),
+  expiresIn: z.number().int().positive().optional(),
+});
+
+export type GetAssetScreenshotParams = z.infer<typeof GetAssetScreenshotParamsSchema>;
+
+async function resolveBucketKey(
+  provider: StorageProvider,
+  assetType: AssetType,
+  assetId: string,
+  organizationId: string
+): Promise<string | null> {
+  const baseKey = `screenshots/${organizationId}/${assetType}-${assetId}`;
+
+  try {
+    const objects = await provider.list(baseKey, { maxKeys: 5 });
+    if (objects.length > 0) {
+      const sorted = [...objects].sort((a, b) => {
+        const aTime = a.lastModified?.getTime?.() ?? 0;
+        const bTime = b.lastModified?.getTime?.() ?? 0;
+        return bTime - aTime;
+      });
+      return sorted[0]?.key ?? null;
+    }
+  } catch (error) {
+    console.error('Failed to list screenshot objects', {
+      baseKey,
+      assetType,
+      assetId,
+      error,
+    });
+  }
+
+  return null;
+}
+
+export async function getAssetScreenshotSignedUrl(
+  params: Readonly<GetAssetScreenshotParams>
+): Promise<string> {
+  const { assetType, assetId, expiresIn, organizationId } =
+    GetAssetScreenshotParamsSchema.parse(params);
+  const provider = await getProviderForOrganization(organizationId);
+  const resolvedKey = await resolveBucketKey(provider, assetType, assetId, organizationId);
+
+  if (!resolvedKey) {
+    throw new Error('Screenshot not found for asset');
+  }
+
+  return provider.getSignedUrl(resolvedKey, expiresIn ?? 900);
+}

packages/search/src/asset-search/index.ts

@@ -1 +1,2 @@
 export * from './search';
+export * from './get-asset-screenshot';

packages/search/src/asset-search/search.test.ts

@@ -6,7 +6,7 @@ import { performTextSearch } from './search';
 vi.mock('@buster/database/queries', () => ({
   getUserOrganizationId: vi.fn(),
   searchText: vi.fn(),
-  getAssetAncestors: vi.fn(),
+  getAssetAncestorsForAssets: vi.fn(),
 }));
 
 vi.mock('./text-processing-helpers', () => ({
@@ -14,7 +14,11 @@ vi.mock('./text-processing-helpers', () => ({
 }));
 
 // Import the mocked functions
-import { getAssetAncestors, getUserOrganizationId, searchText } from '@buster/database/queries';
+import {
+  getAssetAncestorsForAssets,
+  getUserOrganizationId,
+  searchText,
+} from '@buster/database/queries';
 import { processSearchResultText } from './text-processing-helpers';
 
 describe('search.ts - Unit Tests', () => {
@@ -59,6 +63,11 @@ describe('search.ts - Unit Tests', () => {
     reports: [],
   };
 
+  const mockAncestorsForAssets = {
+    'asset-1': mockAncestors,
+    'asset-2': mockAncestors,
+  };
+
   beforeEach(() => {
     vi.clearAllMocks();
     (getUserOrganizationId as Mock).mockResolvedValue(mockUserOrg);
@@ -67,7 +76,7 @@ describe('search.ts - Unit Tests', () => {
       processedTitle: `<b>${title}</b>`,
       processedAdditionalText: `<b>${additionalText}</b>`,
     }));
-    (getAssetAncestors as Mock).mockResolvedValue(mockAncestors);
+    (getAssetAncestorsForAssets as Mock).mockResolvedValue(mockAncestorsForAssets);
   });
 
   describe('performTextSearch', () => {
@@ -222,19 +231,15 @@ describe('search.ts - Unit Tests', () => {
 
       const result = await performTextSearch(mockUserId, searchRequestWithAncestors);
 
-      expect(getAssetAncestors).toHaveBeenCalledTimes(2);
-      expect(getAssetAncestors).toHaveBeenCalledWith(
-        'asset-1',
-        'chat',
-        mockUserId,
-        mockOrganizationId
-      );
-      expect(getAssetAncestors).toHaveBeenCalledWith(
-        'asset-2',
-        'metric_file',
-        mockUserId,
-        mockOrganizationId
-      );
+      expect(getAssetAncestorsForAssets).toHaveBeenCalledTimes(1);
+      expect(getAssetAncestorsForAssets).toHaveBeenCalledWith({
+        assets: [
+          { assetId: 'asset-1', assetType: 'chat' },
+          { assetId: 'asset-2', assetType: 'metric_file' },
+        ],
+        userId: mockUserId,
+        organizationId: mockOrganizationId,
+      });
 
       expect(result.data[0]).toHaveProperty('ancestors', mockAncestors);
       expect(result.data[1]).toHaveProperty('ancestors', mockAncestors);
@@ -243,7 +248,7 @@ describe('search.ts - Unit Tests', () => {
     it('should not include ancestors when not requested', async () => {
       const result = await performTextSearch(mockUserId, basicSearchRequest);
 
-      expect(getAssetAncestors).not.toHaveBeenCalled();
+      expect(getAssetAncestorsForAssets).not.toHaveBeenCalled();
       expect(result.data[0]).not.toHaveProperty('ancestors');
       expect(result.data[1]).not.toHaveProperty('ancestors');
     });
@@ -264,7 +269,7 @@ describe('search.ts - Unit Tests', () => {
 
       expect(result).toEqual(emptySearchResponse);
       expect(processSearchResultText).not.toHaveBeenCalled();
-      expect(getAssetAncestors).not.toHaveBeenCalled();
+      expect(getAssetAncestorsForAssets).not.toHaveBeenCalled();
     });
 
     it('should handle null/undefined additional text', async () => {
@@ -342,6 +347,12 @@ describe('search.ts - Unit Tests', () => {
         },
       });
 
+      // Set up mock ancestors for all 12 assets
+      const manyAncestorsForAssets = Object.fromEntries(
+        Array.from({ length: 12 }, (_, i) => [`asset-${i + 1}`, mockAncestors])
+      );
+      (getAssetAncestorsForAssets as Mock).mockResolvedValue(manyAncestorsForAssets);
+
       const searchRequestWithAncestors: SearchTextRequest = {
         ...basicSearchRequest,
         includeAssetAncestors: true,
@@ -349,8 +360,8 @@ describe('search.ts - Unit Tests', () => {
 
       const result = await performTextSearch(mockUserId, searchRequestWithAncestors);
 
-      // Should call getAssetAncestors for each result
-      expect(getAssetAncestors).toHaveBeenCalledTimes(12);
+      // Should call getAssetAncestorsForAssets once for the batch
+      expect(getAssetAncestorsForAssets).toHaveBeenCalledTimes(1);
 
       // Results should have ancestors added
       expect(result.data).toHaveLength(12);
@@ -365,7 +376,7 @@ describe('search.ts - Unit Tests', () => {
         includeAssetAncestors: true,
       };
 
-      (getAssetAncestors as Mock).mockRejectedValue(new Error('Ancestor lookup failed'));
+      (getAssetAncestorsForAssets as Mock).mockRejectedValue(new Error('Ancestor lookup failed'));
 
       await expect(performTextSearch(mockUserId, searchRequestWithAncestors)).rejects.toThrow(
         'Ancestor lookup failed'

packages/search/src/asset-search/search.ts

@@ -1,7 +1,6 @@
 import {
   type SearchFilters,
-  getAssetAncestors,
-  getAssetAncestorsWithTransaction,
+  getAssetAncestorsForAssets,
   getUserOrganizationId,
   searchText,
 } from '@buster/database/queries';
@@ -124,7 +123,7 @@ async function addAncestorsToSearchResults(
   userId: string,
   organizationId: string
 ): Promise<SearchTextData[]> {
-  const chunkSize = 25;
+  const chunkSize = 50;
   const resultsWithAncestors: SearchTextData[] = [];
   const totalChunks = Math.ceil(searchResults.length / chunkSize);
 
@@ -134,21 +133,24 @@ async function addAncestorsToSearchResults(
 
   for (let i = 0; i < searchResults.length; i += chunkSize) {
     const chunk = searchResults.slice(i, i + chunkSize);
-    const chunkResults = await Promise.all(
-      chunk.map(async (searchResult) => {
-        const ancestors = await getAssetAncestors(
-          searchResult.assetId,
-          searchResult.assetType as AssetType,
-          userId,
-          organizationId
-        );
+    const ancestorsByAssetId = await getAssetAncestorsForAssets({
+      assets: chunk.map((searchResult) => ({
+        assetId: searchResult.assetId,
+        assetType: searchResult.assetType as AssetType,
+      })),
+      userId,
+      organizationId,
+    });
 
-        return {
-          ...searchResult,
-          ancestors,
-        };
-      })
-    );
+    const chunkResults = chunk.map((searchResult) => ({
+      ...searchResult,
+      ancestors: ancestorsByAssetId[searchResult.assetId] ?? {
+        chats: [],
+        collections: [],
+        dashboards: [],
+        reports: [],
+      },
+    }));
     resultsWithAncestors.push(...chunkResults);
   }
 

packages/server-shared/package.json

@@ -103,6 +103,10 @@
     "./search": {
       "types": "./dist/search/index.d.ts",
       "default": "./dist/search/index.js"
+    },
+    "./screenshots": {
+      "types": "./dist/screenshots/index.d.ts",
+      "default": "./dist/screenshots/index.js"
     }
   },
   "module": "src/index.ts",

packages/server-shared/src/index.ts

@@ -22,6 +22,7 @@ export * from './metrics';
 export * from './organization';
 export * from './public-chat';
 export * from './s3-integrations';
+export * from './screenshots';
 export * from './search';
 export * from './security';
 // Export share module (has some naming conflicts with chats and metrics)

packages/server-shared/src/screenshots/index.ts

@@ -0,0 +1 @@
+export * from './screenshots';

packages/server-shared/src/screenshots/screenshots.ts

@@ -0,0 +1,25 @@
+import { z } from 'zod';
+
+export const AssetIdParamsSchema = z.object({
+  id: z.string().uuid('Asset ID must be a valid UUID'),
+});
+
+export const PutScreenshotRequestSchema = z.object({
+  base64Image: z.string().min(1, 'Base64 image is required'),
+});
+
+export type PutScreenshotRequest = z.infer<typeof PutScreenshotRequestSchema>;
+
+export const PutScreenshotResponseSchema = z.object({
+  success: z.literal(true),
+  bucketKey: z.string().min(1, 'Bucket key is required'),
+});
+
+export type PutScreenshotResponse = z.infer<typeof PutScreenshotResponseSchema>;
+
+export const GetScreenshotResponseSchema = z.object({
+  success: z.boolean(),
+  url: z.string().url(),
+});
+
+export type GetScreenshotResponse = z.infer<typeof GetScreenshotResponseSchema>;