From 011a7fb60bd016047da2618bbb61ffa85eb52028 Mon Sep 17 00:00:00 2001 From: Igor Ilic Date: Tue, 11 Nov 2025 13:53:19 +0100 Subject: [PATCH 1/4] fix: Resolve multi user migration --- ..._expand_dataset_database_for_multi_user.py | 54 ++++++------------- 1 file changed, 15 insertions(+), 39 deletions(-) diff --git a/alembic/versions/76625596c5c3_expand_dataset_database_for_multi_user.py b/alembic/versions/76625596c5c3_expand_dataset_database_for_multi_user.py index cd19d09c8..7e13898ae 100644 --- a/alembic/versions/76625596c5c3_expand_dataset_database_for_multi_user.py +++ b/alembic/versions/76625596c5c3_expand_dataset_database_for_multi_user.py @@ -14,7 +14,7 @@ import sqlalchemy as sa # revision identifiers, used by Alembic. revision: str = "76625596c5c3" -down_revision: Union[str, None] = "211ab850ef3d" +down_revision: Union[str, None] = "c946955da633" branch_labels: 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() 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( insp, "dataset_database", "vector_database_provider" ) if not vector_database_provider_column: op.add_column( "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( insp, "dataset_database", "graph_database_provider" @@ -71,18 +51,14 @@ def upgrade() -> None: if not graph_database_provider_column: op.add_column( "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") if not vector_database_url_column: From bb8de7b336a23b6053de3459764d71d7e45b40be Mon Sep 17 00:00:00 2001 From: Igor Ilic <30923996+dexters1@users.noreply.github.com> Date: Tue, 11 Nov 2025 13:56:16 +0100 Subject: [PATCH 2/4] Apply suggestion from @dexters1 --- .../databases/utils/get_or_create_dataset_database.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/cognee/infrastructure/databases/utils/get_or_create_dataset_database.py b/cognee/infrastructure/databases/utils/get_or_create_dataset_database.py index 61d7840c0..0aa836174 100644 --- a/cognee/infrastructure/databases/utils/get_or_create_dataset_database.py +++ b/cognee/infrastructure/databases/utils/get_or_create_dataset_database.py @@ -35,7 +35,7 @@ async def get_or_create_dataset_database( vector_config = get_vectordb_config() graph_config = get_graph_config() - +Note: for hybrid databases both graph and vector DB name have to be the same if graph_config.graph_database_provider == "kuzu": graph_db_name = f"{dataset_id}.pkl" else: From 20d49eeb76ca81057e902fad67658d5db2dcb0a0 Mon Sep 17 00:00:00 2001 From: Igor Ilic <30923996+dexters1@users.noreply.github.com> Date: Tue, 11 Nov 2025 13:56:35 +0100 Subject: [PATCH 3/4] Apply suggestion from @dexters1 --- .../databases/utils/get_or_create_dataset_database.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/cognee/infrastructure/databases/utils/get_or_create_dataset_database.py b/cognee/infrastructure/databases/utils/get_or_create_dataset_database.py index 0aa836174..635734d89 100644 --- a/cognee/infrastructure/databases/utils/get_or_create_dataset_database.py +++ b/cognee/infrastructure/databases/utils/get_or_create_dataset_database.py @@ -35,7 +35,7 @@ async def get_or_create_dataset_database( vector_config = get_vectordb_config() graph_config = get_graph_config() -Note: for hybrid databases both graph and vector DB name have to be the same +# Note: for hybrid databases both graph and vector DB name have to be the same if graph_config.graph_database_provider == "kuzu": graph_db_name = f"{dataset_id}.pkl" else: From 41b844a31c3b938a09a07462e59b70b39593313f Mon Sep 17 00:00:00 2001 From: Igor Ilic <30923996+dexters1@users.noreply.github.com> Date: Tue, 11 Nov 2025 13:56:59 +0100 Subject: [PATCH 4/4] Apply suggestion from @dexters1 --- .../databases/utils/get_or_create_dataset_database.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/cognee/infrastructure/databases/utils/get_or_create_dataset_database.py b/cognee/infrastructure/databases/utils/get_or_create_dataset_database.py index 635734d89..1822221cb 100644 --- a/cognee/infrastructure/databases/utils/get_or_create_dataset_database.py +++ b/cognee/infrastructure/databases/utils/get_or_create_dataset_database.py @@ -35,7 +35,7 @@ async def get_or_create_dataset_database( vector_config = get_vectordb_config() graph_config = get_graph_config() -# Note: for hybrid databases both graph and vector DB name have to be the same + # Note: for hybrid databases both graph and vector DB name have to be the same if graph_config.graph_database_provider == "kuzu": graph_db_name = f"{dataset_id}.pkl" else: