1  from alembic import op
 2  import sqlalchemy as sa
 3  
 4  revision = '026'
 5  down_revision = '025'
 6  branch_labels = None
 7  depends_on = None
 8  
 9  
10  def upgrade():
11      try:
12          op.create_table(
13              'kg_entities',
14              sa.Column('id', sa.Integer(), primary_key=True, index=True),
15              sa.Column('project_id', sa.Integer(), sa.ForeignKey('projects.id'), nullable=False, index=True),
16              sa.Column('name', sa.String(255), nullable=False),
17              sa.Column('normalized', sa.String(255), nullable=False, index=True),
18              sa.Column('entity_type', sa.String(50), nullable=False),
19              sa.Column('mention_count', sa.Integer(), nullable=False, server_default='0'),
20              sa.Column('created_at', sa.DateTime(), nullable=False),
21              sa.Column('updated_at', sa.DateTime(), nullable=False),
22              sa.UniqueConstraint('project_id', 'normalized', 'entity_type', name='uq_kg_entities_project_norm_type'),
23          )
24      except Exception as e:
25          print(f"Error creating kg_entities: {e}")
26  
27      try:
28          op.create_table(
29              'kg_entity_mentions',
30              sa.Column('id', sa.Integer(), primary_key=True, index=True),
31              sa.Column('entity_id', sa.Integer(), sa.ForeignKey('kg_entities.id'), nullable=False, index=True),
32              sa.Column('project_id', sa.Integer(), sa.ForeignKey('projects.id'), nullable=False, index=True),
33              sa.Column('source', sa.String(500), nullable=False, index=True),
34              sa.Column('mention_count', sa.Integer(), nullable=False, server_default='1'),
35              sa.Column('created_at', sa.DateTime(), nullable=False),
36              sa.UniqueConstraint('entity_id', 'source', name='uq_kg_mentions_entity_source'),
37          )
38      except Exception as e:
39          print(f"Error creating kg_entity_mentions: {e}")
40  
41      try:
42          op.create_table(
43              'kg_entity_relationships',
44              sa.Column('id', sa.Integer(), primary_key=True, index=True),
45              sa.Column('project_id', sa.Integer(), sa.ForeignKey('projects.id'), nullable=False, index=True),
46              sa.Column('from_entity_id', sa.Integer(), sa.ForeignKey('kg_entities.id'), nullable=False),
47              sa.Column('to_entity_id', sa.Integer(), sa.ForeignKey('kg_entities.id'), nullable=False),
48              sa.Column('weight', sa.Integer(), nullable=False, server_default='1'),
49              sa.Column('created_at', sa.DateTime(), nullable=False),
50              sa.UniqueConstraint('project_id', 'from_entity_id', 'to_entity_id', name='uq_kg_rel_project_from_to'),
51          )
52      except Exception as e:
53          print(f"Error creating kg_entity_relationships: {e}")
54  
55  
56  def downgrade():
57      for table in ('kg_entity_relationships', 'kg_entity_mentions', 'kg_entities'):
58          try:
59              op.drop_table(table)
60          except Exception as e:
61              print(f"Error dropping {table}: {e}")