Permissions db migration (#992)

<!-- .github/pull_request_template.md -->

## Description
Remove cognee models from permissions migration

## DCO Affirmation
I affirm that all code in every commit of this pull request conforms to
the terms of the Topoteretes Developer Certificate of Origin.
This commit is contained in:
Igor Ilic 2025-06-17 16:27:26 +02:00 committed by GitHub
parent f3c8b66c3a
commit f3455e9f86
No known key found for this signature in database
GPG key ID: B5690EEEBB952194

View file

@ -24,34 +24,86 @@ def _now():
return datetime.now(timezone.utc) 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: def _ensure_permission(conn, permission_name) -> str:
""" """
Return the permission.id for the given name, creating the row if needed. Return the permission.id for the given name, creating the row if needed.
""" """
from cognee.modules.users.models import Permission permissions_table = sa.Table(
"permissions",
row = conn.execute(sa.select(Permission).filter(Permission.name == permission_name)).fetchone() 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: if row is None:
permission_id = uuid4() permission_id = uuid4()
permissions_table = sa.Table(
"acls",
sa.MetaData(),
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")),
)
op.bulk_insert( op.bulk_insert(
permissions_table, permissions_table,
[ [
@ -59,7 +111,6 @@ def _ensure_permission(conn, permission_name) -> str:
"id": permission_id, "id": permission_id,
"name": permission_name, "name": permission_name,
"created_at": _now(), "created_at": _now(),
"updated_at": _now(),
} }
], ],
) )
@ -73,7 +124,6 @@ def _build_acl_row(*, user_id, target_id, permission_id, target_col) -> dict:
return { return {
"id": uuid4(), "id": uuid4(),
"created_at": _now(), "created_at": _now(),
"updated_at": _now(),
"principal_id": user_id, "principal_id": user_id,
target_col: target_id, target_col: target_id,
"permission_id": permission_id, "permission_id": permission_id,
@ -114,9 +164,10 @@ def upgrade() -> None:
sa.Column("dataset_id", UUID, sa.ForeignKey("datasets.id", ondelete="CASCADE")), sa.Column("dataset_id", UUID, sa.ForeignKey("datasets.id", ondelete="CASCADE")),
) )
from cognee.modules.data.models import Dataset # 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
datasets = conn.execute(sa.select(Dataset)).fetchall() dataset_table = _define_dataset_table()
datasets = conn.execute(sa.select(dataset_table)).fetchall()
if not datasets: if not datasets:
return return
@ -152,9 +203,10 @@ def downgrade() -> None:
sa.Column("data_id", UUID, sa.ForeignKey("data.id", ondelete="CASCADE")), sa.Column("data_id", UUID, sa.ForeignKey("data.id", ondelete="CASCADE")),
) )
from cognee.modules.data.models import Data # 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 = conn.execute(sa.select(Data)).fetchall() data_table = _define_data_table()
data = conn.execute(sa.select(data_table)).fetchall()
if not data: if not data:
return return