From 527161202af1c674022c1cfa995546683d6824b9 Mon Sep 17 00:00:00 2001 From: xiaozeyu Date: Fri, 9 Aug 2024 10:17:24 +0800 Subject: [PATCH] fix: join primary key to unique constraint --- ...0_join_primary_key_to_unique_constraint.py | 155 ++++++++++++++++++ api/models/account.py | 6 +- api/models/dataset.py | 5 +- api/models/model.py | 2 +- api/models/provider.py | 4 +- api/models/task.py | 10 +- api/models/tool.py | 2 +- api/models/tools.py | 12 +- 8 files changed, 179 insertions(+), 17 deletions(-) create mode 100644 api/migrations/versions/b889f0fceba0_join_primary_key_to_unique_constraint.py diff --git a/api/migrations/versions/b889f0fceba0_join_primary_key_to_unique_constraint.py b/api/migrations/versions/b889f0fceba0_join_primary_key_to_unique_constraint.py new file mode 100644 index 00000000000000..466ef5714de87f --- /dev/null +++ b/api/migrations/versions/b889f0fceba0_join_primary_key_to_unique_constraint.py @@ -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 ### diff --git a/api/models/account.py b/api/models/account.py index d36b2b9fda3278..05c5202863b7e7 100644 --- a/api/models/account.py +++ b/api/models/account.py @@ -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()')) @@ -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()')) diff --git a/api/models/dataset.py b/api/models/dataset.py index 40f9f4cf83ae96..b625d9cfc39908 100644 --- a/api/models/dataset.py +++ b/api/models/dataset.py @@ -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")) @@ -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') ) diff --git a/api/models/model.py b/api/models/model.py index a6f517ea6b181f..e24ea56ce6fd94 100644 --- a/api/models/model.py +++ b/api/models/model.py @@ -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()')) diff --git a/api/models/provider.py b/api/models/provider.py index 4c14c33f095cee..d23712e1da9dcf 100644 --- a/api/models/provider.py +++ b/api/models/provider.py @@ -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()')) @@ -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()')) diff --git a/api/models/task.py b/api/models/task.py index 618d831d8ed4e9..159596bfd56775 100644 --- a/api/models/task.py +++ b/api/models/task.py @@ -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), @@ -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) diff --git a/api/models/tool.py b/api/models/tool.py index f322944f5f0a8e..ad591189b7919e 100644 --- a/api/models/tool.py +++ b/api/models/tool.py @@ -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()')) diff --git a/api/models/tools.py b/api/models/tools.py index 49212916ec5195..2deb3adff5eaa4 100644 --- a/api/models/tools.py +++ b/api/models/tools.py @@ -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 @@ -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 @@ -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()')) @@ -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()')) @@ -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()'))