Skip to content
Merged
Show file tree
Hide file tree
Changes from 1 commit
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
94 changes: 92 additions & 2 deletions apps/sim/lib/copilot/tools/server/table/user-table.test.ts
Original file line number Diff line number Diff line change
Expand Up @@ -12,6 +12,9 @@ const {
mockBatchInsertRows,
mockReplaceTableRows,
mockAddWorkflowGroup,
mockCreateTable,
mockDeleteTable,
mockGetWorkspaceTableLimits,
fakeEnrichment,
} = vi.hoisted(() => ({
mockResolveWorkspaceFileReference: vi.fn(),
Expand All @@ -20,6 +23,9 @@ const {
mockBatchInsertRows: vi.fn(),
mockReplaceTableRows: vi.fn(),
mockAddWorkflowGroup: vi.fn(),
mockCreateTable: vi.fn(),
mockDeleteTable: vi.fn(),
mockGetWorkspaceTableLimits: vi.fn(),
fakeEnrichment: {
id: 'work-email',
name: 'Work Email',
Expand Down Expand Up @@ -54,13 +60,13 @@ vi.mock('@/lib/table/service', () => ({
addWorkflowGroup: mockAddWorkflowGroup,
batchInsertRows: mockBatchInsertRows,
batchUpdateRows: vi.fn(),
createTable: vi.fn(),
createTable: mockCreateTable,
deleteColumn: vi.fn(),
deleteColumns: vi.fn(),
deleteRow: vi.fn(),
deleteRowsByFilter: vi.fn(),
deleteRowsByIds: vi.fn(),
deleteTable: vi.fn(),
deleteTable: mockDeleteTable,
getRowById: vi.fn(),
getTableById: mockGetTableById,
insertRow: vi.fn(),
Expand All @@ -74,6 +80,10 @@ vi.mock('@/lib/table/service', () => ({
updateRowsByFilter: vi.fn(),
}))

vi.mock('@/lib/table/billing', () => ({
getWorkspaceTableLimits: mockGetWorkspaceTableLimits,
}))

import { userTableServerTool } from '@/lib/copilot/tools/server/table/user-table'

function buildTable(overrides: Partial<TableDefinition> = {}): TableDefinition {
Expand Down Expand Up @@ -232,6 +242,86 @@ describe('userTableServerTool.import_file', () => {
})
})

describe('userTableServerTool.create_from_file', () => {
beforeEach(() => {
vi.clearAllMocks()
mockResolveWorkspaceFileReference.mockResolvedValue({ name: 'people.csv', type: 'text/csv' })
mockDownloadWorkspaceFile.mockResolvedValue(Buffer.from('name,age\nAlice,30\nBob,40'))
mockGetWorkspaceTableLimits.mockResolvedValue({ maxRowsPerTable: 1000, maxTables: 3 })
mockCreateTable.mockResolvedValue(buildTable({ id: 'tbl_new', name: 'people' }))
mockBatchInsertRows.mockImplementation(async (data: { rows: unknown[] }) =>
data.rows.map((_, i) => ({ id: `row_${i}` }))
)
})

it('stamps the workspace plan limits on the created table', async () => {
const result = await userTableServerTool.execute(
{ operation: 'create_from_file', args: { fileId: 'file-1' } },
{ userId: 'user-1', workspaceId: 'workspace-1' }
)

expect(result.success).toBe(true)
expect(mockGetWorkspaceTableLimits).toHaveBeenCalledWith('workspace-1')
expect(mockCreateTable).toHaveBeenCalledTimes(1)
const createArgs = mockCreateTable.mock.calls[0][0] as { maxRows: number; maxTables: number }
expect(createArgs.maxRows).toBe(1000)
expect(createArgs.maxTables).toBe(3)
})

it('rejects a file exceeding the plan row limit without creating a table', async () => {
mockGetWorkspaceTableLimits.mockResolvedValueOnce({ maxRowsPerTable: 1, maxTables: 3 })

const result = await userTableServerTool.execute(
{ operation: 'create_from_file', args: { fileId: 'file-1' } },
{ userId: 'user-1', workspaceId: 'workspace-1' }
)

expect(result.success).toBe(false)
expect(result.message).toMatch(/exceeds this plan's limit/i)
expect(mockCreateTable).not.toHaveBeenCalled()
expect(mockDeleteTable).not.toHaveBeenCalled()
})

it('deletes the created table when row insertion fails', async () => {
mockBatchInsertRows.mockRejectedValueOnce(new Error('Maximum row limit (1000) reached'))

const result = await userTableServerTool.execute(
{ operation: 'create_from_file', args: { fileId: 'file-1' } },
{ userId: 'user-1', workspaceId: 'workspace-1' }
)

expect(result.success).toBe(false)
expect(mockDeleteTable).toHaveBeenCalledWith('tbl_new', expect.any(String))
})
})

describe('userTableServerTool.create', () => {
beforeEach(() => {
vi.clearAllMocks()
mockGetWorkspaceTableLimits.mockResolvedValue({ maxRowsPerTable: 1000, maxTables: 3 })
mockCreateTable.mockResolvedValue(buildTable({ id: 'tbl_new', name: 'People' }))
})

it('stamps the workspace plan limits on the created table', async () => {
const result = await userTableServerTool.execute(
{
operation: 'create',
args: {
name: 'People',
schema: { columns: [{ name: 'name', type: 'string', required: true }] },
},
},
{ userId: 'user-1', workspaceId: 'workspace-1' }
)

expect(result.success).toBe(true)
expect(mockGetWorkspaceTableLimits).toHaveBeenCalledWith('workspace-1')
const createArgs = mockCreateTable.mock.calls[0][0] as { maxRows: number; maxTables: number }
expect(createArgs.maxRows).toBe(1000)
expect(createArgs.maxTables).toBe(3)
})
})

describe('userTableServerTool.list_enrichments', () => {
beforeEach(() => {
vi.clearAllMocks()
Expand Down
23 changes: 22 additions & 1 deletion apps/sim/lib/copilot/tools/server/table/user-table.ts
Original file line number Diff line number Diff line change
Expand Up @@ -14,6 +14,7 @@ import {
type CsvHeaderMapping,
CsvImportValidationError,
coerceRowsForTable,
getWorkspaceTableLimits,
inferSchemaFromCsv,
parseCsvBuffer,
sanitizeName,
Expand Down Expand Up @@ -263,13 +264,16 @@ export const userTableServerTool: BaseServerTool<UserTableArgs, UserTableResult>

const requestId = generateId().slice(0, 8)
assertNotAborted()
const planLimits = await getWorkspaceTableLimits(workspaceId)
const table = await createTable(
{
name: args.name,
description: args.description,
schema: args.schema,
workspaceId,
userId: context.userId,
maxRows: planLimits.maxRowsPerTable,
maxTables: planLimits.maxTables,
},
requestId
)
Expand Down Expand Up @@ -761,19 +765,36 @@ export const userTableServerTool: BaseServerTool<UserTableArgs, UserTableResult>
const tableName = args.name || file.name.replace(/\.[^.]+$/, '')
const requestId = generateId().slice(0, 8)
assertNotAborted()
const planLimits = await getWorkspaceTableLimits(workspaceId)

if (rows.length > planLimits.maxRowsPerTable) {
return {
success: false,
message: `"${file.name}" has ${rows.length.toLocaleString()} rows, which exceeds this plan's limit of ${planLimits.maxRowsPerTable.toLocaleString()} rows per table.`,
}
}

const table = await createTable(
{
name: tableName,
description: args.description || `Imported from ${file.name}`,
schema: { columns },
workspaceId,
userId: context.userId,
maxRows: planLimits.maxRowsPerTable,
maxTables: planLimits.maxTables,
},
requestId
)

const coerced = coerceRowsForTable(rows, { columns }, headerToColumn)
const inserted = await batchInsertAll(table.id, coerced, table, workspaceId, context)
let inserted: number
try {
inserted = await batchInsertAll(table.id, coerced, table, workspaceId, context)
} catch (insertError) {
await deleteTable(table.id, requestId).catch(() => {})
throw insertError
}
Comment thread
greptile-apps[bot] marked this conversation as resolved.
Comment thread
greptile-apps[bot] marked this conversation as resolved.
Comment thread
greptile-apps[bot] marked this conversation as resolved.

logger.info('Table created from file', {
tableId: table.id,
Expand Down
Loading