gongxue-base/apps/server/src/migrations/1784860000000-EnhanceAiChatForAntDesignX.ts

import {
  MigrationInterface,
  QueryRunner,
  Table,
  TableColumn,
  TableForeignKey,
  TableIndex,
} from 'typeorm';

export class EnhanceAiChatForAntDesignX1784860000000 implements MigrationInterface {
  async up(queryRunner: QueryRunner): Promise<void> {
    await this.addColumn(queryRunner, 'ai_config', new TableColumn({
      name: 'supports_vision',
      type: 'boolean',
      default: false,
    }));
    await this.addColumn(queryRunner, 'ai_conversations', new TableColumn({
      name: 'locked_skill_key',
      type: 'varchar',
      length: '50',
      isNullable: true,
    }));
    await this.addColumn(queryRunner, 'ai_messages', new TableColumn({
      name: 'reply_to_message_id',
      type: 'integer',
      isNullable: true,
    }));
    await this.addColumn(queryRunner, 'ai_messages', new TableColumn({
      name: 'feedback',
      type: 'varchar',
      length: '20',
      isNullable: true,
    }));
    await this.addColumn(queryRunner, 'ai_messages', new TableColumn({
      name: 'feedback_reason',
      type: 'varchar',
      length: '500',
      isNullable: true,
    }));
    await this.addColumn(queryRunner, 'ai_messages', new TableColumn({
      name: 'metadata',
      type: 'text',
      isNullable: true,
    }));
    await this.addColumn(queryRunner, 'ai_tool_runs', new TableColumn({
      name: 'skill_key',
      type: 'varchar',
      length: '50',
      isNullable: true,
    }));
    await this.addColumn(queryRunner, 'ai_tool_runs', new TableColumn({
      name: 'arguments_data',
      type: 'text',
      isNullable: true,
    }));
    await this.addColumn(queryRunner, 'ai_tool_runs', new TableColumn({
      name: 'result_data',
      type: 'text',
      isNullable: true,
    }));

    const messagesTable = await queryRunner.getTable('ai_messages');
    if (
      messagesTable &&
      !messagesTable.foreignKeys.some((key) => key.name === 'fk_ai_messages_reply_to')
    ) {
      await queryRunner.createForeignKey(
        'ai_messages',
        new TableForeignKey({
          name: 'fk_ai_messages_reply_to',
          columnNames: ['reply_to_message_id'],
          referencedTableName: 'ai_messages',
          referencedColumnNames: ['id'],
          onDelete: 'SET NULL',
        }),
      );
    }

    if (!(await queryRunner.hasTable('ai_attachments'))) {
      await queryRunner.createTable(
        new Table({
          name: 'ai_attachments',
          columns: [
            { name: 'id', type: 'integer', isPrimary: true, isGenerated: true, generationStrategy: 'increment' },
            { name: 'user_id', type: 'integer' },
            { name: 'original_name', type: 'varchar', length: '255' },
            { name: 'mime_type', type: 'varchar', length: '100' },
            { name: 'size', type: 'integer' },
            { name: 'storage_key', type: 'varchar', length: '255', isUnique: true },
            { name: 'processing_status', type: 'varchar', length: '20', default: "'processing'" },
            { name: 'extracted_text', type: 'text', isNullable: true },
            { name: 'processing_error', type: 'varchar', length: '200', isNullable: true },
            { name: 'image_width', type: 'integer', isNullable: true },
            { name: 'image_height', type: 'integer', isNullable: true },
            { name: 'created_at', type: 'datetime', default: 'CURRENT_TIMESTAMP' },
            { name: 'updated_at', type: 'datetime', default: 'CURRENT_TIMESTAMP' },
          ],
          indices: [
            { name: 'idx_ai_attachments_user_created', columnNames: ['user_id', 'created_at'] },
          ],
          foreignKeys: [
            {
              name: 'fk_ai_attachments_user',
              columnNames: ['user_id'],
              referencedTableName: 'users',
              referencedColumnNames: ['id'],
              onDelete: 'CASCADE',
            },
          ],
        }),
      );
    }

    if (!(await queryRunner.hasTable('ai_message_attachments'))) {
      await queryRunner.createTable(
        new Table({
          name: 'ai_message_attachments',
          columns: [
            { name: 'message_id', type: 'integer', isPrimary: true },
            { name: 'attachment_id', type: 'integer', isPrimary: true },
          ],
          foreignKeys: [
            {
              name: 'fk_ai_message_attachments_attachment',
              columnNames: ['attachment_id'],
              referencedTableName: 'ai_attachments',
              referencedColumnNames: ['id'],
              onDelete: 'CASCADE',
            },
            {
              name: 'fk_ai_message_attachments_message',
              columnNames: ['message_id'],
              referencedTableName: 'ai_messages',
              referencedColumnNames: ['id'],
              onDelete: 'CASCADE',
            },
          ],
        }),
      );
      await queryRunner.createIndex(
        'ai_message_attachments',
        new TableIndex({
          name: 'idx_ai_message_attachments_message',
          columnNames: ['message_id'],
        }),
      );
    }
  }

  async down(queryRunner: QueryRunner): Promise<void> {
    for (const table of ['ai_message_attachments', 'ai_attachments']) {
      if (await queryRunner.hasTable(table)) await queryRunner.dropTable(table);
    }
    const messagesTable = await queryRunner.getTable('ai_messages');
    const replyForeignKey = messagesTable?.foreignKeys.find(
      (key) => key.name === 'fk_ai_messages_reply_to',
    );
    if (replyForeignKey) await queryRunner.dropForeignKey('ai_messages', replyForeignKey);
    const columns: Array<[string, string]> = [
      ['ai_tool_runs', 'result_data'],
      ['ai_tool_runs', 'arguments_data'],
      ['ai_tool_runs', 'skill_key'],
      ['ai_messages', 'metadata'],
      ['ai_messages', 'feedback_reason'],
      ['ai_messages', 'feedback'],
      ['ai_messages', 'reply_to_message_id'],
      ['ai_conversations', 'locked_skill_key'],
      ['ai_config', 'supports_vision'],
    ];
    for (const [table, column] of columns) {
      if (await queryRunner.hasColumn(table, column)) await queryRunner.dropColumn(table, column);
    }
  }

  private async addColumn(
    queryRunner: QueryRunner,
    table: string,
    column: TableColumn,
  ): Promise<void> {
    if ((await queryRunner.hasTable(table)) && !(await queryRunner.hasColumn(table, column.name))) {
      await queryRunner.addColumn(table, column);
    }
  }
}