diff --git a/api/migrations/versions/2025_08_09_1553-e8446f481c1e_add_provider_credential_pool_support.py b/api/migrations/versions/2025_08_09_1553-e8446f481c1e_add_provider_credential_pool_support.py index 87b42346d..7bcdc8f49 100644 --- a/api/migrations/versions/2025_08_09_1553-e8446f481c1e_add_provider_credential_pool_support.py +++ b/api/migrations/versions/2025_08_09_1553-e8446f481c1e_add_provider_credential_pool_support.py @@ -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) - ) \ No newline at end of file + ) diff --git a/api/migrations/versions/2025_08_13_1605-0e154742a5fa_add_provider_model_multi_credential.py b/api/migrations/versions/2025_08_13_1605-0e154742a5fa_add_provider_model_multi_credential.py index bec1a4540..aa7331ec6 100644 --- a/api/migrations/versions/2025_08_13_1605-0e154742a5fa_add_provider_model_multi_credential.py +++ b/api/migrations/versions/2025_08_13_1605-0e154742a5fa_add_provider_model_multi_credential.py @@ -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( diff --git a/api/models/provider.py b/api/models/provider.py index e75b26fd3..18bf0ac5a 100644 --- a/api/models/provider.py +++ b/api/models/provider.py @@ -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) diff --git a/api/tests/test_containers_integration_tests/conftest.py b/api/tests/test_containers_integration_tests/conftest.py index 0369a5cbd..076a2a826 100644 --- a/api/tests/test_containers_integration_tests/conftest.py +++ b/api/tests/test_containers_integration_tests/conftest.py @@ -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")