Merge pull request #1770 from rocket-admin/backend_ai_table_schema_context

Artuomka · web-flow · commit 99f2a58e64f4 · 2026-05-14T17:09:27.000+03:00
feat: implement schema change chat functionality with message persistence and thread management
diff --git a/backend/src/common/application/global-database-context.interface.ts b/backend/src/common/application/global-database-context.interface.ts
@@ -37,6 +37,10 @@ import { TableFiltersEntity } from '../../entities/table-filters/table-filters.e
 import { TableInfoEntity } from '../../entities/table-info/table-info.entity.js';
 import { ITableLogsRepository } from '../../entities/table-logs/repository/table-logs-repository.interface.js';
 import { ITableSchemaChangeRepository } from '../../entities/table-schema/repository/table-schema-change.repository.interface.js';
+import { ISchemaChangeChatRepository } from '../../entities/table-schema/schema-change-chat/schema-change-chat/repository/schema-change-chat-repository.interface.js';
+import { SchemaChangeChatEntity } from '../../entities/table-schema/schema-change-chat/schema-change-chat/schema-change-chat.entity.js';
+import { ISchemaChangeChatMessageRepository } from '../../entities/table-schema/schema-change-chat/schema-change-chat-message/repository/schema-change-chat-message-repository.interface.js';
+import { SchemaChangeChatMessageEntity } from '../../entities/table-schema/schema-change-chat/schema-change-chat-message/schema-change-chat-message.entity.js';
 import { TableSchemaChangeEntity } from '../../entities/table-schema/table-schema-change.entity.js';
 import { ITableSettingsRepository } from '../../entities/table-settings/common-table-settings/repository/table-settings.repository.interface.js';
 import { TableSettingsEntity } from '../../entities/table-settings/common-table-settings/table-settings.entity.js';
@@ -106,4 +110,6 @@ export interface IGlobalDatabaseContext extends IDatabaseContext {
 	userAiChatRepository: Repository<UserAiChatEntity> & IUserAiChatRepository;
 	aiChatMessageRepository: Repository<AiChatMessageEntity> & IAiChatMessageRepository;
 	tableSchemaChangeRepository: Repository<TableSchemaChangeEntity> & ITableSchemaChangeRepository;
+	schemaChangeChatRepository: Repository<SchemaChangeChatEntity> & ISchemaChangeChatRepository;
+	schemaChangeChatMessageRepository: Repository<SchemaChangeChatMessageEntity> & ISchemaChangeChatMessageRepository;
 }
diff --git a/backend/src/common/application/global-database-context.ts b/backend/src/common/application/global-database-context.ts
@@ -64,6 +64,12 @@ import { ITableLogsRepository } from '../../entities/table-logs/repository/table
 import { TableLogsEntity } from '../../entities/table-logs/table-logs.entity.js';
 import { customTableSchemaChangeRepositoryExtension } from '../../entities/table-schema/repository/custom-table-schema-change-repository-extension.js';
 import { ITableSchemaChangeRepository } from '../../entities/table-schema/repository/table-schema-change.repository.interface.js';
+import { schemaChangeChatRepositoryExtension } from '../../entities/table-schema/schema-change-chat/schema-change-chat/repository/schema-change-chat-repository.extension.js';
+import { ISchemaChangeChatRepository } from '../../entities/table-schema/schema-change-chat/schema-change-chat/repository/schema-change-chat-repository.interface.js';
+import { SchemaChangeChatEntity } from '../../entities/table-schema/schema-change-chat/schema-change-chat/schema-change-chat.entity.js';
+import { schemaChangeChatMessageRepositoryExtension } from '../../entities/table-schema/schema-change-chat/schema-change-chat-message/repository/schema-change-chat-message-repository.extension.js';
+import { ISchemaChangeChatMessageRepository } from '../../entities/table-schema/schema-change-chat/schema-change-chat-message/repository/schema-change-chat-message-repository.interface.js';
+import { SchemaChangeChatMessageEntity } from '../../entities/table-schema/schema-change-chat/schema-change-chat-message/schema-change-chat-message.entity.js';
 import { TableSchemaChangeEntity } from '../../entities/table-schema/table-schema-change.entity.js';
 import { ITableSettingsRepository } from '../../entities/table-settings/common-table-settings/repository/table-settings.repository.interface.js';
 import { tableSettingsCustomRepositoryExtension } from '../../entities/table-settings/common-table-settings/repository/table-settings-custom-repository-extension.js';
@@ -157,6 +163,9 @@ export class GlobalDatabaseContext implements IGlobalDatabaseContext {
 	private _userAiChatRepository: Repository<UserAiChatEntity> & IUserAiChatRepository;
 	private _aiChatMessageRepository: Repository<AiChatMessageEntity> & IAiChatMessageRepository;
 	private _tableSchemaChangeRepository: Repository<TableSchemaChangeEntity> & ITableSchemaChangeRepository;
+	private _schemaChangeChatRepository: Repository<SchemaChangeChatEntity> & ISchemaChangeChatRepository;
+	private _schemaChangeChatMessageRepository: Repository<SchemaChangeChatMessageEntity> &
+		ISchemaChangeChatMessageRepository;
 
 	public constructor(
 		@Inject(BaseType.DATA_SOURCE)
@@ -264,6 +273,12 @@ export class GlobalDatabaseContext implements IGlobalDatabaseContext {
 		this._tableSchemaChangeRepository = this.appDataSource
 			.getRepository(TableSchemaChangeEntity)
 			.extend(customTableSchemaChangeRepositoryExtension);
+		this._schemaChangeChatRepository = this.appDataSource
+			.getRepository(SchemaChangeChatEntity)
+			.extend(schemaChangeChatRepositoryExtension);
+		this._schemaChangeChatMessageRepository = this.appDataSource
+			.getRepository(SchemaChangeChatMessageEntity)
+			.extend(schemaChangeChatMessageRepositoryExtension);
 	}
 
 	public get userRepository(): Repository<UserEntity> & IUserRepository {
@@ -428,6 +443,15 @@ export class GlobalDatabaseContext implements IGlobalDatabaseContext {
 		return this._tableSchemaChangeRepository;
 	}
 
+	public get schemaChangeChatRepository(): Repository<SchemaChangeChatEntity> & ISchemaChangeChatRepository {
+		return this._schemaChangeChatRepository;
+	}
+
+	public get schemaChangeChatMessageRepository(): Repository<SchemaChangeChatMessageEntity> &
+		ISchemaChangeChatMessageRepository {
+		return this._schemaChangeChatMessageRepository;
+	}
+
 	public startTransaction(): Promise<void> {
 		this._queryRunner = this.appDataSource.createQueryRunner();
 		this._queryRunner.startTransaction();
diff --git a/backend/src/entities/table-schema/application/data-structures/generate-schema-change.ds.ts b/backend/src/entities/table-schema/application/data-structures/generate-schema-change.ds.ts
@@ -3,4 +3,5 @@ export class GenerateSchemaChangeDs {
 	userPrompt: string;
 	userId: string;
 	masterPassword?: string;
+	threadId?: string | null;
 }
diff --git a/backend/src/entities/table-schema/application/data-transfer-objects/generate-schema-change.dto.ts b/backend/src/entities/table-schema/application/data-transfer-objects/generate-schema-change.dto.ts
@@ -1,5 +1,5 @@
 import { ApiProperty } from '@nestjs/swagger';
-import { IsNotEmpty, IsString, MaxLength, MinLength } from 'class-validator';
+import { IsNotEmpty, IsOptional, IsString, IsUUID, MaxLength, MinLength } from 'class-validator';
 
 export class GenerateSchemaChangeDto {
 	@ApiProperty({
@@ -15,4 +15,15 @@ export class GenerateSchemaChangeDto {
 	@MinLength(1)
 	@MaxLength(2000)
 	userPrompt: string;
+
+	@ApiProperty({
+		type: String,
+		required: false,
+		nullable: true,
+		description:
+			'Optional thread ID to continue an existing conversation. When supplied, prior turns are prepended to the AI prompt, giving the model context for iterative refinement (e.g. "now also add an index", "rename it to created_at"). Omit to start a fresh thread; the returned threadId can be passed back on the next call.',
+	})
+	@IsOptional()
+	@IsUUID()
+	threadId?: string;
 }
diff --git a/backend/src/entities/table-schema/application/data-transfer-objects/schema-change-batch-response.dto.ts b/backend/src/entities/table-schema/application/data-transfer-objects/schema-change-batch-response.dto.ts
@@ -13,4 +13,12 @@ export class SchemaChangeBatchResponseDto {
 		description: 'Generated changes ordered by orderInBatch (dependency order — parents first).',
 	})
 	changes: SchemaChangeResponseDto[];
+
+	@ApiProperty({
+		required: false,
+		nullable: true,
+		description:
+			'Conversation thread ID. Present when the request used or created a chat thread. Pass it back as the threadId query param on subsequent generate calls to continue the conversation with full prior context.',
+	})
+	threadId?: string | null;
 }
diff --git a/backend/src/entities/table-schema/schema-change-chat/schema-change-chat-message/repository/schema-change-chat-message-repository.extension.ts b/backend/src/entities/table-schema/schema-change-chat/schema-change-chat-message/repository/schema-change-chat-message-repository.extension.ts
@@ -0,0 +1,35 @@
+import { MessageRole } from '../../../../ai/ai-conversation-history/ai-chat-messages/message-role.enum.js';
+import { SchemaChangeChatMessageEntity } from '../schema-change-chat-message.entity.js';
+import { ISchemaChangeChatMessageRepository } from './schema-change-chat-message-repository.interface.js';
+
+export const schemaChangeChatMessageRepositoryExtension: ISchemaChangeChatMessageRepository = {
+	async findMessagesForChat(chatId: string): Promise<SchemaChangeChatMessageEntity[]> {
+		return await this.createQueryBuilder('schema_change_chat_message')
+			.where('schema_change_chat_message.chat_id = :chatId', { chatId })
+			.orderBy('schema_change_chat_message.created_at', 'ASC')
+			.getMany();
+	},
+
+	async deleteMessagesForChat(chatId: string): Promise<void> {
+		await this.createQueryBuilder()
+			.delete()
+			.from('schema_change_chat_message')
+			.where('chat_id = :chatId', { chatId })
+			.execute();
+	},
+
+	async saveMessage(
+		chatId: string,
+		message: string,
+		role: MessageRole,
+		batchId?: string | null,
+	): Promise<SchemaChangeChatMessageEntity> {
+		const newMessage = this.create({
+			chat_id: chatId,
+			message,
+			role,
+			batch_id: batchId ?? null,
+		});
+		return await this.save(newMessage);
+	},
+};
diff --git a/backend/src/entities/table-schema/schema-change-chat/schema-change-chat-message/repository/schema-change-chat-message-repository.interface.ts b/backend/src/entities/table-schema/schema-change-chat/schema-change-chat-message/repository/schema-change-chat-message-repository.interface.ts
@@ -0,0 +1,13 @@
+import { MessageRole } from '../../../../ai/ai-conversation-history/ai-chat-messages/message-role.enum.js';
+import { SchemaChangeChatMessageEntity } from '../schema-change-chat-message.entity.js';
+
+export interface ISchemaChangeChatMessageRepository {
+	findMessagesForChat(chatId: string): Promise<SchemaChangeChatMessageEntity[]>;
+	deleteMessagesForChat(chatId: string): Promise<void>;
+	saveMessage(
+		chatId: string,
+		message: string,
+		role: MessageRole,
+		batchId?: string | null,
+	): Promise<SchemaChangeChatMessageEntity>;
+}
diff --git a/backend/src/entities/table-schema/schema-change-chat/schema-change-chat-message/schema-change-chat-message.entity.ts b/backend/src/entities/table-schema/schema-change-chat/schema-change-chat-message/schema-change-chat-message.entity.ts
@@ -0,0 +1,44 @@
+import {
+	Column,
+	CreateDateColumn,
+	Entity,
+	JoinColumn,
+	ManyToOne,
+	PrimaryGeneratedColumn,
+	Relation,
+	UpdateDateColumn,
+} from 'typeorm';
+import { MessageRole } from '../../../ai/ai-conversation-history/ai-chat-messages/message-role.enum.js';
+import { SchemaChangeChatEntity } from '../schema-change-chat/schema-change-chat.entity.js';
+
+@Entity('schema_change_chat_message')
+export class SchemaChangeChatMessageEntity {
+	@PrimaryGeneratedColumn('uuid')
+	id: string;
+
+	@Column({ default: null, type: 'text' })
+	message: string;
+
+	@Column({ nullable: true, default: null, type: 'enum', enum: MessageRole })
+	role: MessageRole;
+
+	@Column({ type: 'uuid', nullable: true, default: null })
+	batch_id: string | null;
+
+	@CreateDateColumn({ type: 'timestamp', default: () => 'CURRENT_TIMESTAMP' })
+	created_at: Date;
+
+	@UpdateDateColumn({ type: 'timestamp', nullable: true, default: null })
+	updated_at: Date;
+
+	@ManyToOne(
+		() => SchemaChangeChatEntity,
+		(chat) => chat.messages,
+		{ onDelete: 'CASCADE' },
+	)
+	@JoinColumn({ name: 'chat_id' })
+	chat: Relation<SchemaChangeChatEntity>;
+
+	@Column()
+	chat_id: string;
+}
diff --git a/backend/src/entities/table-schema/schema-change-chat/schema-change-chat/repository/schema-change-chat-repository.extension.ts b/backend/src/entities/table-schema/schema-change-chat/schema-change-chat/repository/schema-change-chat-repository.extension.ts
@@ -0,0 +1,53 @@
+import { SchemaChangeChatEntity } from '../schema-change-chat.entity.js';
+import { ISchemaChangeChatRepository } from './schema-change-chat-repository.interface.js';
+
+export const schemaChangeChatRepositoryExtension: ISchemaChangeChatRepository = {
+	async findChatByIdAndUserId(chatId: string, userId: string): Promise<SchemaChangeChatEntity | null> {
+		return await this.createQueryBuilder('schema_change_chat')
+			.where('schema_change_chat.id = :chatId', { chatId })
+			.andWhere('schema_change_chat.user_id = :userId', { userId })
+			.getOne();
+	},
+
+	async findChatWithMessagesByIdAndUserId(chatId: string, userId: string): Promise<SchemaChangeChatEntity | null> {
+		return await this.createQueryBuilder('schema_change_chat')
+			.leftJoinAndSelect('schema_change_chat.messages', 'messages')
+			.where('schema_change_chat.id = :chatId', { chatId })
+			.andWhere('schema_change_chat.user_id = :userId', { userId })
+			.orderBy('messages.created_at', 'ASC')
+			.getOne();
+	},
+
+	async findChatsForConnection(connectionId: string, userId: string): Promise<SchemaChangeChatEntity[]> {
+		return await this.createQueryBuilder('schema_change_chat')
+			.where('schema_change_chat.connection_id = :connectionId', { connectionId })
+			.andWhere('schema_change_chat.user_id = :userId', { userId })
+			.orderBy('schema_change_chat.created_at', 'DESC')
+			.getMany();
+	},
+
+	async createChatForUser(userId: string, connectionId: string, name?: string): Promise<SchemaChangeChatEntity> {
+		const newChat = this.create({
+			user_id: userId,
+			connection_id: connectionId,
+			name: name || null,
+		});
+		return await this.save(newChat);
+	},
+
+	async updateChatName(chatId: string, name: string): Promise<void> {
+		await this.createQueryBuilder()
+			.update(SchemaChangeChatEntity)
+			.set({ name })
+			.where('id = :chatId', { chatId })
+			.execute();
+	},
+
+	async updateLastBatchId(chatId: string, batchId: string): Promise<void> {
+		await this.createQueryBuilder()
+			.update(SchemaChangeChatEntity)
+			.set({ last_batch_id: batchId })
+			.where('id = :chatId', { chatId })
+			.execute();
+	},
+};
diff --git a/backend/src/entities/table-schema/schema-change-chat/schema-change-chat/repository/schema-change-chat-repository.interface.ts b/backend/src/entities/table-schema/schema-change-chat/schema-change-chat/repository/schema-change-chat-repository.interface.ts
@@ -0,0 +1,10 @@
+import { SchemaChangeChatEntity } from '../schema-change-chat.entity.js';
+
+export interface ISchemaChangeChatRepository {
+	findChatByIdAndUserId(chatId: string, userId: string): Promise<SchemaChangeChatEntity | null>;
+	findChatWithMessagesByIdAndUserId(chatId: string, userId: string): Promise<SchemaChangeChatEntity | null>;
+	findChatsForConnection(connectionId: string, userId: string): Promise<SchemaChangeChatEntity[]>;
+	createChatForUser(userId: string, connectionId: string, name?: string): Promise<SchemaChangeChatEntity>;
+	updateChatName(chatId: string, name: string): Promise<void>;
+	updateLastBatchId(chatId: string, batchId: string): Promise<void>;
+}
diff --git a/backend/src/entities/table-schema/schema-change-chat/schema-change-chat/schema-change-chat.entity.ts b/backend/src/entities/table-schema/schema-change-chat/schema-change-chat/schema-change-chat.entity.ts
@@ -0,0 +1,52 @@
+import {
+	Column,
+	CreateDateColumn,
+	Entity,
+	JoinColumn,
+	ManyToOne,
+	OneToMany,
+	PrimaryGeneratedColumn,
+	Relation,
+	UpdateDateColumn,
+} from 'typeorm';
+import { ConnectionEntity } from '../../../connection/connection.entity.js';
+import { UserEntity } from '../../../user/user.entity.js';
+import { SchemaChangeChatMessageEntity } from '../schema-change-chat-message/schema-change-chat-message.entity.js';
+
+@Entity('schema_change_chat')
+export class SchemaChangeChatEntity {
+	@PrimaryGeneratedColumn('uuid')
+	id: string;
+
+	@Column({ default: null })
+	name: string;
+
+	@CreateDateColumn({ type: 'timestamp', default: () => 'CURRENT_TIMESTAMP' })
+	created_at: Date;
+
+	@UpdateDateColumn({ type: 'timestamp', nullable: true, default: null })
+	updated_at: Date;
+
+	@ManyToOne(() => UserEntity, { onDelete: 'CASCADE' })
+	@JoinColumn({ name: 'user_id' })
+	user: Relation<UserEntity>;
+
+	@Column()
+	user_id: string;
+
+	@ManyToOne(() => ConnectionEntity, { onDelete: 'CASCADE' })
+	@JoinColumn({ name: 'connection_id' })
+	connection: Relation<ConnectionEntity>;
+
+	@Column({ type: 'varchar', length: 38 })
+	connection_id: string;
+
+	@Column({ type: 'uuid', nullable: true, default: null })
+	last_batch_id: string | null;
+
+	@OneToMany(
+		() => SchemaChangeChatMessageEntity,
+		(message) => message.chat,
+	)
+	messages: Relation<SchemaChangeChatMessageEntity>[];
+}
diff --git a/backend/src/entities/table-schema/table-schema.controller.ts b/backend/src/entities/table-schema/table-schema.controller.ts
@@ -71,7 +71,7 @@ export class TableSchemaController {
 
 	@ApiOperation({
 		summary:
-			'Generate one or more schema changes from a natural-language prompt. The response is always a batch envelope; single-change prompts return a length-1 array.',
+			'Generate one or more schema changes from a natural-language prompt. The response is always a batch envelope; single-change prompts return a length-1 array. Pass an optional threadId in the body to continue an existing conversation; the response returns the threadId to use for follow-up turns.',
 	})
 	@ApiParam({ name: 'connectionId', type: String })
 	@ApiBody({ type: GenerateSchemaChangeDto })
@@ -91,6 +91,7 @@ export class TableSchemaController {
 			userPrompt: body.userPrompt,
 			userId,
 			masterPassword,
+			threadId: body.threadId ?? null,
 		});
 	}
 
diff --git a/backend/src/entities/table-schema/table-schema.module.ts b/backend/src/entities/table-schema/table-schema.module.ts
@@ -8,6 +8,8 @@ import { SchemaChangeOwnershipGuard } from '../../guards/schema-change-ownership
 import { ConnectionEntity } from '../connection/connection.entity.js';
 import { LogOutEntity } from '../log-out/log-out.entity.js';
 import { UserEntity } from '../user/user.entity.js';
+import { SchemaChangeChatEntity } from './schema-change-chat/schema-change-chat/schema-change-chat.entity.js';
+import { SchemaChangeChatMessageEntity } from './schema-change-chat/schema-change-chat-message/schema-change-chat-message.entity.js';
 import { TableSchemaController } from './table-schema.controller.js';
 import { TableSchemaChangeEntity } from './table-schema-change.entity.js';
 import { ApproveAndApplySchemaChangeUseCase } from './use-cases/approve-and-apply-schema-change.use-case.js';
@@ -22,7 +24,16 @@ import { RollbackBatchSchemaChangesUseCase } from './use-cases/rollback-batch-sc
 import { RollbackSchemaChangeUseCase } from './use-cases/rollback-schema-change.use-case.js';
 
 @Module({
-	imports: [TypeOrmModule.forFeature([TableSchemaChangeEntity, ConnectionEntity, UserEntity, LogOutEntity])],
+	imports: [
+		TypeOrmModule.forFeature([
+			TableSchemaChangeEntity,
+			SchemaChangeChatEntity,
+			SchemaChangeChatMessageEntity,
+			ConnectionEntity,
+			UserEntity,
+			LogOutEntity,
+		]),
+	],
 	providers: [
 		{
 			provide: BaseType.GLOBAL_DB_CONTEXT,
diff --git a/backend/src/entities/table-schema/use-cases/generate-schema-change.use-case.ts b/backend/src/entities/table-schema/use-cases/generate-schema-change.use-case.ts
diff --git a/backend/src/migrations/1778767036234-AddSchemaChangeChatEntities.ts b/backend/src/migrations/1778767036234-AddSchemaChangeChatEntities.ts
diff --git a/backend/test/ava-tests/non-saas-tests/non-saas-table-schema-postgres-e2e.test.ts b/backend/test/ava-tests/non-saas-tests/non-saas-table-schema-postgres-e2e.test.ts

Original file line number	Diff line number	Diff line change
`@@ -3,4 +3,5 @@ export class GenerateSchemaChangeDs {`
`3`	`3`	`userPrompt: string;`
`4`	`4`	`userId: string;`
`5`	`5`	`masterPassword?: string;`
	`6`	`+ threadId?: string \| null;`
`6`	`7`	`}`