Skip to content

Commit

Permalink
fix: join primary key to unique constraint
Browse files Browse the repository at this point in the history
  • Loading branch information
xiaozeyu committed Aug 9, 2024
1 parent 633808d commit 5271612
Show file tree
Hide file tree
Showing 8 changed files with 179 additions and 17 deletions.
155 changes: 155 additions & 0 deletions api/migrations/versions/b889f0fceba0_join_primary_key_to_unique_constraint.py
View file
Original file line number Diff line number Diff line change
@@ -0,0 +1,155 @@
"""join_primary_key_to_unique_constraint
Revision ID: b889f0fceba0
Revises: eeb2e349e6ac
Create Date: 2024-08-09 02:07:45.256491
"""
from alembic import op
import models as models
import sqlalchemy as sa


# revision identifiers, used by Alembic.
revision = 'b889f0fceba0'
down_revision = 'eeb2e349e6ac'
branch_labels = None
depends_on = None


def upgrade():
# ### commands auto generated by Alembic - please adjust! ###
with op.batch_alter_table('account_integrates', schema=None) as batch_op:
batch_op.drop_constraint('unique_account_provider', type_='unique')
batch_op.create_unique_constraint('unique_account_provider', ['account_id', 'provider', 'id'])
batch_op.drop_constraint('unique_provider_open_id', type_='unique')
batch_op.create_unique_constraint('unique_provider_open_id', ['provider', 'open_id', 'id'])

with op.batch_alter_table('celery_taskmeta', schema=None) as batch_op:
batch_op.drop_constraint('celery_taskmeta_task_id_key', type_='unique')
batch_op.create_unique_constraint('unique_task_id', ['task_id', 'id'])

with op.batch_alter_table('celery_tasksetmeta', schema=None) as batch_op:
batch_op.drop_constraint('celery_tasksetmeta_taskset_id_key', type_='unique')
batch_op.create_unique_constraint('unique_taskset_id', ['taskset_id', 'id'])

with op.batch_alter_table('dataset_keyword_tables', schema=None) as batch_op:
batch_op.drop_constraint('dataset_keyword_tables_dataset_id_key', type_='unique')
batch_op.create_unique_constraint('unique_dataset_id', ['dataset_id', 'id'])

with op.batch_alter_table('embeddings', schema=None) as batch_op:
batch_op.drop_constraint('embedding_hash_idx', type_='unique')
batch_op.create_unique_constraint('embedding_hash_idx', ['model_name', 'hash', 'provider_name', 'id'])

with op.batch_alter_table('installed_apps', schema=None) as batch_op:
batch_op.drop_constraint('unique_tenant_app', type_='unique')
batch_op.create_unique_constraint('unique_tenant_app', ['tenant_id', 'app_id', 'id'])

with op.batch_alter_table('provider_models', schema=None) as batch_op:
batch_op.drop_constraint('unique_provider_model_name', type_='unique')
batch_op.create_unique_constraint('unique_provider_model_name', ['tenant_id', 'provider_name', 'model_name', 'model_type', 'id'])

with op.batch_alter_table('providers', schema=None) as batch_op:
batch_op.drop_constraint('unique_provider_name_type_quota', type_='unique')
batch_op.create_unique_constraint('unique_provider_name_type_quota', ['tenant_id', 'provider_name', 'provider_type', 'quota_type', 'id'])

with op.batch_alter_table('tenant_account_joins', schema=None) as batch_op:
batch_op.drop_constraint('unique_tenant_account_join', type_='unique')
batch_op.create_unique_constraint('unique_tenant_account_join', ['tenant_id', 'account_id', 'id'])

with op.batch_alter_table('tool_api_providers', schema=None) as batch_op:
batch_op.drop_constraint('unique_api_tool_provider', type_='unique')
batch_op.create_unique_constraint('unique_api_tool_provider', ['name', 'tenant_id', 'id'])

with op.batch_alter_table('tool_builtin_providers', schema=None) as batch_op:
batch_op.drop_constraint('unique_builtin_tool_provider', type_='unique')
batch_op.create_unique_constraint('unique_builtin_tool_provider', ['tenant_id', 'provider', 'id'])

with op.batch_alter_table('tool_label_bindings', schema=None) as batch_op:
batch_op.drop_constraint('unique_tool_label_bind', type_='unique')
batch_op.create_unique_constraint('unique_tool_label_bind', ['tool_id', 'label_name', 'id'])

with op.batch_alter_table('tool_providers', schema=None) as batch_op:
batch_op.drop_constraint('unique_tool_provider_tool_name', type_='unique')
batch_op.create_unique_constraint('unique_tool_provider_tool_name', ['tenant_id', 'tool_name', 'id'])

with op.batch_alter_table('tool_published_apps', schema=None) as batch_op:
batch_op.drop_constraint('unique_published_app_tool', type_='unique')
batch_op.create_unique_constraint('unique_published_app_tool', ['app_id', 'user_id', 'id'])

with op.batch_alter_table('tool_workflow_providers', schema=None) as batch_op:
batch_op.drop_constraint('unique_workflow_tool_provider', type_='unique')
batch_op.create_unique_constraint('unique_workflow_tool_provider', ['name', 'tenant_id', 'id'])
batch_op.drop_constraint('unique_workflow_tool_provider_app_id', type_='unique')
batch_op.create_unique_constraint('unique_workflow_tool_provider_app_id', ['tenant_id', 'app_id', 'id'])

# ### end Alembic commands ###


def downgrade():
# ### commands auto generated by Alembic - please adjust! ###
with op.batch_alter_table('tool_workflow_providers', schema=None) as batch_op:
batch_op.drop_constraint('unique_workflow_tool_provider_app_id', type_='unique')
batch_op.create_unique_constraint('unique_workflow_tool_provider_app_id', ['tenant_id', 'app_id'])
batch_op.drop_constraint('unique_workflow_tool_provider', type_='unique')
batch_op.create_unique_constraint('unique_workflow_tool_provider', ['name', 'tenant_id'])

with op.batch_alter_table('tool_published_apps', schema=None) as batch_op:
batch_op.drop_constraint('unique_published_app_tool', type_='unique')
batch_op.create_unique_constraint('unique_published_app_tool', ['app_id', 'user_id'])

with op.batch_alter_table('tool_providers', schema=None) as batch_op:
batch_op.drop_constraint('unique_tool_provider_tool_name', type_='unique')
batch_op.create_unique_constraint('unique_tool_provider_tool_name', ['tenant_id', 'tool_name'])

with op.batch_alter_table('tool_label_bindings', schema=None) as batch_op:
batch_op.drop_constraint('unique_tool_label_bind', type_='unique')
batch_op.create_unique_constraint('unique_tool_label_bind', ['tool_id', 'label_name'])

with op.batch_alter_table('tool_builtin_providers', schema=None) as batch_op:
batch_op.drop_constraint('unique_builtin_tool_provider', type_='unique')
batch_op.create_unique_constraint('unique_builtin_tool_provider', ['tenant_id', 'provider'])

with op.batch_alter_table('tool_api_providers', schema=None) as batch_op:
batch_op.drop_constraint('unique_api_tool_provider', type_='unique')
batch_op.create_unique_constraint('unique_api_tool_provider', ['name', 'tenant_id'])

with op.batch_alter_table('tenant_account_joins', schema=None) as batch_op:
batch_op.drop_constraint('unique_tenant_account_join', type_='unique')
batch_op.create_unique_constraint('unique_tenant_account_join', ['tenant_id', 'account_id'])

with op.batch_alter_table('providers', schema=None) as batch_op:
batch_op.drop_constraint('unique_provider_name_type_quota', type_='unique')
batch_op.create_unique_constraint('unique_provider_name_type_quota', ['tenant_id', 'provider_name', 'provider_type', 'quota_type'])

with op.batch_alter_table('provider_models', schema=None) as batch_op:
batch_op.drop_constraint('unique_provider_model_name', type_='unique')
batch_op.create_unique_constraint('unique_provider_model_name', ['tenant_id', 'provider_name', 'model_name', 'model_type'])

with op.batch_alter_table('installed_apps', schema=None) as batch_op:
batch_op.drop_constraint('unique_tenant_app', type_='unique')
batch_op.create_unique_constraint('unique_tenant_app', ['tenant_id', 'app_id'])

with op.batch_alter_table('embeddings', schema=None) as batch_op:
batch_op.drop_constraint('embedding_hash_idx', type_='unique')
batch_op.create_unique_constraint('embedding_hash_idx', ['model_name', 'hash', 'provider_name'])

with op.batch_alter_table('dataset_keyword_tables', schema=None) as batch_op:
batch_op.drop_constraint('unique_dataset_id', type_='unique')
batch_op.create_unique_constraint('dataset_keyword_tables_dataset_id_key', ['dataset_id'])

with op.batch_alter_table('celery_tasksetmeta', schema=None) as batch_op:
batch_op.drop_constraint('unique_taskset_id', type_='unique')
batch_op.create_unique_constraint('celery_tasksetmeta_taskset_id_key', ['taskset_id'])

with op.batch_alter_table('celery_taskmeta', schema=None) as batch_op:
batch_op.drop_constraint('unique_task_id', type_='unique')
batch_op.create_unique_constraint('celery_taskmeta_task_id_key', ['task_id'])

with op.batch_alter_table('account_integrates', schema=None) as batch_op:
batch_op.drop_constraint('unique_provider_open_id', type_='unique')
batch_op.create_unique_constraint('unique_provider_open_id', ['provider', 'open_id'])
batch_op.drop_constraint('unique_account_provider', type_='unique')
batch_op.create_unique_constraint('unique_account_provider', ['account_id', 'provider'])

# ### end Alembic commands ###
6 changes: 3 additions & 3 deletions api/models/account.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -200,7 +200,7 @@ class TenantAccountJoin(db.Model):
db.PrimaryKeyConstraint('id', name='tenant_account_join_pkey'),
db.Index('tenant_account_join_account_id_idx', 'account_id'),
db.Index('tenant_account_join_tenant_id_idx', 'tenant_id'),
db.UniqueConstraint('tenant_id', 'account_id', name='unique_tenant_account_join')
db.UniqueConstraint('tenant_id', 'account_id', 'id', name='unique_tenant_account_join')
)

id = db.Column(StringUUID, server_default=db.text('uuid_generate_v4()'))
Expand All @@ -217,8 +217,8 @@ class AccountIntegrate(db.Model):
__tablename__ = 'account_integrates'
__table_args__ = (
db.PrimaryKeyConstraint('id', name='account_integrate_pkey'),
db.UniqueConstraint('account_id', 'provider', name='unique_account_provider'),
db.UniqueConstraint('provider', 'open_id', name='unique_provider_open_id')
db.UniqueConstraint('account_id', 'provider', 'id', name='unique_account_provider'),
db.UniqueConstraint('provider', 'open_id', 'id', name='unique_provider_open_id')
)

id = db.Column(StringUUID, server_default=db.text('uuid_generate_v4()'))
Expand Down
5 changes: 3 additions & 2 deletions api/models/dataset.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -584,11 +584,12 @@ class DatasetKeywordTable(db.Model):
__tablename__ = 'dataset_keyword_tables'
__table_args__ = (
db.PrimaryKeyConstraint('id', name='dataset_keyword_table_pkey'),
db.UniqueConstraint('dataset_id', 'id', name='unique_dataset_id'),
db.Index('dataset_keyword_table_dataset_id_idx', 'dataset_id'),
)

id = db.Column(StringUUID, primary_key=True, server_default=db.text('uuid_generate_v4()'))
dataset_id = db.Column(StringUUID, nullable=False, unique=True)
dataset_id = db.Column(StringUUID, nullable=False)
keyword_table = db.Column(db.Text, nullable=False)
data_source_type = db.Column(db.String(255), nullable=False,
server_default=db.text("'database'::character varying"))
Expand Down Expand Up @@ -630,7 +631,7 @@ class Embedding(db.Model):
__tablename__ = 'embeddings'
__table_args__ = (
db.PrimaryKeyConstraint('id', name='embedding_pkey'),
db.UniqueConstraint('model_name', 'hash', 'provider_name', name='embedding_hash_idx'),
db.UniqueConstraint('model_name', 'hash', 'provider_name', 'id', name='embedding_hash_idx'),
db.Index('created_at_idx', 'created_at')
)

Expand Down
2 changes: 1 addition & 1 deletion api/models/model.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -461,7 +461,7 @@ class InstalledApp(db.Model):
db.PrimaryKeyConstraint('id', name='installed_app_pkey'),
db.Index('installed_app_tenant_id_idx', 'tenant_id'),
db.Index('installed_app_app_id_idx', 'app_id'),
db.UniqueConstraint('tenant_id', 'app_id', name='unique_tenant_app')
db.UniqueConstraint('tenant_id', 'app_id', 'id', name='unique_tenant_app')
)

id = db.Column(StringUUID, server_default=db.text('uuid_generate_v4()'))
Expand Down
4 changes: 2 additions & 2 deletions api/models/provider.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -42,7 +42,7 @@ class Provider(db.Model):
__table_args__ = (
db.PrimaryKeyConstraint('id', name='provider_pkey'),
db.Index('provider_tenant_id_provider_idx', 'tenant_id', 'provider_name'),
db.UniqueConstraint('tenant_id', 'provider_name', 'provider_type', 'quota_type', name='unique_provider_name_type_quota')
db.UniqueConstraint('tenant_id', 'provider_name', 'provider_type', 'quota_type', 'id', name='unique_provider_name_type_quota')
)

id = db.Column(StringUUID, server_default=db.text('uuid_generate_v4()'))
Expand Down Expand Up @@ -89,7 +89,7 @@ class ProviderModel(db.Model):
__table_args__ = (
db.PrimaryKeyConstraint('id', name='provider_model_pkey'),
db.Index('provider_model_tenant_id_provider_idx', 'tenant_id', 'provider_name'),
db.UniqueConstraint('tenant_id', 'provider_name', 'model_name', 'model_type', name='unique_provider_model_name')
db.UniqueConstraint('tenant_id', 'provider_name', 'model_name', 'model_type', 'id', name='unique_provider_model_name')
)

id = db.Column(StringUUID, server_default=db.text('uuid_generate_v4()'))
Expand Down
10 changes: 8 additions & 2 deletions api/models/task.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -9,10 +9,13 @@ class CeleryTask(db.Model):
"""Task result/status."""

__tablename__ = 'celery_taskmeta'
__table_args__ = (
db.UniqueConstraint('task_id', 'id', name='unique_task_id'),
)

id = db.Column(db.Integer, db.Sequence('task_id_sequence'),
primary_key=True, autoincrement=True)
task_id = db.Column(db.String(155), unique=True)
task_id = db.Column(db.String(155))
status = db.Column(db.String(50), default=states.PENDING)
result = db.Column(db.PickleType, nullable=True)
date_done = db.Column(db.DateTime, default=lambda: datetime.now(timezone.utc).replace(tzinfo=None),
Expand All @@ -30,10 +33,13 @@ class CeleryTaskSet(db.Model):
"""TaskSet result."""

__tablename__ = 'celery_tasksetmeta'
__table_args__ = (
db.UniqueConstraint('taskset_id', 'id', name='unique_taskset_id'),
)

id = db.Column(db.Integer, db.Sequence('taskset_id_sequence'),
autoincrement=True, primary_key=True)
taskset_id = db.Column(db.String(155), unique=True)
taskset_id = db.Column(db.String(155))
result = db.Column(db.PickleType, nullable=True)
date_done = db.Column(db.DateTime, default=lambda: datetime.now(timezone.utc).replace(tzinfo=None),
nullable=True)
2 changes: 1 addition & 1 deletion api/models/tool.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -20,7 +20,7 @@ class ToolProvider(db.Model):
__tablename__ = 'tool_providers'
__table_args__ = (
db.PrimaryKeyConstraint('id', name='tool_provider_pkey'),
db.UniqueConstraint('tenant_id', 'tool_name', name='unique_tool_provider_tool_name')
db.UniqueConstraint('tenant_id', 'tool_name','id', name='unique_tool_provider_tool_name')
)

id = db.Column(StringUUID, server_default=db.text('uuid_generate_v4()'))
Expand Down
12 changes: 6 additions & 6 deletions api/models/tools.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -18,7 +18,7 @@ class BuiltinToolProvider(db.Model):
__table_args__ = (
db.PrimaryKeyConstraint('id', name='tool_builtin_provider_pkey'),
# one tenant can only have one tool provider with the same name
db.UniqueConstraint('tenant_id', 'provider', name='unique_builtin_tool_provider')
db.UniqueConstraint('tenant_id', 'provider', 'id', name='unique_builtin_tool_provider')
)

# id of the tool provider
Expand All @@ -45,7 +45,7 @@ class PublishedAppTool(db.Model):
__tablename__ = 'tool_published_apps'
__table_args__ = (
db.PrimaryKeyConstraint('id', name='published_app_tool_pkey'),
db.UniqueConstraint('app_id', 'user_id', name='unique_published_app_tool')
db.UniqueConstraint('app_id', 'user_id', 'id', name='unique_published_app_tool')
)

# id of the tool provider
Expand Down Expand Up @@ -84,7 +84,7 @@ class ApiToolProvider(db.Model):
__tablename__ = 'tool_api_providers'
__table_args__ = (
db.PrimaryKeyConstraint('id', name='tool_api_provider_pkey'),
db.UniqueConstraint('name', 'tenant_id', name='unique_api_tool_provider')
db.UniqueConstraint('name', 'tenant_id', 'id', name='unique_api_tool_provider')
)

id = db.Column(StringUUID, server_default=db.text('uuid_generate_v4()'))
Expand Down Expand Up @@ -140,7 +140,7 @@ class ToolLabelBinding(db.Model):
__tablename__ = 'tool_label_bindings'
__table_args__ = (
db.PrimaryKeyConstraint('id', name='tool_label_bind_pkey'),
db.UniqueConstraint('tool_id', 'label_name', name='unique_tool_label_bind'),
db.UniqueConstraint('tool_id', 'label_name', 'id', name='unique_tool_label_bind'),
)

id = db.Column(StringUUID, server_default=db.text('uuid_generate_v4()'))
Expand All @@ -158,8 +158,8 @@ class WorkflowToolProvider(db.Model):
__tablename__ = 'tool_workflow_providers'
__table_args__ = (
db.PrimaryKeyConstraint('id', name='tool_workflow_provider_pkey'),
db.UniqueConstraint('name', 'tenant_id', name='unique_workflow_tool_provider'),
db.UniqueConstraint('tenant_id', 'app_id', name='unique_workflow_tool_provider_app_id'),
db.UniqueConstraint('name', 'tenant_id', 'id', name='unique_workflow_tool_provider'),
db.UniqueConstraint('tenant_id', 'app_id', 'id', name='unique_workflow_tool_provider_app_id'),
)

id = db.Column(StringUUID, server_default=db.text('uuid_generate_v4()'))
Expand Down

0 comments on commit 5271612

Please sign in to comment.