gmakstutis/cognee
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

chore: Add multi tenant migration

Browse source
This commit is contained in:
Igor Ilic 2025-11-05 17:24:11 +01:00
parent c4807a0c67
commit 9b6cbaf389
2 changed files with 26 additions and 18 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

6
alembic/versions/211ab850ef3d_add_sync_operations_table.py
View file

@ -36,7 +36,8 @@ def upgrade() -> None:
"FAILED", "FAILED",
"CANCELLED", "CANCELLED",
name="syncstatus", name="syncstatus",
create_type=False, create_type=True,
checkfirst=True,
) )
else: else:
syncstatus = sa.Enum( syncstatus = sa.Enum(
@ -46,7 +47,8 @@ def upgrade() -> None:
"FAILED", "FAILED",
"CANCELLED", "CANCELLED",
name="syncstatus", name="syncstatus",
create_type=False, create_type=True,
checkfirst=True,
) )
# Table doesn't exist, create it normally # Table doesn't exist, create it normally

38
alembic/versions/c946955da633_multi_tenant_support.py
View file

@ -20,6 +20,10 @@ branch_labels: Union[str, Sequence[str], None] = None
depends_on: Union[str, Sequence[str], None] = None depends_on: Union[str, Sequence[str], None] = None
def _now():
return datetime.now(timezone.utc)
def _define_user_table() -> sa.Table: def _define_user_table() -> sa.Table:
table = sa.Table( table = sa.Table(
"users", "users",
@ -76,27 +80,29 @@ def upgrade() -> None:
user = _define_user_table() user = _define_user_table()
if "user_tenants" not in insp.get_table_names(): if "user_tenants" not in insp.get_table_names():
tenant_id_from_user = sa.select(user.c.tenant_id).scalar_subquery()
# Define table with all necessary columns including primary key # Define table with all necessary columns including primary key
user_tenants = op.create_table( user_tenants = op.create_table(
"user_tenants", "user_tenants",
sa.Column("user_id", sa.UUID, sa.ForeignKey("users.id"), primary_key=True), sa.Column("user_id", sa.UUID, sa.ForeignKey("users.id"), primary_key=True),
sa.Column("tenant_id", sa.UUID, sa.ForeignKey("tenants.id"), primary_key=True), sa.Column("tenant_id", sa.UUID, sa.ForeignKey("tenants.id"), primary_key=True),
sa.Column("created_at", sa.DateTime(), default=lambda: datetime.now(timezone.utc)), sa.Column(
"created_at", sa.DateTime(timezone=True), default=lambda: datetime.now(timezone.utc)
),
) )
if op.get_context().dialect.name == "sqlite":
# If column doesn't exist create new original_extension column and update from values of extension column # Get all users with their tenant_id
with op.batch_alter_table("user_tenants") as batch_op: user_data = conn.execute(
batch_op.execute( sa.select(user.c.id, user.c.tenant_id).where(user.c.tenant_id.isnot(None))
user_tenants.update().values( ).fetchall()
tenant_id=tenant_id_from_user,
user_id=user.c.id, # Insert into user_tenants table
) if user_data:
) op.bulk_insert(
else: user_tenants,
conn = op.get_bind() [
conn.execute( {"user_id": user_id, "tenant_id": tenant_id, "created_at": _now()}
user_tenants.update().values(tenant_id=tenant_id_from_user, user_id=user.c.id) for user_id, tenant_id in user_data
],
) )
tenant_id_column = _get_column(insp, "datasets", "tenant_id") tenant_id_column = _get_column(insp, "datasets", "tenant_id")
@ -125,6 +131,6 @@ def upgrade() -> None:
def downgrade() -> None: def downgrade() -> None:
# ### commands auto generated by Alembic - please adjust! ### # ### commands auto generated by Alembic - please adjust! ###
op.drop_table("user_tenants")
op.drop_column("datasets", "tenant_id") op.drop_column("datasets", "tenant_id")
# ### end Alembic commands ### # ### end Alembic commands ###