Skip to content

Commit 13879f8

Browse files
ArtuomkaArtuomka
committed
feat: implement AI auto-fix functionality for MongoDB, Elasticsearch, and Postgres schema changes
1 parent df1adbf commit 13879f8

4 files changed

Lines changed: 460 additions & 4 deletions

File tree

‎backend/test/ava-tests/non-saas-tests/non-saas-table-schema-dynamodb-e2e.test.ts‎

Lines changed: 107 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -47,6 +47,7 @@ interface DynamoProposedChange {
4747
}
4848

4949
let nextProposal: DynamoProposedChange | null = null;
50+
let nextFixProposal: DynamoProposedChange | null = null;
5051

5152
function createDynamoProposalStream(proposal: DynamoProposedChange) {
5253
return {
@@ -64,8 +65,21 @@ function createDynamoProposalStream(proposal: DynamoProposedChange) {
6465
};
6566
}
6667

68+
function isFixCall(messages: unknown): boolean {
69+
if (!Array.isArray(messages)) return false;
70+
for (const m of messages) {
71+
const content = (m as { content?: unknown })?.content;
72+
if (typeof content === 'string' && content.includes('DDL repair assistant')) return true;
73+
}
74+
return false;
75+
}
76+
6777
const mockAICoreService = {
68-
streamChatWithToolsAndProvider: async () => {
78+
streamChatWithToolsAndProvider: async (_provider: unknown, messages: unknown) => {
79+
if (isFixCall(messages)) {
80+
if (!nextFixProposal) throw new Error('Test invoked the AI fix loop but nextFixProposal was not set.');
81+
return createDynamoProposalStream(nextFixProposal);
82+
}
6983
if (!nextProposal) throw new Error('Test must set nextProposal.');
7084
return createDynamoProposalStream(nextProposal);
7185
},
@@ -689,3 +703,95 @@ test.serial('DynamoDB: runtime failure marks FAILED and attempts auto-rollback',
689703
t.true(record.autoRollbackAttempted);
690704
t.truthy(record.executionError);
691705
});
706+
707+
test.serial('DynamoDB: runtime failure is repaired by AI auto-fix and applied', async (t) => {
708+
const { token } = await registerUserAndReturnUserInfo(app);
709+
const connectionId = await createConnection(token);
710+
const tableName = randomTableName('ra_ddb_fix');
711+
createdTables.push(tableName);
712+
713+
await seedTable(tableName, [{ AttributeName: 'id', KeyType: 'HASH' }], [{ AttributeName: 'id', AttributeType: 'S' }]);
714+
715+
const fixedIndexName = 'gsi_fixed_email';
716+
const originalForwardOp = JSON.stringify({
717+
operation: 'updateTable',
718+
tableName,
719+
globalSecondaryIndexUpdates: [{ delete: { indexName: 'nonexistent_gsi' } }],
720+
});
721+
const originalRollbackOp = JSON.stringify({
722+
operation: 'updateTable',
723+
tableName,
724+
attributeDefinitions: [
725+
{ attributeName: 'id', attributeType: 'S' },
726+
{ attributeName: 'email', attributeType: 'S' },
727+
],
728+
globalSecondaryIndexUpdates: [
729+
{
730+
create: {
731+
indexName: 'nonexistent_gsi',
732+
keySchema: [{ attributeName: 'email', keyType: 'HASH' }],
733+
projection: { projectionType: 'ALL' },
734+
},
735+
},
736+
],
737+
});
738+
const fixedForwardOp = JSON.stringify({
739+
operation: 'updateTable',
740+
tableName,
741+
attributeDefinitions: [
742+
{ attributeName: 'id', attributeType: 'S' },
743+
{ attributeName: 'email', attributeType: 'S' },
744+
],
745+
globalSecondaryIndexUpdates: [
746+
{
747+
create: {
748+
indexName: fixedIndexName,
749+
keySchema: [{ attributeName: 'email', keyType: 'HASH' }],
750+
projection: { projectionType: 'ALL' },
751+
},
752+
},
753+
],
754+
});
755+
const fixedRollbackOp = JSON.stringify({
756+
operation: 'updateTable',
757+
tableName,
758+
globalSecondaryIndexUpdates: [{ delete: { indexName: fixedIndexName } }],
759+
});
760+
761+
nextProposal = {
762+
forwardOp: originalForwardOp,
763+
rollbackOp: originalRollbackOp,
764+
changeType: SchemaChangeTypeEnum.DYNAMODB_UPDATE_TABLE,
765+
targetTableName: tableName,
766+
isReversible: true,
767+
summary: 'delete missing GSI',
768+
reasoning: '',
769+
};
770+
nextFixProposal = {
771+
forwardOp: fixedForwardOp,
772+
rollbackOp: fixedRollbackOp,
773+
changeType: SchemaChangeTypeEnum.DYNAMODB_UPDATE_TABLE,
774+
targetTableName: tableName,
775+
isReversible: true,
776+
summary: 'AI repaired to create a real GSI',
777+
reasoning: '',
778+
};
779+
780+
const generateResp = await request(app.getHttpServer())
781+
.post(`/table-schema/${connectionId}/generate`)
782+
.set('Cookie', token)
783+
.send({ userPrompt: 'update the table' });
784+
const changeId = JSON.parse(generateResp.text).changes[0].id;
785+
786+
const approveResp = await request(app.getHttpServer())
787+
.post(`/table-schema/change/${changeId}/approve`)
788+
.set('Cookie', token)
789+
.send({});
790+
t.is(approveResp.status, 200);
791+
const applied = JSON.parse(approveResp.text);
792+
t.is(applied.status, SchemaChangeStatusEnum.APPLIED);
793+
t.true(applied.aiAutoFixApplied);
794+
t.truthy(applied.aiAutoFixOriginalError);
795+
t.is(applied.aiAutoFixOriginalForwardSql, originalForwardOp);
796+
t.is(applied.forwardSql, fixedForwardOp);
797+
});

‎backend/test/ava-tests/non-saas-tests/non-saas-table-schema-elasticsearch-e2e.test.ts‎

Lines changed: 78 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -39,6 +39,7 @@ interface ElasticProposedChange {
3939
}
4040

4141
let nextProposal: ElasticProposedChange | null = null;
42+
let nextFixProposal: ElasticProposedChange | null = null;
4243

4344
function createElasticProposalStream(proposal: ElasticProposedChange) {
4445
return {
@@ -56,8 +57,21 @@ function createElasticProposalStream(proposal: ElasticProposedChange) {
5657
};
5758
}
5859

60+
function isFixCall(messages: unknown): boolean {
61+
if (!Array.isArray(messages)) return false;
62+
for (const m of messages) {
63+
const content = (m as { content?: unknown })?.content;
64+
if (typeof content === 'string' && content.includes('DDL repair assistant')) return true;
65+
}
66+
return false;
67+
}
68+
5969
const mockAICoreService = {
60-
streamChatWithToolsAndProvider: async () => {
70+
streamChatWithToolsAndProvider: async (_provider: unknown, messages: unknown) => {
71+
if (isFixCall(messages)) {
72+
if (!nextFixProposal) throw new Error('Test invoked the AI fix loop but nextFixProposal was not set.');
73+
return createElasticProposalStream(nextFixProposal);
74+
}
6175
if (!nextProposal) throw new Error('Test must set nextProposal.');
6276
return createElasticProposalStream(nextProposal);
6377
},
@@ -386,6 +400,69 @@ test.serial('Elasticsearch: invalid op marks FAILED and attempts auto-rollback',
386400
t.truthy(record.executionError);
387401
});
388402

403+
test.serial('Elasticsearch: invalid op is repaired by AI auto-fix and applied', async (t) => {
404+
const { token } = await registerUserAndReturnUserInfo(app);
405+
const connectionId = await createConnection(token);
406+
const indexName = randomIndexName('ra_es_fix');
407+
createdIndices.push(indexName);
408+
409+
await createIndexWithMapping(indexName, { properties: { name: { type: 'keyword' } } });
410+
411+
const originalForwardOp = JSON.stringify({
412+
operation: 'updateMapping',
413+
indexName,
414+
properties: { bad_field: { type: 'totally_made_up_type' } },
415+
});
416+
const originalRollbackOp = JSON.stringify({ operation: 'deleteIndex', indexName });
417+
const fixedForwardOp = JSON.stringify({
418+
operation: 'updateMapping',
419+
indexName,
420+
properties: { phone: { type: 'keyword' } },
421+
});
422+
const fixedRollbackOp = JSON.stringify({ operation: 'deleteIndex', indexName });
423+
424+
nextProposal = {
425+
forwardOp: originalForwardOp,
426+
rollbackOp: originalRollbackOp,
427+
changeType: SchemaChangeTypeEnum.ELASTICSEARCH_UPDATE_MAPPING,
428+
targetTableName: indexName,
429+
isReversible: false,
430+
summary: 'mapping with bogus type',
431+
reasoning: '',
432+
};
433+
nextFixProposal = {
434+
forwardOp: fixedForwardOp,
435+
rollbackOp: fixedRollbackOp,
436+
changeType: SchemaChangeTypeEnum.ELASTICSEARCH_UPDATE_MAPPING,
437+
targetTableName: indexName,
438+
isReversible: false,
439+
summary: 'AI repaired the bogus type to keyword',
440+
reasoning: '',
441+
};
442+
443+
const generateResp = await request(app.getHttpServer())
444+
.post(`/table-schema/${connectionId}/generate`)
445+
.set('Cookie', token)
446+
.send({ userPrompt: 'add a field' });
447+
const changeId = JSON.parse(generateResp.text).changes[0].id;
448+
449+
const approveResp = await request(app.getHttpServer())
450+
.post(`/table-schema/change/${changeId}/approve`)
451+
.set('Cookie', token)
452+
.send({ confirmedDestructive: true });
453+
t.is(approveResp.status, 200);
454+
const applied = JSON.parse(approveResp.text);
455+
t.is(applied.status, SchemaChangeStatusEnum.APPLIED);
456+
t.true(applied.aiAutoFixApplied);
457+
t.truthy(applied.aiAutoFixOriginalError);
458+
t.is(applied.aiAutoFixOriginalForwardSql, originalForwardOp);
459+
t.is(applied.forwardSql, fixedForwardOp);
460+
461+
const mapping = await getMapping(indexName);
462+
const properties = (mapping?.properties as Record<string, { type?: string }>) ?? {};
463+
t.is(properties.phone?.type, 'keyword');
464+
});
465+
389466
test.serial('Elasticsearch: userModifiedSql JSON op is validated and applied', async (t) => {
390467
const { token } = await registerUserAndReturnUserInfo(app);
391468
const connectionId = await createConnection(token);

‎backend/test/ava-tests/non-saas-tests/non-saas-table-schema-mongodb-e2e.test.ts‎

Lines changed: 85 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -39,6 +39,7 @@ interface MongoProposedChange {
3939
}
4040

4141
let nextProposal: MongoProposedChange | null = null;
42+
let nextFixProposal: MongoProposedChange | null = null;
4243

4344
function createMongoProposalStream(proposal: MongoProposedChange) {
4445
return {
@@ -56,8 +57,21 @@ function createMongoProposalStream(proposal: MongoProposedChange) {
5657
};
5758
}
5859

60+
function isFixCall(messages: unknown): boolean {
61+
if (!Array.isArray(messages)) return false;
62+
for (const m of messages) {
63+
const content = (m as { content?: unknown })?.content;
64+
if (typeof content === 'string' && content.includes('DDL repair assistant')) return true;
65+
}
66+
return false;
67+
}
68+
5969
const mockAICoreService = {
60-
streamChatWithToolsAndProvider: async () => {
70+
streamChatWithToolsAndProvider: async (_provider: unknown, messages: unknown) => {
71+
if (isFixCall(messages)) {
72+
if (!nextFixProposal) throw new Error('Test invoked the AI fix loop but nextFixProposal was not set.');
73+
return createMongoProposalStream(nextFixProposal);
74+
}
6175
if (!nextProposal) throw new Error('Test must set nextProposal.');
6276
return createMongoProposalStream(nextProposal);
6377
},
@@ -557,6 +571,76 @@ test.serial('MongoDB: invalid op marks FAILED and attempts auto-rollback', async
557571
t.truthy(record.executionError);
558572
});
559573

574+
test.serial('MongoDB: invalid op is repaired by AI auto-fix and applied', async (t) => {
575+
const { token } = await registerUserAndReturnUserInfo(app);
576+
const connectionId = await createConnection(token);
577+
const collectionName = randomCollectionName('ra_mfix');
578+
createdCollections.push(collectionName);
579+
580+
const mongoParams = getTestData(mockFactory).mongoDbConnection;
581+
await seedCollection(mongoParams, collectionName, [{ email: 'a@a.test' }]);
582+
await withMongoClient(mongoParams, async (db) => {
583+
await db.collection(collectionName).createIndex({ email: 1 }, { name: 'idx_real' });
584+
});
585+
586+
const originalForwardOp = JSON.stringify({ operation: 'dropIndex', collectionName, indexName: 'idx_missing' });
587+
const originalRollbackOp = JSON.stringify({
588+
operation: 'createIndex',
589+
collectionName,
590+
indexName: 'idx_missing',
591+
indexSpec: { x: 1 },
592+
indexOptions: { name: 'idx_missing' },
593+
});
594+
const fixedForwardOp = JSON.stringify({ operation: 'dropIndex', collectionName, indexName: 'idx_real' });
595+
const fixedRollbackOp = JSON.stringify({
596+
operation: 'createIndex',
597+
collectionName,
598+
indexName: 'idx_real',
599+
indexSpec: { email: 1 },
600+
indexOptions: { name: 'idx_real' },
601+
});
602+
603+
nextProposal = {
604+
forwardOp: originalForwardOp,
605+
rollbackOp: originalRollbackOp,
606+
changeType: SchemaChangeTypeEnum.MONGO_DROP_INDEX,
607+
targetTableName: collectionName,
608+
isReversible: true,
609+
summary: 'drop wrong index',
610+
reasoning: '',
611+
};
612+
nextFixProposal = {
613+
forwardOp: fixedForwardOp,
614+
rollbackOp: fixedRollbackOp,
615+
changeType: SchemaChangeTypeEnum.MONGO_DROP_INDEX,
616+
targetTableName: collectionName,
617+
isReversible: true,
618+
summary: 'corrected index name',
619+
reasoning: 'AI repaired indexName',
620+
};
621+
622+
const generateResp = await request(app.getHttpServer())
623+
.post(`/table-schema/${connectionId}/generate`)
624+
.set('Cookie', token)
625+
.send({ userPrompt: 'drop the email index' });
626+
const changeId = JSON.parse(generateResp.text).changes[0].id;
627+
628+
const approveResp = await request(app.getHttpServer())
629+
.post(`/table-schema/change/${changeId}/approve`)
630+
.set('Cookie', token)
631+
.send({});
632+
t.is(approveResp.status, 200);
633+
const applied = JSON.parse(approveResp.text);
634+
t.is(applied.status, SchemaChangeStatusEnum.APPLIED);
635+
t.true(applied.aiAutoFixApplied);
636+
t.truthy(applied.aiAutoFixOriginalError);
637+
t.is(applied.aiAutoFixOriginalForwardSql, originalForwardOp);
638+
t.is(applied.forwardSql, fixedForwardOp);
639+
640+
const indexes = await getIndexes(mongoParams, collectionName);
641+
t.falsy(indexes.find((idx) => idx.name === 'idx_real'));
642+
});
643+
560644
test.serial('MongoDB: tool/changeType mismatch is rejected', async (t) => {
561645
const { token } = await registerUserAndReturnUserInfo(app);
562646
const connectionId = await createConnection(token);

0 commit comments

Comments
 (0)