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

refactor: use batch insert for SQLite as well

Browse source
This commit is contained in:
Igor Ilic 2025-11-06 16:23:54 +01:00
parent ac751bacf0
commit ef3a382669
1 changed files with 8 additions and 18 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

26
alembic/versions/c946955da633_multi_tenant_support.py
View file

@ -97,23 +97,13 @@ def upgrade() -> None:
# Insert into user_tenants table # Insert into user_tenants table
if user_data: if user_data:
if op.get_context().dialect.name == "sqlite": op.bulk_insert(
insert_stmt = user_tenants.insert().values( user_tenants,
[ [
{"user_id": user_id, "tenant_id": tenant_id, "created_at": _now()} {"user_id": user_id, "tenant_id": tenant_id, "created_at": _now()}
for user_id, tenant_id in user_data for user_id, tenant_id in user_data
] ],
) )
conn.execute(insert_stmt)
conn.commit()
else:
op.bulk_insert(
user_tenants,
[
{"user_id": user_id, "tenant_id": tenant_id, "created_at": _now()}
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")
if not tenant_id_column: if not tenant_id_column:
@ -141,7 +131,7 @@ 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_index(op.f("ix_datasets_tenant_id"), table_name="datasets") op.drop_index(op.f("ix_datasets_tenant_id"), table_name="datasets")
op.drop_column("datasets", "tenant_id") op.drop_column("datasets", "tenant_id")
op.drop_table("user_tenants")
# ### end Alembic commands ### # ### end Alembic commands ###