cognee 0.5.1__py3-none-any.whl → 0.5.2__py3-none-any.whl

cognee/alembic/versions/ab7e313804ae_permission_system_rework.py

@@ -0,0 +1,236 @@
+"""permission_system_rework
+
+Revision ID: ab7e313804ae
+Revises: 1d0bb7fede17
+Create Date: 2025-06-16 15:20:43.118246
+
+"""
+
+from typing import Sequence, Union
+from alembic import op
+import sqlalchemy as sa
+from sqlalchemy import UUID
+from datetime import datetime, timezone
+from uuid import uuid4
+
+# revision identifiers, used by Alembic.
+revision: str = "ab7e313804ae"
+down_revision: Union[str, None] = "1d0bb7fede17"
+branch_labels: Union[str, Sequence[str], None] = None
+depends_on: Union[str, Sequence[str], None] = None
+
+
+def _now():
+    return datetime.now(timezone.utc)
+
+
+def _define_dataset_table() -> sa.Table:
+    # Note: We can't use any Cognee model info to gather data (as it can change) in database so we must use our own table
+    #       definition or load what is in the database
+    table = sa.Table(
+        "datasets",
+        sa.MetaData(),
+        sa.Column("id", UUID, primary_key=True, default=uuid4),
+        sa.Column("name", sa.Text),
+        sa.Column(
+            "created_at",
+            sa.DateTime(timezone=True),
+            default=lambda: datetime.now(timezone.utc),
+        ),
+        sa.Column(
+            "updated_at",
+            sa.DateTime(timezone=True),
+            onupdate=lambda: datetime.now(timezone.utc),
+        ),
+        sa.Column("owner_id", UUID, sa.ForeignKey("principals.id"), index=True),
+    )
+
+    return table
+
+
+def _define_data_table() -> sa.Table:
+    # Note: We can't use any Cognee model info to gather data (as it can change) in database so we must use our own table
+    #       definition or load what is in the database
+    table = sa.Table(
+        "data",
+        sa.MetaData(),
+        sa.Column("id", UUID, primary_key=True, default=uuid4),
+        sa.Column("name", sa.String),
+        sa.Column("extension", sa.String),
+        sa.Column("mime_type", sa.String),
+        sa.Column("raw_data_location", sa.String),
+        sa.Column("owner_id", UUID, index=True),
+        sa.Column("content_hash", sa.String),
+        sa.Column("external_metadata", sa.JSON),
+        sa.Column("node_set", sa.JSON, nullable=True),  # list of strings
+        sa.Column("token_count", sa.Integer),
+        sa.Column(
+            "created_at",
+            sa.DateTime(timezone=True),
+            default=lambda: datetime.now(timezone.utc),
+        ),
+        sa.Column(
+            "updated_at",
+            sa.DateTime(timezone=True),
+            onupdate=lambda: datetime.now(timezone.utc),
+        ),
+    )
+
+    return table
+
+
+def _ensure_permission(conn, permission_name) -> str:
+    """
+    Return the permission.id for the given name, creating the row if needed.
+    """
+    permissions_table = sa.Table(
+        "permissions",
+        sa.MetaData(),
+        sa.Column("id", UUID, primary_key=True, index=True, default=uuid4),
+        sa.Column(
+            "created_at", sa.DateTime(timezone=True), default=lambda: datetime.now(timezone.utc)
+        ),
+        sa.Column(
+            "updated_at",
+            sa.DateTime(timezone=True),
+            onupdate=lambda: datetime.now(timezone.utc),
+        ),
+        sa.Column("name", sa.String, unique=True, nullable=False, index=True),
+    )
+    row = conn.execute(
+        sa.select(permissions_table).filter(permissions_table.c.name == permission_name)
+    ).fetchone()
+
+    if row is None:
+        permission_id = uuid4()
+
+        op.bulk_insert(
+            permissions_table,
+            [
+                {
+                    "id": permission_id,
+                    "name": permission_name,
+                    "created_at": _now(),
+                }
+            ],
+        )
+        return permission_id
+
+    return row.id
+
+
+def _build_acl_row(*, user_id, target_id, permission_id, target_col) -> dict:
+    """Create a dict with the correct column names for the ACL row."""
+    return {
+        "id": uuid4(),
+        "created_at": _now(),
+        "principal_id": user_id,
+        target_col: target_id,
+        "permission_id": permission_id,
+    }
+
+
+def _create_dataset_permission(conn, user_id, dataset_id, permission_name):
+    perm_id = _ensure_permission(conn, permission_name)
+    return _build_acl_row(
+        user_id=user_id, target_id=dataset_id, permission_id=perm_id, target_col="dataset_id"
+    )
+
+
+def _create_data_permission(conn, user_id, data_id, permission_name):
+    perm_id = _ensure_permission(conn, permission_name)
+    return _build_acl_row(
+        user_id=user_id, target_id=data_id, permission_id=perm_id, target_col="data_id"
+    )
+
+
+def _get_column(inspector, table, name, schema=None):
+    for col in inspector.get_columns(table, schema=schema):
+        if col["name"] == name:
+            return col
+    return None
+
+
+def upgrade() -> None:
+    conn = op.get_bind()
+    insp = sa.inspect(conn)
+
+    dataset_id_column = _get_column(insp, "acls", "dataset_id")
+    if not dataset_id_column:
+        # Recreate ACLs table with default permissions set to datasets instead of documents
+        op.drop_table("acls")
+
+        acls_table = op.create_table(
+            "acls",
+            sa.Column("id", UUID, primary_key=True, default=uuid4),
+            sa.Column(
+                "created_at", sa.DateTime(timezone=True), default=lambda: datetime.now(timezone.utc)
+            ),
+            sa.Column(
+                "updated_at",
+                sa.DateTime(timezone=True),
+                onupdate=lambda: datetime.now(timezone.utc),
+            ),
+            sa.Column("principal_id", UUID, sa.ForeignKey("principals.id")),
+            sa.Column("permission_id", UUID, sa.ForeignKey("permissions.id")),
+            sa.Column("dataset_id", UUID, sa.ForeignKey("datasets.id", ondelete="CASCADE")),
+        )
+
+        # Note: We can't use any Cognee model info to gather data (as it can change) in database so we must use our own table
+        #       definition or load what is in the database
+        dataset_table = _define_dataset_table()
+        datasets = conn.execute(sa.select(dataset_table)).fetchall()
+
+        if not datasets:
+            return
+
+        acl_list = []
+
+        for dataset in datasets:
+            acl_list.append(_create_dataset_permission(conn, dataset.owner_id, dataset.id, "read"))
+            acl_list.append(_create_dataset_permission(conn, dataset.owner_id, dataset.id, "write"))
+            acl_list.append(_create_dataset_permission(conn, dataset.owner_id, dataset.id, "share"))
+            acl_list.append(
+                _create_dataset_permission(conn, dataset.owner_id, dataset.id, "delete")
+            )
+
+        if acl_list:
+            op.bulk_insert(acls_table, acl_list)
+
+
+def downgrade() -> None:
+    conn = op.get_bind()
+
+    op.drop_table("acls")
+
+    acls_table = op.create_table(
+        "acls",
+        sa.Column("id", UUID, primary_key=True, nullable=False, default=uuid4),
+        sa.Column(
+            "created_at", sa.DateTime(timezone=True), default=lambda: datetime.now(timezone.utc)
+        ),
+        sa.Column(
+            "updated_at", sa.DateTime(timezone=True), onupdate=lambda: datetime.now(timezone.utc)
+        ),
+        sa.Column("principal_id", UUID, sa.ForeignKey("principals.id")),
+        sa.Column("permission_id", UUID, sa.ForeignKey("permissions.id")),
+        sa.Column("data_id", UUID, sa.ForeignKey("data.id", ondelete="CASCADE")),
+    )
+
+    # Note: We can't use any Cognee model info to gather data (as it can change) in database so we must use our own table
+    #       definition or load what is in the database
+    data_table = _define_data_table()
+    data = conn.execute(sa.select(data_table)).fetchall()
+
+    if not data:
+        return
+
+    acl_list = []
+    for single_data in data:
+        acl_list.append(_create_data_permission(conn, single_data.owner_id, single_data.id, "read"))
+        acl_list.append(
+            _create_data_permission(conn, single_data.owner_id, single_data.id, "write")
+        )
+
+    if acl_list:
+        op.bulk_insert(acls_table, acl_list)

cognee/alembic/versions/b9274c27a25a_kuzu_11_migration.py

@@ -0,0 +1,75 @@
+"""kuzu-11-migration
+
+Revision ID: b9274c27a25a
+Revises: e4ebee1091e7
+Create Date: 2025-07-24 17:11:52.174737
+
+"""
+
+import os
+from typing import Sequence, Union
+
+from cognee.infrastructure.databases.graph.kuzu.kuzu_migrate import (
+    kuzu_migration,
+    read_kuzu_storage_version,
+)
+import kuzu
+
+# revision identifiers, used by Alembic.
+revision: str = "b9274c27a25a"
+down_revision: Union[str, None] = "e4ebee1091e7"
+branch_labels: Union[str, Sequence[str], None] = None
+depends_on: Union[str, Sequence[str], None] = None
+
+
+def upgrade() -> None:
+    # This migration is only for multi-user Cognee mode
+    if os.getenv("ENABLE_BACKEND_ACCESS_CONTROL", "false").lower() == "true":
+        from cognee.base_config import get_base_config
+
+        base_config = get_base_config()
+
+        databases_root = os.path.join(base_config.system_root_directory, "databases")
+        if not os.path.isdir(databases_root):
+            raise FileNotFoundError(f"Directory not found: {databases_root}")
+
+        for current_path, dirnames, _ in os.walk(databases_root):
+            # If file is kuzu graph database
+            if ".pkl" in current_path[-4:]:
+                kuzu_db_version = read_kuzu_storage_version(current_path)
+                if (
+                    kuzu_db_version == "0.9.0" or kuzu_db_version == "0.8.2"
+                ) and kuzu_db_version != kuzu.__version__:
+                    # Try to migrate kuzu database to latest version
+                    kuzu_migration(
+                        new_db=current_path + "_new",
+                        old_db=current_path,
+                        new_version=kuzu.__version__,
+                        old_version=kuzu_db_version,
+                        overwrite=True,
+                    )
+    else:
+        from cognee.infrastructure.databases.graph import get_graph_config
+
+        graph_config = get_graph_config()
+        if graph_config.graph_database_provider.lower() == "kuzu":
+            if os.path.exists(graph_config.graph_file_path):
+                kuzu_db_version = read_kuzu_storage_version(graph_config.graph_file_path)
+                if (
+                    kuzu_db_version == "0.9.0" or kuzu_db_version == "0.8.2"
+                ) and kuzu_db_version != kuzu.__version__:
+                    # Try to migrate kuzu database to latest version
+                    kuzu_migration(
+                        new_db=graph_config.graph_file_path + "_new",
+                        old_db=graph_config.graph_file_path,
+                        new_version=kuzu.__version__,
+                        old_version=kuzu_db_version,
+                        overwrite=True,
+                    )
+
+
+def downgrade() -> None:
+    # To downgrade you will have to manually change the backup old kuzu graph databases
+    # stored in the user folder to its previous name and remove the new kuzu graph
+    # database that replaced it
+    pass

cognee/alembic/versions/c946955da633_multi_tenant_support.py

@@ -0,0 +1,137 @@
+"""Multi Tenant Support
+
+Revision ID: c946955da633
+Revises: 211ab850ef3d
+Create Date: 2025-11-04 18:11:09.325158
+
+"""
+
+from typing import Sequence, Union
+from datetime import datetime, timezone
+from uuid import uuid4
+
+from alembic import op
+import sqlalchemy as sa
+
+# revision identifiers, used by Alembic.
+revision: str = "c946955da633"
+down_revision: Union[str, None] = "211ab850ef3d"
+branch_labels: 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:
+    table = sa.Table(
+        "users",
+        sa.MetaData(),
+        sa.Column(
+            "id",
+            sa.UUID,
+            sa.ForeignKey("principals.id", ondelete="CASCADE"),
+            primary_key=True,
+            nullable=False,
+        ),
+        sa.Column("tenant_id", sa.UUID, sa.ForeignKey("tenants.id"), index=True, nullable=True),
+    )
+    return table
+
+
+def _define_dataset_table() -> sa.Table:
+    # Note: We can't use any Cognee model info to gather data (as it can change) in database so we must use our own table
+    #       definition or load what is in the database
+    table = sa.Table(
+        "datasets",
+        sa.MetaData(),
+        sa.Column("id", sa.UUID, primary_key=True, default=uuid4),
+        sa.Column("name", sa.Text),
+        sa.Column(
+            "created_at",
+            sa.DateTime(timezone=True),
+            default=lambda: datetime.now(timezone.utc),
+        ),
+        sa.Column(
+            "updated_at",
+            sa.DateTime(timezone=True),
+            onupdate=lambda: datetime.now(timezone.utc),
+        ),
+        sa.Column("owner_id", sa.UUID(), sa.ForeignKey("principals.id"), index=True),
+        sa.Column("tenant_id", sa.UUID(), sa.ForeignKey("tenants.id"), index=True, nullable=True),
+    )
+
+    return table
+
+
+def _get_column(inspector, table, name, schema=None):
+    for col in inspector.get_columns(table, schema=schema):
+        if col["name"] == name:
+            return col
+    return None
+
+
+def upgrade() -> None:
+    conn = op.get_bind()
+    insp = sa.inspect(conn)
+
+    dataset = _define_dataset_table()
+    user = _define_user_table()
+
+    if "user_tenants" not in insp.get_table_names():
+        # Define table with all necessary columns including primary key
+        user_tenants = op.create_table(
+            "user_tenants",
+            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(
+                "created_at", sa.DateTime(timezone=True), default=lambda: datetime.now(timezone.utc)
+            ),
+        )
+
+        # Get all users with their tenant_id
+        user_data = conn.execute(
+            sa.select(user.c.id, user.c.tenant_id).where(user.c.tenant_id.isnot(None))
+        ).fetchall()
+
+        # Insert into user_tenants table
+        if user_data:
+            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")
+    if not tenant_id_column:
+        op.add_column("datasets", sa.Column("tenant_id", sa.UUID(), nullable=True))
+
+        # Build subquery, select users.tenant_id for each dataset.owner_id
+        tenant_id_from_dataset_owner = (
+            sa.select(user.c.tenant_id).where(user.c.id == dataset.c.owner_id).scalar_subquery()
+        )
+
+        if op.get_context().dialect.name == "sqlite":
+            # If column doesn't exist create new original_extension column and update from values of extension column
+            with op.batch_alter_table("datasets") as batch_op:
+                batch_op.execute(
+                    dataset.update().values(
+                        tenant_id=tenant_id_from_dataset_owner,
+                    )
+                )
+        else:
+            conn = op.get_bind()
+            conn.execute(dataset.update().values(tenant_id=tenant_id_from_dataset_owner))
+
+        op.create_index(op.f("ix_datasets_tenant_id"), "datasets", ["tenant_id"])
+
+
+def downgrade() -> None:
+    # ### 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_column("datasets", "tenant_id")
+    # ### end Alembic commands ###

cognee/alembic/versions/e1ec1dcb50b6_add_last_accessed_to_data.py

@@ -0,0 +1,51 @@
+"""add_last_accessed_to_data
+
+Revision ID: e1ec1dcb50b6
+Revises: 211ab850ef3d
+Create Date: 2025-11-04 21:45:52.642322
+
+"""
+
+import os
+from typing import Sequence, Union
+
+from alembic import op
+import sqlalchemy as sa
+
+
+# revision identifiers, used by Alembic.
+revision: str = "e1ec1dcb50b6"
+down_revision: Union[str, None] = "a1b2c3d4e5f6"
+branch_labels: Union[str, Sequence[str], None] = None
+depends_on: Union[str, Sequence[str], None] = None
+
+
+def _get_column(inspector, table, name, schema=None):
+    for col in inspector.get_columns(table, schema=schema):
+        if col["name"] == name:
+            return col
+    return None
+
+
+def upgrade() -> None:
+    conn = op.get_bind()
+    insp = sa.inspect(conn)
+
+    last_accessed_column = _get_column(insp, "data", "last_accessed")
+    if not last_accessed_column:
+        # Always create the column for schema consistency
+        op.add_column("data", sa.Column("last_accessed", sa.DateTime(timezone=True), nullable=True))
+
+        # Only initialize existing records if feature is enabled
+        enable_last_accessed = os.getenv("ENABLE_LAST_ACCESSED", "false").lower() == "true"
+        if enable_last_accessed:
+            op.execute("UPDATE data SET last_accessed = CURRENT_TIMESTAMP")
+
+
+def downgrade() -> None:
+    conn = op.get_bind()
+    insp = sa.inspect(conn)
+
+    last_accessed_column = _get_column(insp, "data", "last_accessed")
+    if last_accessed_column:
+        op.drop_column("data", "last_accessed")

cognee/alembic/versions/e4ebee1091e7_expand_data_model_info.py

@@ -0,0 +1,140 @@
+"""Expand data model info
+
+Revision ID: e4ebee1091e7
+Revises: ab7e313804ae
+Create Date: 2025-07-24 13:21:30.738486
+
+"""
+
+from typing import Sequence, Union
+
+from alembic import op
+import sqlalchemy as sa
+from sqlalchemy.dialects import postgresql
+
+# revision identifiers, used by Alembic.
+revision: str = "e4ebee1091e7"
+down_revision: Union[str, None] = "ab7e313804ae"
+branch_labels: Union[str, Sequence[str], None] = None
+depends_on: Union[str, Sequence[str], None] = None
+
+
+def _get_column(inspector, table, name, schema=None):
+    for col in inspector.get_columns(table, schema=schema):
+        if col["name"] == name:
+            return col
+    return None
+
+
+def _index_exists(inspector, table, name, schema=None):
+    return any(ix["name"] == name for ix in inspector.get_indexes(table, schema=schema))
+
+
+def upgrade() -> None:
+    TABLES_TO_DROP = [
+        "file_metadata",
+        "_dlt_loads",
+        "_dlt_version",
+        "_dlt_pipeline_state",
+    ]
+
+    conn = op.get_bind()
+    insp = sa.inspect(conn)
+    existing = set(insp.get_table_names())
+
+    for tbl in TABLES_TO_DROP:
+        if tbl in existing:
+            op.drop_table(tbl)
+
+    DATA_TABLE = "data"
+    DATA_TENANT_COL = "tenant_id"
+    DATA_SIZE_COL = "data_size"
+    DATA_TENANT_IDX = "ix_data_tenant_id"
+
+    # --- tenant_id ---
+    col = _get_column(insp, DATA_TABLE, DATA_TENANT_COL)
+    if col is None:
+        op.add_column(
+            DATA_TABLE,
+            sa.Column(DATA_TENANT_COL, postgresql.UUID(as_uuid=True), nullable=True),
+        )
+    else:
+        # Column exists – fix nullability if needed
+        if col.get("nullable", True) is False:
+            op.alter_column(
+                DATA_TABLE,
+                DATA_TENANT_COL,
+                existing_type=postgresql.UUID(as_uuid=True),
+                nullable=True,
+            )
+
+    # --- data_size ---
+    col = _get_column(insp, DATA_TABLE, DATA_SIZE_COL)
+    if col is None:
+        op.add_column(DATA_TABLE, sa.Column(DATA_SIZE_COL, sa.Integer(), nullable=True))
+    else:
+        # If you also need to change nullability for data_size, do it here
+        if col.get("nullable", True) is False:
+            op.alter_column(
+                DATA_TABLE,
+                DATA_SIZE_COL,
+                existing_type=sa.Integer(),
+                nullable=True,
+            )
+
+    # --- index on tenant_id ---
+    if not _index_exists(insp, DATA_TABLE, DATA_TENANT_IDX):
+        op.create_index(DATA_TENANT_IDX, DATA_TABLE, [DATA_TENANT_COL], unique=False)
+
+
+def downgrade() -> None:
+    op.drop_index(op.f("ix_data_tenant_id"), table_name="data")
+    op.drop_column("data", "data_size")
+    op.drop_column("data", "tenant_id")
+    op.create_table(
+        "_dlt_pipeline_state",
+        sa.Column("version", sa.BIGINT(), autoincrement=False, nullable=False),
+        sa.Column("engine_version", sa.BIGINT(), autoincrement=False, nullable=False),
+        sa.Column("pipeline_name", sa.TEXT(), autoincrement=False, nullable=False),
+        sa.Column("state", sa.TEXT(), autoincrement=False, nullable=False),
+        sa.Column(
+            "created_at", postgresql.TIMESTAMP(timezone=True), autoincrement=False, nullable=False
+        ),
+        sa.Column("version_hash", sa.TEXT(), autoincrement=False, nullable=True),
+        sa.Column("_dlt_load_id", sa.TEXT(), autoincrement=False, nullable=False),
+        sa.Column("_dlt_id", sa.VARCHAR(length=128), autoincrement=False, nullable=False),
+    )
+    op.create_table(
+        "_dlt_version",
+        sa.Column("version", sa.BIGINT(), autoincrement=False, nullable=False),
+        sa.Column("engine_version", sa.BIGINT(), autoincrement=False, nullable=False),
+        sa.Column(
+            "inserted_at", postgresql.TIMESTAMP(timezone=True), autoincrement=False, nullable=False
+        ),
+        sa.Column("schema_name", sa.TEXT(), autoincrement=False, nullable=False),
+        sa.Column("version_hash", sa.TEXT(), autoincrement=False, nullable=False),
+        sa.Column("schema", sa.TEXT(), autoincrement=False, nullable=False),
+    )
+    op.create_table(
+        "_dlt_loads",
+        sa.Column("load_id", sa.TEXT(), autoincrement=False, nullable=False),
+        sa.Column("schema_name", sa.TEXT(), autoincrement=False, nullable=True),
+        sa.Column("status", sa.BIGINT(), autoincrement=False, nullable=False),
+        sa.Column(
+            "inserted_at", postgresql.TIMESTAMP(timezone=True), autoincrement=False, nullable=False
+        ),
+        sa.Column("schema_version_hash", sa.TEXT(), autoincrement=False, nullable=True),
+    )
+    op.create_table(
+        "file_metadata",
+        sa.Column("id", sa.TEXT(), autoincrement=False, nullable=False),
+        sa.Column("name", sa.TEXT(), autoincrement=False, nullable=True),
+        sa.Column("file_path", sa.TEXT(), autoincrement=False, nullable=True),
+        sa.Column("extension", sa.TEXT(), autoincrement=False, nullable=True),
+        sa.Column("mime_type", sa.TEXT(), autoincrement=False, nullable=True),
+        sa.Column("content_hash", sa.TEXT(), autoincrement=False, nullable=True),
+        sa.Column("owner_id", sa.TEXT(), autoincrement=False, nullable=True),
+        sa.Column("_dlt_load_id", sa.TEXT(), autoincrement=False, nullable=False),
+        sa.Column("_dlt_id", sa.VARCHAR(length=128), autoincrement=False, nullable=False),
+        sa.Column("node_set", sa.TEXT(), autoincrement=False, nullable=True),
+    )