Skip to content

[PR 10] Add validateSchemaNode Error Handling for DML #2486

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
wants to merge 5 commits into from
Closed
Show file tree
Hide file tree
Changes from all commits
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
Original file line number Diff line number Diff line change
Expand Up @@ -17,4 +17,9 @@ export const WORKFLOW_RETRY_CONFIG = {
* by going back to the designSchema node with the error information
*/
MAX_DDL_EXECUTION_RETRIES: 1,
/**
* Maximum number of retries for DML execution failures
* When DML execution fails with retryable errors, the workflow will retry
*/
MAX_DML_EXECUTION_RETRIES: 1,
} as const
Original file line number Diff line number Diff line change
Expand Up @@ -231,4 +231,157 @@ describe('validateSchemaNode', () => {
expect(executeQuery).not.toHaveBeenCalled()
expect(result).toEqual(state)
})

it('should handle missing DML statements', async () => {
const state = createMockState({
dmlStatements: undefined,
})

const repositories = createMockRepositories()
const result = await validateSchemaNode(state, {
configurable: { repositories, logger: mockLogger },
})

expect(executeQuery).not.toHaveBeenCalled()
expect(result).toEqual(state)
})

it('should retry execution on retryable errors', async () => {
const mockResults: SqlResult[] = [
{
success: false,
sql: 'INSERT INTO users VALUES (1);',
result: { error: 'foreign key constraint violation' },
id: 'result-1',
metadata: {
executionTime: 5,
timestamp: new Date().toISOString(),
},
},
]

vi.mocked(executeQuery).mockResolvedValue(mockResults)

const state = createMockState({
dmlStatements: 'INSERT INTO users VALUES (1);',
})

const repositories = createMockRepositories()
const result = await validateSchemaNode(state, {
configurable: { repositories, logger: mockLogger },
})

expect(result.shouldRetryDmlExecution).toBe(true)
expect(result.dmlRetryReason).toContain('foreign key constraint violation')
expect(result.retryCount['dmlExecutionRetry']).toBe(1)
expect(mockLogger.log).toHaveBeenCalledWith(
'[validateSchemaNode] Execution failed with retryable error, scheduling retry',
)
})

it('should not retry execution on non-retryable errors', async () => {
const mockResults: SqlResult[] = [
{
success: false,
sql: 'INSERT INTO users VALUES;',
result: { error: 'syntax error at or near "VALUES"' },
id: 'result-1',
metadata: {
executionTime: 2,
timestamp: new Date().toISOString(),
},
},
]

vi.mocked(executeQuery).mockResolvedValue(mockResults)

const state = createMockState({
dmlStatements: 'INSERT INTO users VALUES;',
})

const repositories = createMockRepositories()
const result = await validateSchemaNode(state, {
configurable: { repositories, logger: mockLogger },
})

expect(result.shouldRetryDmlExecution).toBeUndefined()
expect(result.dmlExecutionErrors).toContain('syntax error')
expect(mockLogger.log).toHaveBeenCalledWith(
'[validateSchemaNode] Completed with errors',
)
})

it('should not retry after max retries exceeded', async () => {
const mockResults: SqlResult[] = [
{
success: false,
sql: 'INSERT INTO users VALUES (1);',
result: { error: 'foreign key constraint violation' },
id: 'result-1',
metadata: {
executionTime: 5,
timestamp: new Date().toISOString(),
},
},
]

vi.mocked(executeQuery).mockResolvedValue(mockResults)

const state = createMockState({
dmlStatements: 'INSERT INTO users VALUES (1);',
retryCount: { dmlExecutionRetry: 1 }, // Already at max retries
})

const repositories = createMockRepositories()
const result = await validateSchemaNode(state, {
configurable: { repositories, logger: mockLogger },
})

expect(result.shouldRetryDmlExecution).toBeUndefined()
expect(result.dmlExecutionErrors).toContain(
'foreign key constraint violation',
)
expect(result.retryCount['dmlExecutionRetry']).toBe(1)
})

it('should handle mixed success and failure in execution', async () => {
const mockResults: SqlResult[] = [
{
success: true,
sql: 'INSERT INTO categories VALUES (1, "test");',
result: { rows: [], columns: [] },
id: 'result-1',
metadata: {
executionTime: 5,
timestamp: new Date().toISOString(),
affectedRows: 1,
},
},
{
success: false,
sql: 'INSERT INTO products VALUES (1, 999);',
result: { error: 'foreign key constraint violation on category_id' },
id: 'result-2',
metadata: {
executionTime: 3,
timestamp: new Date().toISOString(),
},
},
]

vi.mocked(executeQuery).mockResolvedValue(mockResults)

const state = createMockState({
dmlStatements:
'INSERT INTO categories VALUES (1, "test"); INSERT INTO products VALUES (1, 999);',
})

const repositories = createMockRepositories()
const result = await validateSchemaNode(state, {
configurable: { repositories, logger: mockLogger },
})

expect(result.shouldRetryDmlExecution).toBe(true)
expect(result.dmlRetryReason).toContain('foreign key constraint violation')
})
})
Original file line number Diff line number Diff line change
@@ -1,9 +1,27 @@
import type { RunnableConfig } from '@langchain/core/runnables'
import { executeQuery } from '@liam-hq/pglite-server'
import type { SqlResult } from '@liam-hq/pglite-server/src/types'
import { WORKFLOW_RETRY_CONFIG } from '../constants'
import { getConfigurable } from '../shared/getConfigurable'
import type { WorkflowState } from '../types'

const NODE_NAME = 'validateSchemaNode'

// Error patterns that might benefit from retry
const RETRYABLE_ERROR_PATTERNS = [
'constraint violation',
'foreign key',
'duplicate key',
'deadlock',
'lock timeout',
]

function isRetryableError(errorMessage: string): boolean {
const lowerError = errorMessage.toLowerCase()
return RETRYABLE_ERROR_PATTERNS.some((pattern) =>
lowerError.includes(pattern),
)
}
/**
* Validate Schema Node - Combined DDL/DML Execution & Validation
* Executes DDL and DML together in a single query to validate schema with test data
Expand Down Expand Up @@ -54,6 +72,39 @@ export async function validateSchemaNode(
)
.join('; ')

configurableResult.value.logger.error(
`[${NODE_NAME}] Execution failed: ${errorMessages}`,
)

// Check if errors are retryable
const hasRetryableError = isRetryableError(errorMessages)
const currentRetryCount = state.retryCount?.['dmlExecutionRetry'] || 0

if (
hasRetryableError &&
currentRetryCount < WORKFLOW_RETRY_CONFIG.MAX_DML_EXECUTION_RETRIES
) {
configurableResult.value.logger.log(
`[${NODE_NAME}] Execution failed with retryable error, scheduling retry`,
)
configurableResult.value.logger.log(
`[${NODE_NAME}] Completed with retry scheduled`,
)

return {
...state,
dmlExecutionErrors: errorMessages,
shouldRetryDmlExecution: true,
dmlRetryReason: errorMessages,
retryCount: {
...state.retryCount,
dmlExecutionRetry: currentRetryCount + 1,
},
}
}

// Non-retryable error or max retries exceeded
configurableResult.value.logger.log(`[${NODE_NAME}] Completed with errors`)
return {
...state,
dmlExecutionErrors: errorMessages,
Expand Down
2 changes: 2 additions & 0 deletions frontend/internal-packages/agent/src/chat/workflow/types.ts
Original file line number Diff line number Diff line change
Expand Up @@ -33,6 +33,8 @@ export type WorkflowState = {
// DML execution results
dmlExecutionSuccessful?: boolean | undefined
dmlExecutionErrors?: string | undefined
shouldRetryDmlExecution?: boolean | undefined
dmlRetryReason?: string | undefined

// Schema update fields
buildingSchemaId: string
Expand Down