Merge pull request #1239 from buster-so/wells-bus-1994-take-screenshot-api

Search improvements and adding signed url to the search

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

@@ -15,6 +15,14 @@ const app = new Hono().get('/', zValidator('param', AssetIdParamsSchema), async
     throw new HTTPException(404, { message: 'Chat not found' });
   }
 
+  if (!chat.screenshotBucketKey) {
+    const result: GetScreenshotResponse = {
+      success: false,
+      error: 'Screenshot not found',
+    };
+    return c.json(result);
+  }
+
   const permission = await checkPermission({
     userId: user.id,
     assetId: id,
@@ -30,28 +38,27 @@ const app = new Hono().get('/', zValidator('param', AssetIdParamsSchema), async
     });
   }
 
-  let signedUrl = '';
-  let success = true;
-
   try {
-    signedUrl = await getAssetScreenshotSignedUrl({
-      assetType: 'chat',
-      assetId: id,
+    const signedUrl = await getAssetScreenshotSignedUrl({
+      key: chat.screenshotBucketKey,
       organizationId: chat.organizationId,
     });
+    const result: GetScreenshotResponse = {
+      success: true,
+      url: signedUrl,
+    };
+    return c.json(result);
   } catch (error) {
     console.error('Failed to generate chat screenshot URL', {
       chatId: id,
       error,
     });
-    success = false;
+    const result: GetScreenshotResponse = {
+      success: false,
+      error: 'Failed to generate screenshot URL',
+    };
+    return c.json(result);
   }
-
-  const response: GetScreenshotResponse = {
-    success,
-    url: signedUrl,
-  };
-  return c.json(response);
 });
 
 export default app;

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

@@ -15,6 +15,14 @@ const app = new Hono().get('/', zValidator('param', AssetIdParamsSchema), async
     throw new HTTPException(404, { message: 'Dashboard not found' });
   }
 
+  if (!dashboard.screenshotBucketKey) {
+    const result: GetScreenshotResponse = {
+      success: false,
+      error: 'Screenshot not found',
+    };
+    return c.json(result);
+  }
+
   const permission = await checkPermission({
     userId: user.id,
     assetId: dashboardId,
@@ -30,28 +38,27 @@ const app = new Hono().get('/', zValidator('param', AssetIdParamsSchema), async
     });
   }
 
-  let signedUrl = '';
-  let success = true;
-
   try {
-    signedUrl = await getAssetScreenshotSignedUrl({
-      assetType: 'dashboard_file',
-      assetId: dashboardId,
+    const signedUrl = await getAssetScreenshotSignedUrl({
+      key: dashboard.screenshotBucketKey,
       organizationId: dashboard.organizationId,
     });
+    const result: GetScreenshotResponse = {
+      success: true,
+      url: signedUrl,
+    };
+    return c.json(result);
   } catch (error) {
     console.error('Failed to generate dashboard screenshot URL', {
       dashboardId,
       error,
     });
-    success = false;
+    const result: GetScreenshotResponse = {
+      success: false,
+      error: 'Failed to generate screenshot URL',
+    };
+    return c.json(result);
   }
-
-  const response: GetScreenshotResponse = {
-    success,
-    url: signedUrl,
-  };
-  return c.json(response);
 });
 
 export default app;

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

@@ -15,6 +15,14 @@ const app = new Hono().get('/', zValidator('param', AssetIdParamsSchema), async
     throw new HTTPException(404, { message: 'Metric not found' });
   }
 
+  if (!metric.screenshotBucketKey) {
+    const result: GetScreenshotResponse = {
+      success: false,
+      error: 'Screenshot not found',
+    };
+    return c.json(result);
+  }
+
   const permission = await checkPermission({
     userId: user.id,
     assetId: metricId,
@@ -30,28 +38,27 @@ const app = new Hono().get('/', zValidator('param', AssetIdParamsSchema), async
     });
   }
 
-  let signedUrl = '';
-  let success = true;
-
   try {
-    signedUrl = await getAssetScreenshotSignedUrl({
-      assetType: 'metric_file',
-      assetId: metricId,
+    const signedUrl = await getAssetScreenshotSignedUrl({
+      key: metric.screenshotBucketKey,
       organizationId: metric.organizationId,
     });
+    const result: GetScreenshotResponse = {
+      success: true,
+      url: signedUrl,
+    };
+    return c.json(result);
   } catch (error) {
     console.error('Failed to generate metric screenshot URL', {
       metricId,
       error,
     });
-    success = false;
+    const result: GetScreenshotResponse = {
+      success: false,
+      error: 'Failed to generate screenshot URL',
+    };
+    return c.json(result);
   }
-
-  const response: GetScreenshotResponse = {
-    success,
-    url: signedUrl,
-  };
-  return c.json(response);
 });
 
 export default app;

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

@@ -1,5 +1,5 @@
 import { checkPermission } from '@buster/access-controls';
-import { getReportFileById } from '@buster/database/queries';
+import { getAssetScreenshotBucketKey, 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';
@@ -16,6 +16,19 @@ const app = new Hono().get('/', zValidator('param', AssetIdParamsSchema), async
     throw new HTTPException(404, { message: 'Report not found' });
   }
 
+  const existingKey = await getAssetScreenshotBucketKey({
+    assetType: 'report_file',
+    assetId: reportId,
+  });
+
+  if (!existingKey) {
+    const result: GetScreenshotResponse = {
+      success: false,
+      error: 'Screenshot not found',
+    };
+    return c.json(result);
+  }
+
   const permission = await checkPermission({
     userId: user.id,
     assetId: reportId,
@@ -31,26 +44,27 @@ const app = new Hono().get('/', zValidator('param', AssetIdParamsSchema), async
     });
   }
 
-  let signedUrl = '';
-  let success = true;
-
   try {
-    signedUrl = await getAssetScreenshotSignedUrl({
-      assetType: 'report_file',
-      assetId: reportId,
+    const signedUrl = await getAssetScreenshotSignedUrl({
+      key: existingKey,
       organizationId: report.organization_id,
     });
+
+    const result: GetScreenshotResponse = {
+      success: true,
+      url: signedUrl,
+    };
+    return c.json(result);
   } catch (error) {
     console.error('Failed to generate report screenshot URL', {
       reportId,
       error,
     });
-    success = false;
   }
 
   const response: GetScreenshotResponse = {
-    success,
-    url: signedUrl,
+    success: false,
+    error: 'Failed to generate screenshot URL',
   };
   return c.json(response);
 });

apps/server/src/api/v2/security/update-workspace-settings.test.ts

@@ -28,6 +28,14 @@ vi.mock('@buster/database/connection', () => ({
 
 vi.mock('@buster/database/schema', () => ({
   organizations: {},
+  chats: {},
+  collections: {},
+  dashboardFiles: {},
+  reportFiles: {},
+  metricFiles: {},
+  users: {},
+  datasets: {},
+  dataSources: {},
   eq: vi.fn(),
   and: vi.fn(),
   isNull: vi.fn(),

apps/server/src/api/v2/slack/services/slack-oauth-service.test.ts

@@ -26,6 +26,12 @@ vi.mock('@buster/database/schema', () => ({
   permissionGroups: {},
   users: {},
   slackIntegrations: {},
+  chats: {},
+  collections: {},
+  dashboardFiles: {},
+  reportFiles: {},
+  metricFiles: {},
+  dataSources: {},
 }));
 
 import { beforeEach, describe, expect, it, vi } from 'vitest';

apps/server/src/shared-helpers/metric-helpers.test.ts

@@ -100,6 +100,7 @@ describe('metric-helpers', () => {
     workspaceSharingEnabledAt: null,
     workspaceSharingEnabledBy: null,
     deletedAt: null,
+    screenshotBucketKey: null,
     ...overrides,
   });
 

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

@@ -1,4 +1,5 @@
 import { getProviderForOrganization } from '@buster/data-source';
+import { updateAssetScreenshotBucketKey } from '@buster/database/queries';
 import type { AssetType } from '@buster/server-shared/assets';
 import { AssetTypeSchema } from '@buster/server-shared/assets';
 import {
@@ -82,6 +83,12 @@ export async function uploadScreenshotHandler(
     throw new Error(result.error ?? 'Failed to upload screenshot');
   }
 
+  await updateAssetScreenshotBucketKey({
+    assetId,
+    assetType,
+    screenshotBucketKey: result.key,
+  });
+
   return PutScreenshotResponseSchema.parse({
     success: true,
     bucketKey: result.key,

packages/database/drizzle/0116_cuddly_thaddeus_ross.sql

@@ -0,0 +1,6 @@
+ALTER TABLE "asset_search_v2" ADD COLUMN "screenshot_bucket_key" text;--> statement-breakpoint
+ALTER TABLE "chats" ADD COLUMN "screenshot_bucket_key" text;--> statement-breakpoint
+ALTER TABLE "collections" ADD COLUMN "screenshot_bucket_key" text;--> statement-breakpoint
+ALTER TABLE "dashboard_files" ADD COLUMN "screenshot_bucket_key" text;--> statement-breakpoint
+ALTER TABLE "metric_files" ADD COLUMN "screenshot_bucket_key" text;--> statement-breakpoint
+ALTER TABLE "report_files" ADD COLUMN "screenshot_bucket_key" text;

packages/database/drizzle/0117_careful_alex_power.sql

@@ -0,0 +1,162 @@
+-- Custom SQL migration file, put your code below! --
+
+-- Drop existing triggers to recreate them with screenshot_bucket_key support
+DROP TRIGGER IF EXISTS sync_chats_text_search ON chats;--> statement-breakpoint
+DROP TRIGGER IF EXISTS sync_metric_files_text_search ON metric_files;--> statement-breakpoint
+DROP TRIGGER IF EXISTS sync_dashboard_files_text_search ON dashboard_files;--> statement-breakpoint
+DROP TRIGGER IF EXISTS sync_report_files_text_search ON report_files;--> statement-breakpoint
+
+-- Drop existing trigger functions to recreate them with screenshot_bucket_key support
+DROP FUNCTION IF EXISTS sync_chats_to_text_search();--> statement-breakpoint
+DROP FUNCTION IF EXISTS sync_metric_files_to_text_search();--> statement-breakpoint
+DROP FUNCTION IF EXISTS sync_dashboard_files_to_text_search();--> statement-breakpoint
+DROP FUNCTION IF EXISTS sync_report_files_to_text_search();--> statement-breakpoint
+
+-- Function for chats table with screenshot_bucket_key support
+CREATE OR REPLACE FUNCTION sync_chats_to_text_search()
+RETURNS TRIGGER AS $$
+DECLARE
+    messages_text text;
+BEGIN
+    IF TG_OP = 'INSERT' OR TG_OP = 'UPDATE' THEN
+        -- Get all request messages for this chat
+        SELECT string_agg(COALESCE(request_message, ''), E'\n' ORDER BY created_at)
+        INTO messages_text
+        FROM public.messages
+        WHERE chat_id = NEW.id AND deleted_at IS NULL AND request_message IS NOT NULL AND request_message != '';
+        
+        INSERT INTO public.asset_search_v2 (
+            asset_id, asset_type, title, additional_text, organization_id,
+            created_at, updated_at, deleted_at, screenshot_bucket_key
+        )
+        VALUES (
+            NEW.id, 'chat', COALESCE(NEW.title, ''), messages_text,
+            NEW.organization_id,
+            NEW.created_at, NEW.updated_at, NEW.deleted_at, NEW.screenshot_bucket_key
+        )
+        ON CONFLICT (asset_id, asset_type) DO UPDATE SET
+            title = EXCLUDED.title,
+            additional_text = EXCLUDED.additional_text,
+            updated_at = EXCLUDED.updated_at,
+            deleted_at = EXCLUDED.deleted_at,
+            screenshot_bucket_key = EXCLUDED.screenshot_bucket_key;
+    ELSIF TG_OP = 'DELETE' THEN
+        UPDATE public.asset_search_v2
+        SET deleted_at = NOW()
+        WHERE asset_id = OLD.id AND asset_type = 'chat';
+    END IF;
+    RETURN NULL;
+END;
+$$ LANGUAGE plpgsql;
+
+-- Function for metric_files table with screenshot_bucket_key support
+CREATE OR REPLACE FUNCTION sync_metric_files_to_text_search()
+RETURNS TRIGGER AS $$
+BEGIN
+    IF TG_OP = 'INSERT' OR TG_OP = 'UPDATE' THEN
+        INSERT INTO public.asset_search_v2 (
+            asset_id, asset_type, title, additional_text, organization_id,
+            created_at, updated_at, deleted_at, screenshot_bucket_key
+        )
+        VALUES (
+            NEW.id, 'metric_file', COALESCE(NEW.name, ''), 
+            COALESCE(NEW.content ->> 'description', ''),
+            NEW.organization_id,
+            NEW.created_at, NEW.updated_at, NEW.deleted_at, NEW.screenshot_bucket_key
+        )
+        ON CONFLICT (asset_id, asset_type) DO UPDATE SET
+            title = EXCLUDED.title,
+            additional_text = EXCLUDED.additional_text,
+            updated_at = EXCLUDED.updated_at,
+            deleted_at = EXCLUDED.deleted_at,
+            screenshot_bucket_key = EXCLUDED.screenshot_bucket_key;
+    ELSIF TG_OP = 'DELETE' THEN
+        UPDATE public.asset_search_v2
+        SET deleted_at = NOW()
+        WHERE asset_id = OLD.id AND asset_type = 'metric_file';
+    END IF;
+    RETURN NULL;
+END;
+$$ LANGUAGE plpgsql;
+
+-- Function for dashboard_files table with screenshot_bucket_key support
+CREATE OR REPLACE FUNCTION sync_dashboard_files_to_text_search()
+RETURNS TRIGGER AS $$
+BEGIN
+    IF TG_OP = 'INSERT' OR TG_OP = 'UPDATE' THEN
+        INSERT INTO public.asset_search_v2 (
+            asset_id, asset_type, title, additional_text, organization_id,
+            created_at, updated_at, deleted_at, screenshot_bucket_key
+        )
+        VALUES (
+            NEW.id, 'dashboard_file', COALESCE(NEW.name, ''), 
+            COALESCE(NEW.content ->> 'description', ''),
+            NEW.organization_id,
+            NEW.created_at, NEW.updated_at, NEW.deleted_at, NEW.screenshot_bucket_key
+        )
+        ON CONFLICT (asset_id, asset_type) DO UPDATE SET
+            title = EXCLUDED.title,
+            additional_text = EXCLUDED.additional_text,
+            updated_at = EXCLUDED.updated_at,
+            deleted_at = EXCLUDED.deleted_at,
+            screenshot_bucket_key = EXCLUDED.screenshot_bucket_key;
+    ELSIF TG_OP = 'DELETE' THEN
+        UPDATE public.asset_search_v2
+        SET deleted_at = NOW()
+        WHERE asset_id = OLD.id AND asset_type = 'dashboard_file';
+    END IF;
+    RETURN NULL;
+END;
+$$ LANGUAGE plpgsql;
+
+-- Function for report_files table with screenshot_bucket_key support
+CREATE OR REPLACE FUNCTION sync_report_files_to_text_search()
+RETURNS TRIGGER AS $$
+DECLARE
+    cleaned_content text;
+BEGIN
+    IF TG_OP = 'INSERT' OR TG_OP = 'UPDATE' THEN
+        -- Remove <metric> tags (both self-closing and opening/closing pairs) and newlines
+        cleaned_content := regexp_replace(COALESCE(NEW.content, ''), '<metric[^>]*(?:/>|>.*?</metric>)', '', 'g');
+        cleaned_content := regexp_replace(cleaned_content, '\n', '', 'g');
+        
+        INSERT INTO public.asset_search_v2 (
+            asset_id, asset_type, title, additional_text, organization_id,
+            created_at, updated_at, deleted_at, screenshot_bucket_key
+        )
+        VALUES (
+            NEW.id, 'report_file', COALESCE(NEW.name, ''), cleaned_content,
+            NEW.organization_id,
+            NEW.created_at, NEW.updated_at, NEW.deleted_at, NEW.screenshot_bucket_key
+        )
+        ON CONFLICT (asset_id, asset_type) DO UPDATE SET
+            title = EXCLUDED.title,
+            additional_text = EXCLUDED.additional_text,
+            updated_at = EXCLUDED.updated_at,
+            deleted_at = EXCLUDED.deleted_at,
+            screenshot_bucket_key = EXCLUDED.screenshot_bucket_key;
+    ELSIF TG_OP = 'DELETE' THEN
+        UPDATE public.asset_search_v2
+        SET deleted_at = NOW()
+        WHERE asset_id = OLD.id AND asset_type = 'report_file';
+    END IF;
+    RETURN NULL;
+END;
+$$ LANGUAGE plpgsql;
+
+-- Recreate triggers for each table
+CREATE TRIGGER sync_chats_text_search
+AFTER INSERT OR UPDATE OR DELETE ON chats
+FOR EACH ROW EXECUTE FUNCTION sync_chats_to_text_search();
+
+CREATE TRIGGER sync_metric_files_text_search
+AFTER INSERT OR UPDATE OR DELETE ON metric_files
+FOR EACH ROW EXECUTE FUNCTION sync_metric_files_to_text_search();
+
+CREATE TRIGGER sync_dashboard_files_text_search
+AFTER INSERT OR UPDATE OR DELETE ON dashboard_files
+FOR EACH ROW EXECUTE FUNCTION sync_dashboard_files_to_text_search();
+
+CREATE TRIGGER sync_report_files_text_search
+AFTER INSERT OR UPDATE OR DELETE ON report_files
+FOR EACH ROW EXECUTE FUNCTION sync_report_files_to_text_search();

packages/database/drizzle/meta/_journal.json

@@ -806,6 +806,20 @@
       "when": 1759336097951,
       "tag": "0115_glamorous_penance",
       "breakpoints": true
+    },
+    {
+      "idx": 116,
+      "version": "7",
+      "when": 1759356549500,
+      "tag": "0116_cuddly_thaddeus_ross",
+      "breakpoints": true
+    },
+    {
+      "idx": 117,
+      "version": "7",
+      "when": 1759357134333,
+      "tag": "0117_careful_alex_power",
+      "breakpoints": true
     }
   ]
 }

packages/database/src/queries/assets/get-asset-screenshot-bucket-key.ts

@@ -0,0 +1,48 @@
+import { and, eq, isNull } from 'drizzle-orm';
+import { z } from 'zod';
+import { db } from '../../connection';
+import { chats, collections, dashboardFiles, metricFiles, reportFiles } from '../../schema';
+import { type AssetType, AssetTypeSchema } from '../../schema-types';
+
+export const GetAssetScreenshotBucketKeyInputSchema = z.object({
+  assetId: z.string().uuid('Asset ID must be a valid UUID'),
+  assetType: AssetTypeSchema,
+});
+
+export type GetAssetScreenshotBucketKeyInput = z.infer<
+  typeof GetAssetScreenshotBucketKeyInputSchema
+>;
+
+type ScreenshotTable =
+  | typeof chats
+  | typeof collections
+  | typeof dashboardFiles
+  | typeof metricFiles
+  | typeof reportFiles;
+
+const assetTableMap: Record<AssetType, ScreenshotTable> = {
+  chat: chats,
+  collection: collections,
+  dashboard_file: dashboardFiles,
+  metric_file: metricFiles,
+  report_file: reportFiles,
+};
+
+export const getAssetScreenshotBucketKey = async (
+  input: GetAssetScreenshotBucketKeyInput
+): Promise<string | null> => {
+  const { assetId, assetType } = GetAssetScreenshotBucketKeyInputSchema.parse(input);
+
+  const table = assetTableMap[assetType];
+
+  const [existing] = await db
+    .select({ screenshotBucketKey: table.screenshotBucketKey })
+    .from(table)
+    .where(and(eq(table.id, assetId), isNull(table.deletedAt)));
+
+  if (!existing) {
+    throw new Error(`Asset ${assetType} with id ${assetId} not found`);
+  }
+
+  return existing.screenshotBucketKey ?? null;
+};

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

@@ -41,3 +41,15 @@ export {
   type GetUsersWithAssetPermissionsInput,
   type GetUsersWithAssetPermissionsResult,
 } from './get-users-with-asset-permissions';
+
+export {
+  updateAssetScreenshotBucketKey,
+  UpdateAssetScreenshotBucketKeyInputSchema,
+  type UpdateAssetScreenshotBucketKeyInput,
+} from './update-asset-screenshot-bucket-key';
+
+export {
+  getAssetScreenshotBucketKey,
+  GetAssetScreenshotBucketKeyInputSchema,
+  type GetAssetScreenshotBucketKeyInput,
+} from './get-asset-screenshot-bucket-key';

packages/database/src/queries/assets/update-asset-screenshot-bucket-key.ts

@@ -0,0 +1,62 @@
+import { and, eq, isNull } from 'drizzle-orm';
+import { z } from 'zod';
+import { db } from '../../connection';
+import { chats, collections, dashboardFiles, metricFiles, reportFiles } from '../../schema';
+import { type AssetType, AssetTypeSchema } from '../../schema-types';
+
+export const UpdateAssetScreenshotBucketKeyInputSchema = z.object({
+  assetId: z.string().uuid('Asset ID must be a valid UUID'),
+  assetType: AssetTypeSchema,
+  screenshotBucketKey: z.string().min(1, 'Screenshot bucket key is required'),
+});
+
+export type UpdateAssetScreenshotBucketKeyInput = z.infer<
+  typeof UpdateAssetScreenshotBucketKeyInputSchema
+>;
+
+type ScreenshotTable =
+  | typeof chats
+  | typeof collections
+  | typeof dashboardFiles
+  | typeof metricFiles
+  | typeof reportFiles;
+
+const assetTableMap: Record<AssetType, ScreenshotTable> = {
+  chat: chats,
+  collection: collections,
+  dashboard_file: dashboardFiles,
+  metric_file: metricFiles,
+  report_file: reportFiles,
+};
+
+export const updateAssetScreenshotBucketKey = async (
+  input: UpdateAssetScreenshotBucketKeyInput
+): Promise<{ updated: boolean }> => {
+  const { assetId, assetType, screenshotBucketKey } =
+    UpdateAssetScreenshotBucketKeyInputSchema.parse(input);
+
+  const table = assetTableMap[assetType];
+
+  const [existing] = await db
+    .select({ screenshotBucketKey: table.screenshotBucketKey })
+    .from(table)
+    .where(and(eq(table.id, assetId), isNull(table.deletedAt)));
+
+  if (!existing) {
+    throw new Error(`Asset ${assetType} with id ${assetId} not found`);
+  }
+
+  if (existing.screenshotBucketKey === screenshotBucketKey) {
+    return { updated: false };
+  }
+
+  await db
+    .update(table)
+    .set({
+      screenshotBucketKey,
+      updatedAt: new Date().toISOString(),
+    })
+    .where(and(eq(table.id, assetId), isNull(table.deletedAt)));
+
+  return { updated: true };
+};

packages/database/src/queries/search/search-queries.ts

@@ -56,9 +56,10 @@ export async function searchText(input: SearchTextInput): Promise<SearchTextResp
     const paginationCheckCount = page_size + 1;
 
     const filterConditions = [];
+    let fullSearchString = searchString;
 
     if (searchString) {
-      const fullSearchString = `${searchString}*`;
+      fullSearchString = `${searchString}*`;
       filterConditions.push(
         sql`ARRAY[${assetSearchV2.title}, ${assetSearchV2.additionalText}] &@~ ${fullSearchString}`
       );
@@ -89,15 +90,28 @@ export async function searchText(input: SearchTextInput): Promise<SearchTextResp
     // Create the permissioned assets subquery for this user
     const permissionedAssetsSubquery = createPermissionedAssetsSubquery(userId, organizationId);
 
-    // Execute search query with pagination
-    const searchQueryStart = performance.now();
+    const highlightedTitleSql = fullSearchString
+      ? sql<string>`pgroonga_highlight_html(${assetSearchV2.title}, pgroonga_query_extract_keywords(${fullSearchString}))`
+      : assetSearchV2.title;
+
+    const snippetLength = 160;
+    const additionalSnippetSql = fullSearchString
+      ? sql<string>`coalesce(
+           (pgroonga_snippet_html(${assetSearchV2.additionalText},
+                                  pgroonga_query_extract_keywords(${fullSearchString}),
+                                  ${snippetLength}))[1],
+           left(${assetSearchV2.additionalText}, ${snippetLength})
+         )`
+      : sql<string>`left(${assetSearchV2.additionalText}, ${snippetLength})`;
+
     const results = await db
       .select({
         assetId: assetSearchV2.assetId,
         assetType: assetSearchV2.assetType,
-        title: assetSearchV2.title,
-        additionalText: assetSearchV2.additionalText,
+        title: highlightedTitleSql,
+        additionalText: additionalSnippetSql,
         updatedAt: assetSearchV2.updatedAt,
+        screenshotBucketKey: assetSearchV2.screenshotBucketKey,
       })
       .from(assetSearchV2)
       .innerJoin(
@@ -118,11 +132,6 @@ export async function searchText(input: SearchTextInput): Promise<SearchTextResp
       results.pop();
     }
 
-    const searchQueryDuration = performance.now() - searchQueryStart;
-    console.info(
-      `[SEARCH_TIMING] Main search query completed in ${searchQueryDuration.toFixed(2)}ms - found ${results.length} results`
-    );
-
     const paginatedResponse = {
       data: results,
       pagination: {

packages/database/src/schema-types/search.ts

@@ -7,4 +7,5 @@ export const TextSearchResultSchema = z.object({
   title: z.string(),
   additionalText: z.string().nullable(),
   updatedAt: z.string().datetime(),
+  screenshotBucketKey: z.string().nullable(),
 });

packages/database/src/schema.ts

@@ -232,6 +232,7 @@ export const collections = pgTable(
       withTimezone: true,
       mode: 'string',
     }),
+    screenshotBucketKey: text('screenshot_bucket_key'),
   },
   (table) => [
     foreignKey({
@@ -741,6 +742,7 @@ export const dashboardFiles = pgTable(
       withTimezone: true,
       mode: 'string',
     }),
+    screenshotBucketKey: text('screenshot_bucket_key'),
   },
   (table) => [
     index('dashboard_files_created_by_idx').using(
@@ -814,6 +816,7 @@ export const reportFiles = pgTable(
       withTimezone: true,
       mode: 'string',
     }),
+    screenshotBucketKey: text('screenshot_bucket_key'),
   },
   (table) => [
     index('report_files_created_by_idx').using(
@@ -885,6 +888,7 @@ export const chats = pgTable(
       withTimezone: true,
       mode: 'string',
     }),
+    screenshotBucketKey: text('screenshot_bucket_key'),
   },
   (table) => [
     index('chats_created_at_idx').using(
@@ -1035,6 +1039,7 @@ export const metricFiles = pgTable(
       withTimezone: true,
       mode: 'string',
     }),
+    screenshotBucketKey: text('screenshot_bucket_key'),
   },
   (table) => [
     index('metric_files_created_by_idx').using(
@@ -1900,6 +1905,7 @@ export const assetSearchV2 = pgTable(
       .defaultNow()
       .notNull(),
     deletedAt: timestamp('deleted_at', { withTimezone: true, mode: 'string' }),
+    screenshotBucketKey: text('screenshot_bucket_key'),
   },
   (table) => [
     foreignKey({

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

@@ -1,57 +1,24 @@
 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(),
+  key: z.string(),
+  organizationId: z.string().uuid().optional(),
   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>
+  params: GetAssetScreenshotParams,
+  provider?: StorageProvider
 ): 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');
+  const { key, expiresIn, organizationId } = GetAssetScreenshotParamsSchema.parse(params);
+  if (!organizationId && !provider) {
+    throw new Error('Provider or organization ID is required');
   }
+  const s3Client = provider ?? (await getProviderForOrganization(organizationId ?? ''));
 
-  return provider.getSignedUrl(resolvedKey, expiresIn ?? 900);
+  // Asume the key is valid and exists because we have received it from the database
+  return s3Client.getSignedUrl(key, expiresIn ?? 900);
 }

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

@@ -9,17 +9,22 @@ vi.mock('@buster/database/queries', () => ({
   getAssetAncestorsForAssets: vi.fn(),
 }));
 
-vi.mock('./text-processing-helpers', () => ({
-  processSearchResultText: vi.fn(),
+vi.mock('./get-asset-screenshot', () => ({
+  getAssetScreenshotSignedUrl: vi.fn(),
 }));
 
+vi.mock('@buster/data-source', () => ({
+  getProviderForOrganization: vi.fn(),
+}));
+
+import { getProviderForOrganization } from '@buster/data-source';
 // Import the mocked functions
 import {
   getAssetAncestorsForAssets,
   getUserOrganizationId,
   searchText,
 } from '@buster/database/queries';
-import { processSearchResultText } from './text-processing-helpers';
+import { getAssetScreenshotSignedUrl } from './get-asset-screenshot';
 
 describe('search.ts - Unit Tests', () => {
   const mockUserId = 'test-user-id';
@@ -37,6 +42,7 @@ describe('search.ts - Unit Tests', () => {
       title: 'Test Result 1',
       additionalText: 'This is additional text for result 1',
       updatedAt: '2024-01-01T00:00:00.000Z',
+      screenshotBucketKey: 'screenshots/asset-1.png',
     },
     {
       assetId: 'asset-2',
@@ -44,6 +50,7 @@ describe('search.ts - Unit Tests', () => {
       title: 'Test Result 2',
       additionalText: 'This is additional text for result 2',
       updatedAt: '2024-01-01T00:00:00.000Z',
+      screenshotBucketKey: 'screenshots/asset-2.png',
     },
   ];
 
@@ -72,11 +79,9 @@ describe('search.ts - Unit Tests', () => {
     vi.clearAllMocks();
     (getUserOrganizationId as Mock).mockResolvedValue(mockUserOrg);
     (searchText as Mock).mockResolvedValue(mockSearchResponse);
-    (processSearchResultText as Mock).mockImplementation((query, title, additionalText) => ({
-      processedTitle: `<b>${title}</b>`,
-      processedAdditionalText: `<b>${additionalText}</b>`,
-    }));
     (getAssetAncestorsForAssets as Mock).mockResolvedValue(mockAncestorsForAssets);
+    (getAssetScreenshotSignedUrl as Mock).mockResolvedValue('https://example.com/screenshot.png');
+    (getProviderForOrganization as Mock).mockResolvedValue({});
   });
 
   describe('performTextSearch', () => {
@@ -99,18 +104,19 @@ describe('search.ts - Unit Tests', () => {
         page_size: 10,
         filters: {},
       });
+      expect(getAssetScreenshotSignedUrl).not.toHaveBeenCalled();
 
       expect(result).toEqual({
         data: [
           {
             ...mockSearchResults[0],
-            title: '<b>Test Result 1</b>',
-            additionalText: '<b>This is additional text for result 1</b>',
+            title: 'Test Result 1',
+            additionalText: 'This is additional text for result 1',
           },
           {
             ...mockSearchResults[1],
-            title: '<b>Test Result 2</b>',
-            additionalText: '<b>This is additional text for result 2</b>',
+            title: 'Test Result 2',
+            additionalText: 'This is additional text for result 2',
           },
         ],
         pagination: {
@@ -119,18 +125,6 @@ describe('search.ts - Unit Tests', () => {
           has_more: false,
         },
       });
-
-      expect(processSearchResultText).toHaveBeenCalledTimes(2);
-      expect(processSearchResultText).toHaveBeenCalledWith(
-        'test query',
-        'Test Result 1',
-        'This is additional text for result 1'
-      );
-      expect(processSearchResultText).toHaveBeenCalledWith(
-        'test query',
-        'Test Result 2',
-        'This is additional text for result 2'
-      );
     });
 
     it('should handle search with asset type filters', async () => {
@@ -245,6 +239,55 @@ describe('search.ts - Unit Tests', () => {
       expect(result.data[1]).toHaveProperty('ancestors', mockAncestors);
     });
 
+    it('should include screenshots when requested', async () => {
+      const searchRequestWithScreenshots: SearchTextRequest = {
+        ...basicSearchRequest,
+        includeScreenshots: true,
+      };
+
+      const screenshotUrl = 'https://example.com/screenshot.png';
+      (getAssetScreenshotSignedUrl as Mock).mockResolvedValue(screenshotUrl);
+      const mockProvider = {};
+      (getProviderForOrganization as Mock).mockResolvedValue(mockProvider);
+
+      const result = await performTextSearch(mockUserId, searchRequestWithScreenshots);
+
+      expect(getAssetScreenshotSignedUrl).toHaveBeenCalledTimes(mockSearchResults.length);
+      expect(getAssetScreenshotSignedUrl).toHaveBeenNthCalledWith(
+        1,
+        {
+          key: 'screenshots/asset-1.png',
+        },
+        mockProvider
+      );
+      expect(getAssetScreenshotSignedUrl).toHaveBeenNthCalledWith(
+        2,
+        {
+          key: 'screenshots/asset-2.png',
+        },
+        mockProvider
+      );
+
+      expect(result.data[0]).toHaveProperty('screenshotUrl', screenshotUrl);
+      expect(result.data[1]).toHaveProperty('screenshotUrl', screenshotUrl);
+    });
+
+    it('should continue without screenshots if fetching fails', async () => {
+      const searchRequestWithScreenshots: SearchTextRequest = {
+        ...basicSearchRequest,
+        includeScreenshots: true,
+      };
+
+      (getAssetScreenshotSignedUrl as Mock).mockRejectedValueOnce(new Error('missing screenshot'));
+      (getAssetScreenshotSignedUrl as Mock).mockRejectedValueOnce(new Error('missing screenshot'));
+
+      const result = await performTextSearch(mockUserId, searchRequestWithScreenshots);
+
+      expect(getAssetScreenshotSignedUrl).toHaveBeenCalledTimes(mockSearchResults.length);
+      expect(result.data[0]).not.toHaveProperty('screenshotUrl');
+      expect(result.data[1]).not.toHaveProperty('screenshotUrl');
+    });
+
     it('should not include ancestors when not requested', async () => {
       const result = await performTextSearch(mockUserId, basicSearchRequest);
 
@@ -268,7 +311,6 @@ describe('search.ts - Unit Tests', () => {
       const result = await performTextSearch(mockUserId, basicSearchRequest);
 
       expect(result).toEqual(emptySearchResponse);
-      expect(processSearchResultText).not.toHaveBeenCalled();
       expect(getAssetAncestorsForAssets).not.toHaveBeenCalled();
     });
 
@@ -280,6 +322,7 @@ describe('search.ts - Unit Tests', () => {
           title: 'Test Result 1',
           additionalText: null,
           updatedAt: '2024-01-01T00:00:00.000Z',
+          screenshotBucketKey: 'screenshots/asset-1.png',
         },
       ];
 
@@ -289,8 +332,6 @@ describe('search.ts - Unit Tests', () => {
       });
 
       await performTextSearch(mockUserId, basicSearchRequest);
-
-      expect(processSearchResultText).toHaveBeenCalledWith('test query', 'Test Result 1', '');
     });
 
     it('should handle null query', async () => {
@@ -300,12 +341,6 @@ describe('search.ts - Unit Tests', () => {
       };
 
       await performTextSearch(mockUserId, searchRequestWithNullQuery);
-
-      expect(processSearchResultText).toHaveBeenCalledWith(
-        '',
-        'Test Result 1',
-        'This is additional text for result 1'
-      );
     });
 
     it('should throw error when user has no organization', async () => {
@@ -335,6 +370,7 @@ describe('search.ts - Unit Tests', () => {
         title: `Test Result ${i + 1}`,
         additionalText: `Additional text ${i + 1}`,
         updatedAt: '2024-01-01T00:00:00.000Z',
+        screenshotBucketKey: `screenshots/asset-${i + 1}.png`,
       }));
 
       (searchText as Mock).mockResolvedValue({

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

@@ -1,3 +1,4 @@
+import { getProviderForOrganization } from '@buster/data-source';
 import {
   type SearchFilters,
   getAssetAncestorsForAssets,
@@ -5,11 +6,13 @@ import {
   searchText,
 } from '@buster/database/queries';
 import type {
+  AssetAncestors,
   AssetType,
   SearchTextData,
   SearchTextRequest,
   SearchTextResponse,
 } from '@buster/server-shared';
+import { getAssetScreenshotSignedUrl } from './get-asset-screenshot';
 import { processSearchResultText } from './text-processing-helpers';
 
 /**
@@ -28,9 +31,7 @@ export async function performTextSearch(
   );
 
   // Get user's organization
-  const orgLookupStart = performance.now();
   const userOrg = await getUserOrganizationId(userId);
-  const orgLookupDuration = performance.now() - orgLookupStart;
 
   if (!userOrg) {
     throw new Error('User is not associated with an organization');
@@ -64,67 +65,108 @@ export async function performTextSearch(
   });
   const searchDuration = performance.now() - searchStart;
 
-  // Process search result text (highlighting)
-  const textProcessingStart = performance.now();
-  const highlightedResults = await Promise.all(
-    result.data.map(async (searchResult) => {
-      const { processedTitle, processedAdditionalText } = processSearchResultText(
-        searchRequest.query ?? '',
-        searchResult.title,
-        searchResult.additionalText ?? ''
-      );
-      return {
-        ...searchResult,
-        title: processedTitle,
-        additionalText: processedAdditionalText,
-      };
-    })
-  );
+  let screenshotsDuration = 0;
+  let ancestorsDuration = 0;
+
+  const screenshotsStart = performance.now();
+  const screenshotsPromise = searchRequest.includeScreenshots
+    ? fetchScreenshotsByAssetId(result.data, userOrg.organizationId).then((map) => {
+        screenshotsDuration = performance.now() - screenshotsStart;
+        return map;
+      })
+    : Promise.resolve<Record<string, string> | null>(null);
+
+  const ancestorsStart = performance.now();
+  const ancestorsPromise = searchRequest.includeAssetAncestors
+    ? fetchAncestorsByAssetId(result.data, userId, userOrg.organizationId).then((map) => {
+        ancestorsDuration = performance.now() - ancestorsStart;
+        return map;
+      })
+    : Promise.resolve<Record<string, AssetAncestors> | null>(null);
+
+  const [screenshotsByAssetId, ancestorsByAssetId] = await Promise.all([
+    screenshotsPromise,
+    ancestorsPromise,
+  ]);
+
+  // Process screenshots and ancestors in a single loop
+  const processedData = result.data.map((searchResult) => {
+    const updates: Partial<SearchTextData> = {};
+
+    // Add screenshot URL if available
+    if (screenshotsByAssetId) {
+      const screenshotUrl = screenshotsByAssetId[searchResult.assetId];
+      if (screenshotUrl) {
+        updates.screenshotUrl = screenshotUrl;
+      }
+    }
+
+    // Add ancestors if available
+    if (ancestorsByAssetId) {
+      updates.ancestors = ancestorsByAssetId[searchResult.assetId] ?? createEmptyAncestors();
+    }
+
+    // Return original result if no updates, otherwise merge updates
+    return Object.keys(updates).length > 0 ? { ...searchResult, ...updates } : searchResult;
+  });
+
   result = {
     ...result,
-    data: highlightedResults,
+    data: processedData,
   };
-  const textProcessingDuration = performance.now() - textProcessingStart;
+
+  // TODO: Remove this block once we decide if pgroonga highlighting is good enough
+  // // Process search result text (highlighting)
+  // const textProcessingStart = performance.now();
+  // const highlightedResults = await Promise.all(
+  //   result.data.map(async (searchResult) => {
+  //     const { processedTitle, processedAdditionalText } = processSearchResultText(
+  //       searchRequest.query ?? '',
+  //       searchResult.title,
+  //       searchResult.additionalText ?? ''
+  //     );
+  //     return {
+  //       ...searchResult,
+  //       title: processedTitle,
+  //       additionalText: processedAdditionalText,
+  //     };
+  //   })
+  // );
+  // result = {
+  //   ...result,
+  //   data: highlightedResults,
+  // };
+  // const textProcessingDuration = performance.now() - textProcessingStart;
 
   // Add ancestors if requested
-  let ancestorsDuration = 0;
-  if (searchRequest.includeAssetAncestors) {
-    const ancestorsStart = performance.now();
-    const resultsWithAncestors = await addAncestorsToSearchResults(
-      result.data,
-      userId,
-      userOrg.organizationId
-    );
-
-    result = {
-      ...result,
-      data: resultsWithAncestors,
-    };
-    ancestorsDuration = performance.now() - ancestorsStart;
-  }
-
   const totalDuration = performance.now() - startTime;
   console.info(
-    `[SEARCH_PIPELINE_TIMING] performTextSearch completed in ${totalDuration.toFixed(2)}ms total (org: ${orgLookupDuration.toFixed(2)}ms, search: ${searchDuration.toFixed(2)}ms, text: ${textProcessingDuration.toFixed(2)}ms, ancestors: ${ancestorsDuration.toFixed(2)}ms)`
+    `[SEARCH_PIPELINE_TIMING] performTextSearch completed in ${totalDuration.toFixed(2)}ms total (search: ${searchDuration.toFixed(2)}ms, screenshots: ${screenshotsDuration.toFixed(2)}ms, ancestors: ${ancestorsDuration.toFixed(2)}ms)`
   );
 
   return result;
 }
 
-/**
- * Add ancestors to search results in chunks to avoid overwhelming the database
- * @param searchResults - Array of search results to enhance with ancestors
- * @param userId - The user ID making the request
- * @param organizationId - The organization ID
- * @returns Promise<SearchResult[]> - Search results with ancestors added
- */
-async function addAncestorsToSearchResults(
+function createEmptyAncestors(): AssetAncestors {
+  return {
+    chats: [],
+    collections: [],
+    dashboards: [],
+    reports: [],
+  };
+}
+
+async function fetchAncestorsByAssetId(
   searchResults: SearchTextData[],
   userId: string,
   organizationId: string
-): Promise<SearchTextData[]> {
+): Promise<Record<string, AssetAncestors>> {
+  if (searchResults.length === 0) {
+    return {};
+  }
+
   const chunkSize = 50;
-  const resultsWithAncestors: SearchTextData[] = [];
+  const ancestorsByAssetId: Record<string, AssetAncestors> = {};
   const totalChunks = Math.ceil(searchResults.length / chunkSize);
 
   console.info(
@@ -133,7 +175,7 @@ async function addAncestorsToSearchResults(
 
   for (let i = 0; i < searchResults.length; i += chunkSize) {
     const chunk = searchResults.slice(i, i + chunkSize);
-    const ancestorsByAssetId = await getAssetAncestorsForAssets({
+    const chunkAncestors = await getAssetAncestorsForAssets({
       assets: chunk.map((searchResult) => ({
         assetId: searchResult.assetId,
         assetType: searchResult.assetType as AssetType,
@@ -142,17 +184,45 @@ async function addAncestorsToSearchResults(
       organizationId,
     });
 
-    const chunkResults = chunk.map((searchResult) => ({
-      ...searchResult,
-      ancestors: ancestorsByAssetId[searchResult.assetId] ?? {
-        chats: [],
-        collections: [],
-        dashboards: [],
-        reports: [],
-      },
-    }));
-    resultsWithAncestors.push(...chunkResults);
+    Object.assign(ancestorsByAssetId, chunkAncestors);
   }
 
-  return resultsWithAncestors;
+  return ancestorsByAssetId;
+}
+
+async function fetchScreenshotsByAssetId(
+  searchResults: SearchTextData[],
+  organizationId: string
+): Promise<Record<string, string>> {
+  if (searchResults.length === 0) {
+    return {};
+  }
+
+  const provider = await getProviderForOrganization(organizationId);
+
+  const entries = await Promise.all(
+    searchResults.map(async (searchResult) => {
+      try {
+        if (!searchResult.screenshotBucketKey) {
+          return null;
+        }
+        const screenshotUrl = await getAssetScreenshotSignedUrl(
+          { key: searchResult.screenshotBucketKey },
+          provider
+        );
+
+        return [searchResult.assetId, screenshotUrl] as const;
+      } catch {
+        return null;
+      }
+    })
+  );
+
+  return entries.reduce<Record<string, string>>((acc, entry) => {
+    if (entry) {
+      const [assetId, screenshotUrl] = entry;
+      acc[assetId] = screenshotUrl;
+    }
+    return acc;
+  }, {});
 }

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

@@ -19,7 +19,8 @@ export type PutScreenshotResponse = z.infer<typeof PutScreenshotResponseSchema>;
 
 export const GetScreenshotResponseSchema = z.object({
   success: z.boolean(),
-  url: z.string(),
+  url: z.string().optional(),
+  error: z.string().optional(),
 });
 
 export type GetScreenshotResponse = z.infer<typeof GetScreenshotResponseSchema>;

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

@@ -33,6 +33,15 @@ export const SearchTextRequestSchema = z
       }, z.boolean())
       .default(false)
       .optional(),
+    includeScreenshots: z
+      .preprocess((val) => {
+        if (typeof val === 'string') {
+          return val.toLowerCase() === 'true';
+        }
+        return Boolean(val);
+      }, z.boolean())
+      .default(false)
+      .optional(),
     endDate: z.string().datetime().optional(),
     startDate: z.string().datetime().optional(),
   })
@@ -40,6 +49,7 @@ export const SearchTextRequestSchema = z
 
 export const SearchTextDataSchema = TextSearchResultSchema.extend({
   ancestors: AssetAncestorsSchema.optional(),
+  screenshotUrl: z.string().optional(),
 });
 
 export type { AssetAncestors } from '@buster/database/schema-types';