chore(api): Use uuidv7 as PK for new provider crendential tables (#24545)

This commit is contained in:
QuantumGhost
2025-08-27 11:22:08 +08:00
committed by GitHub
parent 726c429772
commit 58189ed9a0
4 changed files with 81 additions and 25 deletions

View File

@@ -6,10 +6,10 @@ Create Date: 2025-08-09 15:53:54.341341
"""
from alembic import op
from libs.uuid_utils import uuidv7
import models as models
import sqlalchemy as sa
from sqlalchemy.sql import table, column
import uuid
# revision identifiers, used by Alembic.
revision = 'e8446f481c1e'
@@ -21,7 +21,7 @@ depends_on = None
def upgrade():
# Create provider_credentials table
op.create_table('provider_credentials',
sa.Column('id', models.types.StringUUID(), server_default=sa.text('uuid_generate_v4()'), nullable=False),
sa.Column('id', models.types.StringUUID(), server_default=sa.text('uuidv7()'), nullable=False),
sa.Column('tenant_id', models.types.StringUUID(), nullable=False),
sa.Column('provider_name', sa.String(length=255), nullable=False),
sa.Column('credential_name', sa.String(length=255), nullable=False),
@@ -63,7 +63,7 @@ def migrate_existing_providers_data():
column('updated_at', sa.DateTime()),
column('credential_id', models.types.StringUUID()),
)
provider_credential_table = table('provider_credentials',
column('id', models.types.StringUUID()),
column('tenant_id', models.types.StringUUID()),
@@ -79,15 +79,15 @@ def migrate_existing_providers_data():
# Query all existing providers data
existing_providers = conn.execute(
sa.select(providers_table.c.id, providers_table.c.tenant_id,
sa.select(providers_table.c.id, providers_table.c.tenant_id,
providers_table.c.provider_name, providers_table.c.encrypted_config,
providers_table.c.created_at, providers_table.c.updated_at)
.where(providers_table.c.encrypted_config.isnot(None))
).fetchall()
# Iterate through each provider and insert into provider_credentials
for provider in existing_providers:
credential_id = str(uuid.uuid4())
credential_id = str(uuidv7())
if not provider.encrypted_config or provider.encrypted_config.strip() == '':
continue
@@ -134,7 +134,7 @@ def downgrade():
def migrate_data_back_to_providers():
"""Migrate data back from provider_credentials to providers table for downgrade"""
# Define table structure for data manipulation
providers_table = table('providers',
column('id', models.types.StringUUID()),
@@ -143,7 +143,7 @@ def migrate_data_back_to_providers():
column('encrypted_config', sa.Text()),
column('credential_id', models.types.StringUUID()),
)
provider_credential_table = table('provider_credentials',
column('id', models.types.StringUUID()),
column('tenant_id', models.types.StringUUID()),
@@ -160,18 +160,18 @@ def migrate_data_back_to_providers():
sa.select(providers_table.c.id, providers_table.c.credential_id)
.where(providers_table.c.credential_id.isnot(None))
).fetchall()
# For each provider, get the credential data and update providers table
for provider in providers_with_credentials:
credential = conn.execute(
sa.select(provider_credential_table.c.encrypted_config)
.where(provider_credential_table.c.id == provider.credential_id)
).fetchone()
if credential:
# Update providers table with encrypted_config from credential
conn.execute(
providers_table.update()
.where(providers_table.c.id == provider.id)
.values(encrypted_config=credential.encrypted_config)
)
)

View File

@@ -5,9 +5,9 @@ Revises: e8446f481c1e
Create Date: 2025-08-13 16:05:42.657730
"""
import uuid
from alembic import op
from libs.uuid_utils import uuidv7
import models as models
import sqlalchemy as sa
from sqlalchemy.sql import table, column
@@ -23,7 +23,7 @@ depends_on = None
def upgrade():
# Create provider_model_credentials table
op.create_table('provider_model_credentials',
sa.Column('id', models.types.StringUUID(), server_default=sa.text('uuid_generate_v4()'), nullable=False),
sa.Column('id', models.types.StringUUID(), server_default=sa.text('uuidv7()'), nullable=False),
sa.Column('tenant_id', models.types.StringUUID(), nullable=False),
sa.Column('provider_name', sa.String(length=255), nullable=False),
sa.Column('model_name', sa.String(length=255), nullable=False),
@@ -71,7 +71,7 @@ def migrate_existing_provider_models_data():
column('updated_at', sa.DateTime()),
column('credential_id', models.types.StringUUID()),
)
provider_model_credentials_table = table('provider_model_credentials',
column('id', models.types.StringUUID()),
column('tenant_id', models.types.StringUUID()),
@@ -90,19 +90,19 @@ def migrate_existing_provider_models_data():
# Query all existing provider_models data with encrypted_config
existing_provider_models = conn.execute(
sa.select(provider_models_table.c.id, provider_models_table.c.tenant_id,
sa.select(provider_models_table.c.id, provider_models_table.c.tenant_id,
provider_models_table.c.provider_name, provider_models_table.c.model_name,
provider_models_table.c.model_type, provider_models_table.c.encrypted_config,
provider_models_table.c.created_at, provider_models_table.c.updated_at)
.where(provider_models_table.c.encrypted_config.isnot(None))
).fetchall()
# Iterate through each provider_model and insert into provider_model_credentials
for provider_model in existing_provider_models:
if not provider_model.encrypted_config or provider_model.encrypted_config.strip() == '':
continue
credential_id = str(uuid.uuid4())
credential_id = str(uuidv7())
# Insert into provider_model_credentials table
conn.execute(
@@ -148,14 +148,14 @@ def downgrade():
def migrate_data_back_to_provider_models():
"""Migrate data back from provider_model_credentials to provider_models table for downgrade"""
# Define table structure for data manipulation
provider_models_table = table('provider_models',
column('id', models.types.StringUUID()),
column('encrypted_config', sa.Text()),
column('credential_id', models.types.StringUUID()),
)
provider_model_credentials_table = table('provider_model_credentials',
column('id', models.types.StringUUID()),
column('encrypted_config', sa.Text()),
@@ -169,14 +169,14 @@ def migrate_data_back_to_provider_models():
sa.select(provider_models_table.c.id, provider_models_table.c.credential_id)
.where(provider_models_table.c.credential_id.isnot(None))
).fetchall()
# For each provider_model, get the credential data and update provider_models table
for provider_model in provider_models_with_credentials:
credential = conn.execute(
sa.select(provider_model_credentials_table.c.encrypted_config)
.where(provider_model_credentials_table.c.id == provider_model.credential_id)
).fetchone()
if credential:
# Update provider_models table with encrypted_config from credential
conn.execute(

View File

@@ -274,7 +274,7 @@ class ProviderCredential(Base):
sa.Index("provider_credential_tenant_provider_idx", "tenant_id", "provider_name"),
)
id: Mapped[str] = mapped_column(StringUUID, server_default=text("uuid_generate_v4()"))
id: Mapped[str] = mapped_column(StringUUID, server_default=text("uuidv7()"))
tenant_id: Mapped[str] = mapped_column(StringUUID, nullable=False)
provider_name: Mapped[str] = mapped_column(String(255), nullable=False)
credential_name: Mapped[str] = mapped_column(String(255), nullable=False)
@@ -300,7 +300,7 @@ class ProviderModelCredential(Base):
),
)
id: Mapped[str] = mapped_column(StringUUID, server_default=text("uuid_generate_v4()"))
id: Mapped[str] = mapped_column(StringUUID, server_default=text("uuidv7()"))
tenant_id: Mapped[str] = mapped_column(StringUUID, nullable=False)
provider_name: Mapped[str] = mapped_column(String(255), nullable=False)
model_name: Mapped[str] = mapped_column(String(255), nullable=False)

View File

@@ -10,11 +10,13 @@ more reliable and realistic test scenarios.
import logging
import os
from collections.abc import Generator
from pathlib import Path
from typing import Optional
import pytest
from flask import Flask
from flask.testing import FlaskClient
from sqlalchemy import Engine, text
from sqlalchemy.orm import Session
from testcontainers.core.container import DockerContainer
from testcontainers.core.waiting_utils import wait_for_logs
@@ -64,7 +66,7 @@ class DifyTestContainers:
# PostgreSQL is used for storing user data, workflows, and application state
logger.info("Initializing PostgreSQL container...")
self.postgres = PostgresContainer(
image="postgres:16-alpine",
image="postgres:14-alpine",
)
self.postgres.start()
db_host = self.postgres.get_container_host_ip()
@@ -116,7 +118,7 @@ class DifyTestContainers:
# Start Redis container for caching and session management
# Redis is used for storing session data, cache entries, and temporary data
logger.info("Initializing Redis container...")
self.redis = RedisContainer(image="redis:latest", port=6379)
self.redis = RedisContainer(image="redis:6-alpine", port=6379)
self.redis.start()
redis_host = self.redis.get_container_host_ip()
redis_port = self.redis.get_exposed_port(6379)
@@ -184,6 +186,57 @@ class DifyTestContainers:
_container_manager = DifyTestContainers()
def _get_migration_dir() -> Path:
conftest_dir = Path(__file__).parent
return conftest_dir.parent.parent / "migrations"
def _get_engine_url(engine: Engine):
try:
return engine.url.render_as_string(hide_password=False).replace("%", "%%")
except AttributeError:
return str(engine.url).replace("%", "%%")
_UUIDv7SQL = r"""
/* Main function to generate a uuidv7 value with millisecond precision */
CREATE FUNCTION uuidv7() RETURNS uuid
AS
$$
-- Replace the first 48 bits of a uuidv4 with the current
-- number of milliseconds since 1970-01-01 UTC
-- and set the "ver" field to 7 by setting additional bits
SELECT encode(
set_bit(
set_bit(
overlay(uuid_send(gen_random_uuid()) placing
substring(int8send((extract(epoch from clock_timestamp()) * 1000)::bigint) from
3)
from 1 for 6),
52, 1),
53, 1), 'hex')::uuid;
$$ LANGUAGE SQL VOLATILE PARALLEL SAFE;
COMMENT ON FUNCTION uuidv7 IS
'Generate a uuid-v7 value with a 48-bit timestamp (millisecond precision) and 74 bits of randomness';
CREATE FUNCTION uuidv7_boundary(timestamptz) RETURNS uuid
AS
$$
/* uuid fields: version=0b0111, variant=0b10 */
SELECT encode(
overlay('\x00000000000070008000000000000000'::bytea
placing substring(int8send(floor(extract(epoch from $1) * 1000)::bigint) from 3)
from 1 for 6),
'hex')::uuid;
$$ LANGUAGE SQL STABLE STRICT PARALLEL SAFE;
COMMENT ON FUNCTION uuidv7_boundary(timestamptz) IS
'Generate a non-random uuidv7 with the given timestamp (first 48 bits) and all random bits to 0.
As the smallest possible uuidv7 for that timestamp, it may be used as a boundary for partitions.';
"""
def _create_app_with_containers() -> Flask:
"""
Create Flask application configured to use test containers.
@@ -211,7 +264,10 @@ def _create_app_with_containers() -> Flask:
# Initialize database schema
logger.info("Creating database schema...")
with app.app_context():
with db.engine.connect() as conn, conn.begin():
conn.execute(text(_UUIDv7SQL))
db.create_all()
logger.info("Database schema created successfully")