Integrate access controls with metric file downloads

- Added permission checks to metric file download handler
- Added permission checks to export-metric-data trigger task
- Created comprehensive tests for access control integration
- Updated AssetType enums to include report_file
- Resolved export conflicts in server-shared package
- Fixed AI package imports for new getPermissionedDatasets API

apps/server/src/api/v2/metric_files/download-metric-file.test.ts

@@ -0,0 +1,356 @@
+import { checkPermission } from '@buster/access-controls';
+import { getUserOrganizationId } from '@buster/database';
+import { tasks } from '@trigger.dev/sdk';
+import { HTTPException } from 'hono/http-exception';
+import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest';
+import { downloadMetricFileHandler } from './download-metric-file';
+
+// Mock all external dependencies
+vi.mock('@buster/access-controls');
+vi.mock('@buster/database');
+vi.mock('@trigger.dev/sdk', () => ({
+  tasks: {
+    trigger: vi.fn(),
+  },
+  runs: {
+    retrieve: vi.fn(),
+  },
+}));
+
+describe('downloadMetricFileHandler', () => {
+  const mockUser = {
+    id: 'user-123',
+    email: 'test@example.com',
+    name: 'Test User',
+  };
+
+  const mockMetricId = 'metric-456';
+  const mockOrganizationId = 'org-789';
+
+  beforeEach(() => {
+    vi.clearAllMocks();
+  });
+
+  afterEach(() => {
+    vi.restoreAllMocks();
+  });
+
+  describe('Access Control', () => {
+    it('should throw 403 if user is not part of an organization', async () => {
+      vi.mocked(getUserOrganizationId).mockResolvedValue(null);
+
+      await expect(downloadMetricFileHandler(mockMetricId, mockUser as any)).rejects.toThrow(
+        HTTPException
+      );
+
+      await expect(downloadMetricFileHandler(mockMetricId, mockUser as any)).rejects.toMatchObject({
+        status: 403,
+        message: 'You must be part of an organization to download metric files',
+      });
+    });
+
+    it('should throw 403 if user does not have permission to view the metric file', async () => {
+      vi.mocked(getUserOrganizationId).mockResolvedValue({
+        organizationId: mockOrganizationId,
+        userId: mockUser.id,
+        role: 'member',
+      } as any);
+
+      vi.mocked(checkPermission).mockResolvedValue({
+        hasAccess: false,
+        effectiveRole: undefined,
+        accessPath: undefined,
+      });
+
+      await expect(downloadMetricFileHandler(mockMetricId, mockUser as any)).rejects.toThrow(
+        HTTPException
+      );
+
+      await expect(downloadMetricFileHandler(mockMetricId, mockUser as any)).rejects.toMatchObject({
+        status: 403,
+        message: 'You do not have permission to download this metric file',
+      });
+
+      // Verify permission check was called with correct parameters
+      expect(checkPermission).toHaveBeenCalledWith({
+        userId: mockUser.id,
+        assetId: mockMetricId,
+        assetType: 'metric_file',
+        requiredRole: 'can_view',
+        organizationId: mockOrganizationId,
+      });
+    });
+
+    it('should proceed with export when user has permission', async () => {
+      const mockTaskHandle = { id: 'task-handle-123' };
+
+      vi.mocked(getUserOrganizationId).mockResolvedValue({
+        organizationId: mockOrganizationId,
+        userId: mockUser.id,
+        role: 'member',
+      } as any);
+
+      vi.mocked(checkPermission).mockResolvedValue({
+        hasAccess: true,
+        effectiveRole: 'can_view',
+        accessPath: 'direct',
+      });
+
+      vi.mocked(tasks.trigger).mockResolvedValue(mockTaskHandle as any);
+
+      // Mock successful task completion
+      const { runs } = await import('@trigger.dev/sdk');
+      vi.mocked(runs.retrieve).mockResolvedValue({
+        status: 'COMPLETED',
+        output: {
+          success: true,
+          downloadUrl: 'https://example.com/download',
+          expiresAt: new Date(Date.now() + 60000).toISOString(),
+          fileSize: 1024,
+          fileName: 'metric-456.csv',
+          rowCount: 100,
+        },
+      } as any);
+
+      const result = await downloadMetricFileHandler(mockMetricId, mockUser as any);
+
+      // Verify permission check was performed
+      expect(checkPermission).toHaveBeenCalledWith({
+        userId: mockUser.id,
+        assetId: mockMetricId,
+        assetType: 'metric_file',
+        requiredRole: 'can_view',
+        organizationId: mockOrganizationId,
+      });
+
+      // Verify task was triggered with correct parameters
+      expect(tasks.trigger).toHaveBeenCalledWith('export-metric-data', {
+        metricId: mockMetricId,
+        userId: mockUser.id,
+        organizationId: mockOrganizationId,
+      });
+
+      // Verify successful response
+      expect(result).toMatchObject({
+        downloadUrl: 'https://example.com/download',
+        fileSize: 1024,
+        fileName: 'metric-456.csv',
+        rowCount: 100,
+        message: 'Download link expires in 60 seconds. Please start your download immediately.',
+      });
+    });
+
+    it('should respect different permission levels', async () => {
+      vi.mocked(getUserOrganizationId).mockResolvedValue({
+        organizationId: mockOrganizationId,
+        userId: mockUser.id,
+        role: 'member',
+      } as any);
+
+      // Test with 'owner' permission
+      vi.mocked(checkPermission).mockResolvedValue({
+        hasAccess: true,
+        effectiveRole: 'owner',
+        accessPath: 'direct',
+      });
+
+      vi.mocked(tasks.trigger).mockResolvedValue({ id: 'task-123' } as any);
+
+      const { runs } = await import('@trigger.dev/sdk');
+      vi.mocked(runs.retrieve).mockResolvedValue({
+        status: 'COMPLETED',
+        output: {
+          success: true,
+          downloadUrl: 'https://example.com/download',
+          expiresAt: new Date(Date.now() + 60000).toISOString(),
+        },
+      } as any);
+
+      await downloadMetricFileHandler(mockMetricId, mockUser as any);
+
+      expect(checkPermission).toHaveBeenCalledWith({
+        userId: mockUser.id,
+        assetId: mockMetricId,
+        assetType: 'metric_file',
+        requiredRole: 'can_view',
+        organizationId: mockOrganizationId,
+      });
+    });
+
+    it('should handle workspace sharing permissions', async () => {
+      vi.mocked(getUserOrganizationId).mockResolvedValue({
+        organizationId: mockOrganizationId,
+        userId: mockUser.id,
+        role: 'member',
+      } as any);
+
+      // Simulate access through workspace sharing
+      vi.mocked(checkPermission).mockResolvedValue({
+        hasAccess: true,
+        effectiveRole: 'can_view',
+        accessPath: 'workspace_sharing',
+      });
+
+      vi.mocked(tasks.trigger).mockResolvedValue({ id: 'task-123' } as any);
+
+      const { runs } = await import('@trigger.dev/sdk');
+      vi.mocked(runs.retrieve).mockResolvedValue({
+        status: 'COMPLETED',
+        output: {
+          success: true,
+          downloadUrl: 'https://example.com/download',
+          expiresAt: new Date(Date.now() + 60000).toISOString(),
+        },
+      } as any);
+
+      await downloadMetricFileHandler(mockMetricId, mockUser as any);
+
+      expect(checkPermission).toHaveBeenCalled();
+      expect(tasks.trigger).toHaveBeenCalled();
+    });
+
+    it('should handle cascading permissions', async () => {
+      vi.mocked(getUserOrganizationId).mockResolvedValue({
+        organizationId: mockOrganizationId,
+        userId: mockUser.id,
+        role: 'member',
+      } as any);
+
+      // Simulate access through cascading permissions (e.g., from parent collection)
+      vi.mocked(checkPermission).mockResolvedValue({
+        hasAccess: true,
+        effectiveRole: 'can_edit',
+        accessPath: 'cascading',
+      });
+
+      vi.mocked(tasks.trigger).mockResolvedValue({ id: 'task-123' } as any);
+
+      const { runs } = await import('@trigger.dev/sdk');
+      vi.mocked(runs.retrieve).mockResolvedValue({
+        status: 'COMPLETED',
+        output: {
+          success: true,
+          downloadUrl: 'https://example.com/download',
+          expiresAt: new Date(Date.now() + 60000).toISOString(),
+        },
+      } as any);
+
+      await downloadMetricFileHandler(mockMetricId, mockUser as any);
+
+      expect(checkPermission).toHaveBeenCalled();
+      expect(tasks.trigger).toHaveBeenCalled();
+    });
+  });
+
+  describe('Error Handling', () => {
+    beforeEach(() => {
+      // Setup basic mocks for permission checks to pass
+      vi.mocked(getUserOrganizationId).mockResolvedValue({
+        organizationId: mockOrganizationId,
+        userId: mockUser.id,
+        role: 'member',
+      } as any);
+
+      vi.mocked(checkPermission).mockResolvedValue({
+        hasAccess: true,
+        effectiveRole: 'can_view',
+        accessPath: 'direct',
+      });
+    });
+
+    it('should handle task failure', async () => {
+      vi.mocked(tasks.trigger).mockResolvedValue({ id: 'task-123' } as any);
+
+      const { runs } = await import('@trigger.dev/sdk');
+      vi.mocked(runs.retrieve).mockResolvedValue({
+        status: 'FAILED',
+      } as any);
+
+      await expect(downloadMetricFileHandler(mockMetricId, mockUser as any)).rejects.toThrow(
+        HTTPException
+      );
+
+      await expect(downloadMetricFileHandler(mockMetricId, mockUser as any)).rejects.toMatchObject({
+        status: 500,
+        message: 'Export task failed',
+      });
+    });
+
+    it('should handle timeout', async () => {
+      vi.mocked(tasks.trigger).mockResolvedValue({ id: 'task-123' } as any);
+
+      const { runs } = await import('@trigger.dev/sdk');
+
+      // Simulate task still in progress after timeout
+      let callCount = 0;
+      vi.mocked(runs.retrieve).mockImplementation(async () => {
+        callCount++;
+        // Return EXECUTING status to simulate timeout
+        return { status: 'EXECUTING' } as any;
+      });
+
+      // Speed up test by reducing timeout
+      vi.useFakeTimers();
+
+      const promise = downloadMetricFileHandler(mockMetricId, mockUser as any);
+
+      // Fast-forward time to trigger timeout
+      await vi.runAllTimersAsync();
+
+      await expect(promise).rejects.toThrow(HTTPException);
+
+      vi.useRealTimers();
+    });
+
+    it('should handle specific error codes from export task', async () => {
+      vi.mocked(tasks.trigger).mockResolvedValue({ id: 'task-123' } as any);
+
+      const { runs } = await import('@trigger.dev/sdk');
+
+      // Test UNAUTHORIZED error code
+      vi.mocked(runs.retrieve).mockResolvedValue({
+        status: 'COMPLETED',
+        output: {
+          success: false,
+          error: 'Access denied to data source',
+          errorCode: 'UNAUTHORIZED',
+        },
+      } as any);
+
+      await expect(downloadMetricFileHandler(mockMetricId, mockUser as any)).rejects.toMatchObject({
+        status: 403,
+        message: 'Access denied to data source',
+      });
+
+      // Test NOT_FOUND error code
+      vi.mocked(runs.retrieve).mockResolvedValue({
+        status: 'COMPLETED',
+        output: {
+          success: false,
+          error: 'Metric not found',
+          errorCode: 'NOT_FOUND',
+        },
+      } as any);
+
+      await expect(downloadMetricFileHandler(mockMetricId, mockUser as any)).rejects.toMatchObject({
+        status: 404,
+        message: 'Metric file not found or data source credentials missing',
+      });
+
+      // Test QUERY_ERROR
+      vi.mocked(runs.retrieve).mockResolvedValue({
+        status: 'COMPLETED',
+        output: {
+          success: false,
+          error: 'SQL syntax error',
+          errorCode: 'QUERY_ERROR',
+        },
+      } as any);
+
+      await expect(downloadMetricFileHandler(mockMetricId, mockUser as any)).rejects.toMatchObject({
+        status: 400,
+        message: 'Query execution failed: SQL syntax error',
+      });
+    });
+  });
+});

apps/server/src/api/v2/metric_files/download-metric-file.ts

@@ -1,3 +1,4 @@
+import { type AssetPermissionCheck, checkPermission } from '@buster/access-controls';
 import type { User } from '@buster/database';
 import { getUserOrganizationId } from '@buster/database';
 import type { ExportMetricDataOutput, MetricDownloadResponse } from '@buster/server-shared/metrics';
@@ -9,9 +10,10 @@ import { HTTPException } from 'hono/http-exception';
  *
  * This handler:
  * 1. Validates user has access to the organization
- * 2. Triggers the export task in Trigger.dev
+ * 2. Checks user has permission to view the metric file
- * 3. Waits for the task to complete (max 2 minutes)
+ * 3. Triggers the export task in Trigger.dev
- * 4. Returns a presigned URL for downloading the CSV file
+ * 4. Waits for the task to complete (max 2 minutes)
+ * 5. Returns a presigned URL for downloading the CSV file
  *
  * The download URL expires after 60 seconds for security
  */
@@ -30,6 +32,23 @@ export async function downloadMetricFileHandler(
 
   const { organizationId } = userOrg;
 
+  // Check if user has permission to view this metric file
+  const permissionCheck: AssetPermissionCheck = {
+    userId: user.id,
+    assetId: metricId,
+    assetType: 'metric_file',
+    requiredRole: 'can_view',
+    organizationId,
+  };
+
+  const permissionResult = await checkPermission(permissionCheck);
+
+  if (!permissionResult.hasAccess) {
+    throw new HTTPException(403, {
+      message: 'You do not have permission to download this metric file',
+    });
+  }
+
   try {
     // Trigger the export task
     const handle = await tasks.trigger('export-metric-data', {

apps/trigger/src/tasks/export-metric-data/export-metric-data.ts

@@ -1,6 +1,7 @@
 import { randomBytes } from 'node:crypto';
 import { GetObjectCommand, PutObjectCommand, S3Client } from '@aws-sdk/client-s3';
 import { getSignedUrl } from '@aws-sdk/s3-request-presigner';
+import { type AssetPermissionCheck, checkPermission } from '@buster/access-controls';
 import { createAdapter } from '@buster/data-source';
 import type { Credentials } from '@buster/data-source';
 import { getDataSourceCredentials, getMetricForExport } from '@buster/database';
@@ -41,13 +42,14 @@ const MAX_FILE_SIZE = 500 * 1024 * 1024; // 500MB max file size
  * Task for exporting metric data to CSV and generating a presigned download URL
  *
  * This task:
- * 1. Fetches metric configuration and validates user access
+ * 1. Fetches metric configuration and validates organization access
- * 2. Retrieves data source credentials from vault
+ * 2. Checks user has permission to view the metric file
- * 3. Executes the metric's SQL query
+ * 3. Retrieves data source credentials from vault
- * 4. Converts results to CSV format
+ * 4. Executes the metric's SQL query
- * 5. Uploads to R2 storage
+ * 5. Converts results to CSV format
- * 6. Generates a 60-second presigned URL for download
+ * 6. Uploads to R2 storage
- * 7. Schedules cleanup after 60 seconds
+ * 7. Generates a 60-second presigned URL for download
+ * 8. Schedules cleanup after 60 seconds
  */
 export const exportMetricData: ReturnType<
   typeof schemaTask<
@@ -96,6 +98,31 @@ export const exportMetricData: ReturnType<
         };
       }
 
+      // Step 1b: Check user has permission to view this metric file
+      const permissionCheck: AssetPermissionCheck = {
+        userId: payload.userId,
+        assetId: payload.metricId,
+        assetType: 'metric_file',
+        requiredRole: 'can_view',
+        organizationId: payload.organizationId,
+      };
+
+      const permissionResult = await checkPermission(permissionCheck);
+
+      if (!permissionResult.hasAccess) {
+        logger.error('User lacks permission to access metric file', {
+          metricId: payload.metricId,
+          userId: payload.userId,
+          organizationId: payload.organizationId,
+        });
+
+        return {
+          success: false,
+          error: 'You do not have permission to export this metric',
+          errorCode: 'UNAUTHORIZED',
+        };
+      }
+
       logger.log('Metric validated', {
         metricName: metric.name,
         dataSourceId: metric.dataSourceId,

packages/access-controls/src/types/asset-permissions.ts

@@ -14,6 +14,7 @@ export const AssetTypeSchema = z.enum([
   'chat',
   'metric_file',
   'dashboard_file',
+  'report_file',
   'data_source',
   'metric',
   'filter',
@@ -32,6 +33,7 @@ export type CascadingAssetType =
   | 'chat'
   | 'metric_file'
   | 'dashboard_file'
+  | 'report_file'
   | 'collection';
 
 // Permission roles - matching database AssetPermissionRole enum

packages/ai/src/agents/think-and-prep-agent/investigation-instructions.ts

@@ -715,11 +715,11 @@ export const getThinkAndPrepInstructions = async ({
   const userId = runtimeContext.get('userId');
   const dataSourceSyntax = runtimeContext.get('dataSourceSyntax');
 
-  const datasets = await getPermissionedDatasets(userId, 0, 1000);
+  const datasetResults = await getPermissionedDatasets({ userId, page: 0, pageSize: 1000 });
 
   // Extract yml_content from each dataset and join with separators
-  const assembledYmlContent = datasets
+  const assembledYmlContent = datasetResults.datasets
-    .map((dataset: { ymlFile: string | null | undefined }) => dataset.ymlFile)
+    .map((dataset: { ymlContent: string | null | undefined }) => dataset.ymlContent)
     .filter((content: string | null | undefined) => content !== null && content !== undefined)
     .join('\n---\n');
 

packages/database/src/queries/asset-permissions/check-asset-permission.ts

@@ -17,7 +17,14 @@ import type { AssetPermissionRole, AssetType, WorkspaceSharing } from '../../sch
 export interface CheckAssetPermissionParams {
   userId: string;
   assetId: string;
-  assetType: 'dashboard' | 'thread' | 'chat' | 'metric_file' | 'dashboard_file' | 'collection';
+  assetType:
+    | 'dashboard'
+    | 'thread'
+    | 'chat'
+    | 'metric_file'
+    | 'dashboard_file'
+    | 'report_file'
+    | 'collection';
   organizationId?: string;
 }
 

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

@@ -6,6 +6,7 @@ export type AssetType =
   | 'chat'
   | 'metric_file'
   | 'dashboard_file'
+  | 'report_file'
   | 'collection'
   | 'data_source'
   | 'metric'

packages/server-shared/src/index.ts

@@ -12,6 +12,8 @@ export * from './message';
 export * from './metrics';
 export * from './organization';
 export * from './security';
+// Export share module (has some naming conflicts with chats and metrics)
+// TODO: Resolve naming conflicts properly
 export * from './share';
 export * from './slack';
 export * from './teams';

packages/server-shared/src/share/assets.ts

@@ -12,6 +12,9 @@ const AssetPermissionRoleEnums: Record<AssetPermissionRoleBase, AssetPermissionR
     can_view: 'can_view',
   });
 
-export const AssetPermissionRoleSchema = z.enum(
+export const ShareAssetPermissionRoleSchema = z.enum(
   Object.values(AssetPermissionRoleEnums) as [AssetPermissionRoleBase, ...AssetPermissionRoleBase[]]
 );
+
+// Keep the old name for backward compatibility but don't export it from index
+const AssetPermissionRoleSchema = ShareAssetPermissionRoleSchema;

packages/server-shared/src/share/individual-permissions.ts

@@ -1,12 +1,12 @@
 import { z } from 'zod';
-import { AssetPermissionRoleSchema } from './assets';
+import { ShareAssetPermissionRoleSchema } from './assets';
 
 export const IndividualPermissionSchema = z.object({
   id: z.string(),
   name: z.string().nullable(),
   email: z.string(),
   avatar_url: z.string().nullable(),
-  role: AssetPermissionRoleSchema,
+  role: ShareAssetPermissionRoleSchema,
 });
 
 export type IndividualPermission = z.infer<typeof IndividualPermissionSchema>;

packages/server-shared/src/share/requests.ts

@@ -13,7 +13,7 @@ export const SharePostRequestSchema = z.array(
 export type SharePostRequest = z.infer<typeof SharePostRequestSchema>;
 
 //Used for updating share permissions for a report, collection, or metric
-export const ShareUpdateRequestSchema = z.object({
+export const SharePermissionsUpdateRequestSchema = z.object({
   publicly_accessible: z.boolean().optional(),
   public_expiry_date: z.string().nullable().optional(),
   public_password: z.string().nullable().optional(),
@@ -28,7 +28,11 @@ export const ShareUpdateRequestSchema = z.object({
     .optional(),
 });
 
-export type ShareUpdateRequest = z.infer<typeof ShareUpdateRequestSchema>;
+export type SharePermissionsUpdateRequest = z.infer<typeof SharePermissionsUpdateRequestSchema>;
+
+// Keep old names for backward compatibility but don't export from index
+const ShareUpdateRequestSchema = SharePermissionsUpdateRequestSchema;
+type ShareUpdateRequest = SharePermissionsUpdateRequest;
 
 //Used for deleting share permissions for a report, collection, or metric
 export const ShareDeleteRequestSchema = z.array(z.string());