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

fix: Resolve multi user migration

Browse source
This commit is contained in:
Igor Ilic 2025-11-11 13:53:19 +01:00
parent ed2d687135
commit 011a7fb60b
1 changed files with 15 additions and 39 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

54
alembic/versions/76625596c5c3_expand_dataset_database_for_multi_user.py
View file

@ -14,7 +14,7 @@ import sqlalchemy as sa
# revision identifiers, used by Alembic. # revision identifiers, used by Alembic.
revision: str = "76625596c5c3" revision: str = "76625596c5c3"
down_revision: Union[str, None] = "211ab850ef3d" down_revision: Union[str, None] = "c946955da633"
branch_labels: Union[str, Sequence[str], None] = None branch_labels: Union[str, Sequence[str], None] = None
depends_on: Union[str, Sequence[str], None] = None depends_on: Union[str, Sequence[str], None] = None
@ -30,40 +30,20 @@ def upgrade() -> None:
conn = op.get_bind() conn = op.get_bind()
insp = sa.inspect(conn) insp = sa.inspect(conn)
data = sa.table(
"dataset_database",
sa.Column("dataset_id", sa.UUID, primary_key=True, index=True), # Critical for SQLite
sa.Column("owner_id", sa.UUID, index=True),
sa.Column("vector_database_name", sa.String(), unique=True, nullable=False),
sa.Column("graph_database_name", sa.String(), unique=True, nullable=False),
sa.Column("vector_database_provider", sa.String(), unique=False, nullable=False),
sa.Column("graph_database_provider", sa.String(), unique=False, nullable=False),
sa.Column("vector_database_url", sa.String(), unique=False, nullable=True),
sa.Column("graph_database_url", sa.String(), unique=False, nullable=True),
sa.Column("vector_database_key", sa.String(), unique=False, nullable=True),
sa.Column("graph_database_key", sa.String(), unique=False, nullable=True),
sa.Column("created_at", sa.DateTime(timezone=True)),
sa.Column("updated_at", sa.DateTime(timezone=True)),
)
vector_database_provider_column = _get_column( vector_database_provider_column = _get_column(
insp, "dataset_database", "vector_database_provider" insp, "dataset_database", "vector_database_provider"
) )
if not vector_database_provider_column: if not vector_database_provider_column:
op.add_column( op.add_column(
"dataset_database", "dataset_database",
sa.Column("vector_database_provider", sa.String(), unique=False, nullable=False), sa.Column(
"vector_database_provider",
sa.String(),
unique=False,
nullable=False,
server_default="lancedb",
),
) )
if op.get_context().dialect.name == "sqlite":
with op.batch_alter_table("dataset_database") as batch_op:
batch_op.execute(
data.update().values(
vector_database_provider="lancedb",
)
)
else:
conn = op.get_bind()
conn.execute(data.update().values(vector_database_provider="lancedb"))
graph_database_provider_column = _get_column( graph_database_provider_column = _get_column(
insp, "dataset_database", "graph_database_provider" insp, "dataset_database", "graph_database_provider"
@ -71,18 +51,14 @@ def upgrade() -> None:
if not graph_database_provider_column: if not graph_database_provider_column:
op.add_column( op.add_column(
"dataset_database", "dataset_database",
sa.Column("graph_database_provider", sa.String(), unique=False, nullable=False), sa.Column(
"graph_database_provider",
sa.String(),
unique=False,
nullable=False,
server_default="kuzu",
),
) )
if op.get_context().dialect.name == "sqlite":
with op.batch_alter_table("dataset_database") as batch_op:
batch_op.execute(
data.update().values(
graph_database_provider="kuzu",
)
)
else:
conn = op.get_bind()
conn.execute(data.update().values(graph_database_provider="kuzu"))
vector_database_url_column = _get_column(insp, "dataset_database", "vector_database_url") vector_database_url_column = _get_column(insp, "dataset_database", "vector_database_url")
if not vector_database_url_column: if not vector_database_url_column: