99 lines
3.2 KiB
Python
99 lines
3.2 KiB
Python
"""update identities unique constraint and properties
|
|
|
|
Revision ID: 549eff097c71
|
|
Revises: a3047a624130
|
|
Create Date: 2025-02-20 09:53:42.743105
|
|
|
|
"""
|
|
|
|
from typing import Sequence, Union
|
|
|
|
import sqlalchemy as sa
|
|
from sqlalchemy.dialects import postgresql
|
|
|
|
from alembic import op
|
|
from letta.settings import settings
|
|
|
|
# revision identifiers, used by Alembic.
|
|
revision: str = "549eff097c71"
|
|
down_revision: Union[str, None] = "a3047a624130"
|
|
branch_labels: Union[str, Sequence[str], None] = None
|
|
depends_on: Union[str, Sequence[str], None] = None
|
|
|
|
|
|
def upgrade() -> None:
|
|
# Skip this migration for SQLite
|
|
if not settings.letta_pg_uri_no_default:
|
|
return
|
|
|
|
# ### commands auto generated by Alembic - please adjust! ###
|
|
# Update unique constraint on identities table
|
|
op.drop_constraint("unique_identifier_pid_org_id", "identities", type_="unique")
|
|
op.create_unique_constraint(
|
|
"unique_identifier_without_project",
|
|
"identities",
|
|
["identifier_key", "project_id", "organization_id"],
|
|
postgresql_nulls_not_distinct=True,
|
|
)
|
|
|
|
# Add properties column to identities table
|
|
op.add_column("identities", sa.Column("properties", postgresql.JSONB, nullable=False, server_default="[]"))
|
|
|
|
# Create identities_agents table for many-to-many relationship
|
|
op.create_table(
|
|
"identities_agents",
|
|
sa.Column("identity_id", sa.String(), nullable=False),
|
|
sa.Column("agent_id", sa.String(), nullable=False),
|
|
sa.ForeignKeyConstraint(["agent_id"], ["agents.id"], ondelete="CASCADE"),
|
|
sa.ForeignKeyConstraint(["identity_id"], ["identities.id"], ondelete="CASCADE"),
|
|
sa.PrimaryKeyConstraint("identity_id", "agent_id"),
|
|
)
|
|
|
|
# Migrate existing relationships
|
|
# First, get existing relationships where identity_id is not null
|
|
op.execute(
|
|
"""
|
|
INSERT INTO identities_agents (identity_id, agent_id)
|
|
SELECT DISTINCT identity_id, id as agent_id
|
|
FROM agents
|
|
WHERE identity_id IS NOT NULL
|
|
"""
|
|
)
|
|
|
|
# Remove old identity_id column from agents
|
|
op.drop_column("agents", "identity_id")
|
|
op.drop_column("agents", "identifier_key")
|
|
# ### end Alembic commands ###
|
|
|
|
|
|
def downgrade() -> None:
|
|
# Skip this migration for SQLite
|
|
if not settings.letta_pg_uri_no_default:
|
|
return
|
|
|
|
# ### commands auto generated by Alembic - please adjust! ###
|
|
# Add back the old columns to agents
|
|
op.add_column("agents", sa.Column("identity_id", sa.String(), nullable=True))
|
|
op.add_column("agents", sa.Column("identifier_key", sa.String(), nullable=True))
|
|
|
|
# Migrate relationships back
|
|
op.execute(
|
|
"""
|
|
UPDATE agents a
|
|
SET identity_id = ia.identity_id
|
|
FROM identities_agents ia
|
|
WHERE a.id = ia.agent_id
|
|
"""
|
|
)
|
|
|
|
# Drop the many-to-many table
|
|
op.drop_table("identities_agents")
|
|
|
|
# Drop properties column
|
|
op.drop_column("identities", "properties")
|
|
|
|
# Restore old unique constraint
|
|
op.drop_constraint("unique_identifier_without_project", "identities", type_="unique")
|
|
op.create_unique_constraint("unique_identifier_pid_org_id", "identities", ["identifier_key", "project_id", "organization_id"])
|
|
# ### end Alembic commands ###
|