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

Fixes to the ACL model

Browse source
This commit is contained in:
Vasilije 2024-07-27 10:09:01 +02:00
parent 9616545de9
commit 218d322c75
7 changed files with 158 additions and 128 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

2
cognee/api/v1/authenticate_user/authenticate_user.py
View file

@ -1,4 +1,4 @@
from cognee.infrastructure.databases.relational.user_authentication.authentication_db import authenticate_user_method
from cognee.infrastructure.databases.relational.user_authentication.users import authenticate_user_method
async def authenticate_user():

20
cognee/api/v1/cognify/cognify_v2.py
View file

@ -7,8 +7,8 @@ from typing import Union
from sqlalchemy.ext.asyncio import AsyncSession
from cognee.infrastructure.databases.graph import get_graph_config
from cognee.infrastructure.databases.relational.user_authentication.authentication_db import async_session_maker
from cognee.infrastructure.databases.relational.user_authentication.users import get_user_permissions, fastapi_users
# from cognee.infrastructure.databases.relational.user_authentication.authentication_db import async_session_maker
# from cognee.infrastructure.databases.relational.user_authentication.users import get_user_permissions, fastapi_users
from cognee.modules.cognify.config import get_cognify_config
from cognee.infrastructure.databases.relational.config import get_relationaldb_config
from cognee.modules.data.processing.document_types.AudioDocument import AudioDocument
@ -37,14 +37,14 @@ class PermissionDeniedException(Exception):
super().__init__(self.message)
async def cognify(datasets: Union[str, list[str]] = None, root_node_id: str = None, user_id:str="default_user"):
session: AsyncSession = async_session_maker()
user = await fastapi_users.get_user_manager.get(user_id)
user_permissions = await get_user_permissions(user, session)
hash_object = hashlib.sha256(user.encode())
hashed_user_id = hash_object.hexdigest()
required_permission = "write"
if required_permission not in user_permissions:
raise PermissionDeniedException("Not enough permissions")
# session: AsyncSession = async_session_maker()
# user = await fastapi_users.get_user_manager.get(user_id)
# user_permissions = await get_user_permissions(user, session)
# hash_object = hashlib.sha256(user.encode())
# hashed_user_id = hash_object.hexdigest()
# required_permission = "write"
# if required_permission not in user_permissions:
# raise PermissionDeniedException("Not enough permissions")
relational_config = get_relationaldb_config()
db_engine = relational_config.database_engine

19
cognee/api/v1/create_user/create_user.py
View file

@ -1,7 +1,22 @@
from cognee.infrastructure.databases.relational.user_authentication.authentication_db import create_user_method
from cognee.infrastructure.databases.relational.user_authentication.users import create_user_method
async def create_user(email: str, password: str, is_superuser: bool = False):
output = await create_user_method(email=email, password=password, is_superuser=is_superuser)
return output
return output
if __name__ == "__main__":
import asyncio
# Define an example user
example_email = "example@example.com"
example_password = "securepassword123"
example_is_superuser = False
# Create an event loop and run the create_user function
loop = asyncio.get_event_loop()
result = loop.run_until_complete(create_user(example_email, example_password, example_is_superuser))
# Print the result
print(result)

2
cognee/api/v1/verify_user_token/verify_user_token.py
View file

@ -1,4 +1,4 @@
from cognee.infrastructure.databases.relational.user_authentication.authentication_db import user_check_token
from cognee.infrastructure.databases.relational.user_authentication.users import user_check_token

17
cognee/infrastructure/data/models/DatasetData.py
View file

@ -1,14 +1,19 @@
from uuid import uuid4
from datetime import datetime, timezone
from sqlalchemy import Column, DateTime, UUID, ForeignKey
from sqlalchemy import Column, DateTime, UUID, ForeignKey, PrimaryKeyConstraint, UniqueConstraint
from sqlalchemy.orm import relationship
from cognee.infrastructure.databases.relational import Base
class DatasetData(Base):
__tablename__ = "dataset_data"
id = Column(UUID, primary_key = True, default = uuid4())
id = Column(UUID, primary_key=True, default=uuid4)
created_at = Column(DateTime, default=datetime.now(timezone.utc))
dataset_id = Column(UUID, ForeignKey("dataset.id"), nullable=False)
data_id = Column(UUID, ForeignKey("data.id"), nullable=False)
__table_args__ = (
UniqueConstraint('dataset_id', 'data_id', name='uix_dataset_data'),
)
created_at = Column(DateTime, default = datetime.now(timezone.utc))
dataset_id = Column("dataset", UUID, ForeignKey("dataset.id"), primary_key = True)
data_id = Column("data", UUID, ForeignKey("data.id"), primary_key = True)
acls = relationship('ACL', back_populates='document')

116
cognee/infrastructure/databases/relational/user_authentication/authentication_db.py
View file

@ -14,7 +14,7 @@ from fastapi_users.exceptions import UserAlreadyExists
from fastapi_users.authentication import JWTStrategy
from sqlalchemy.orm import DeclarativeBase, sessionmaker, Session, relationship
from cognee.infrastructure.databases.relational.user_authentication.schemas import UserRead, UserCreate
from cognee.infrastructure.databases.relational.user_authentication.users import get_user_manager, get_jwt_strategy
# from cognee.infrastructure.databases.relational.user_authentication.users import get_user_manager, get_jwt_strategy
from fastapi.security import OAuth2PasswordRequestForm
# Association table for many-to-many relationship between users and groups
user_group = Table('user_group', Base.metadata,
@ -31,14 +31,15 @@ group_permission = Table('group_permission', Base.metadata,
class User(SQLAlchemyBaseUserTableUUID, Base):
__tablename__ = 'users'
groups = relationship('Group', secondary=user_group, back_populates='users')
acls = relationship('ACL', back_populates='user')
class Group(Base):
__tablename__ = 'groups'
id = Column(UUID, primary_key=True, index=True)
id = Column(UUID, primary_key=True, index=True, default=uuid.uuid4)
name = Column(String, unique=True, index=True)
users = relationship('users', secondary=user_group, back_populates='groups')
users = relationship('User', secondary=user_group, back_populates='groups')
permissions = relationship('Permission', secondary=group_permission, back_populates='groups')
acls = relationship('ACL', back_populates='group')
class Permission(Base):
__tablename__ = 'permissions'
id = Column(UUID, primary_key=True, index=True)
@ -47,15 +48,14 @@ class Permission(Base):
class ACL(Base):
__tablename__ = 'acls'
id = Column(UUID, primary_key=True, index=True)
document_id = Column(UUID, ForeignKey('documents.id'))
id = Column(UUID, primary_key=True, index=True, default=uuid.uuid4)
document_id = Column(UUID, ForeignKey('dataset_data.id'))
user_id = Column(UUID, ForeignKey('users.id'), nullable=True)
group_id = Column(UUID, ForeignKey('groups.id'), nullable=True)
permission = Column(String) # 'read', 'write', 'execute'
document = relationship('dataset_data', back_populates='acls')
user = relationship('users', back_populates='acls')
group = relationship('groups', back_populates='acls')
document = relationship('DatasetData', back_populates='acls')
user = relationship('User', back_populates='acls')
group = relationship('Group', back_populates='acls')
relational_config = get_relationaldb_config()
@ -77,99 +77,3 @@ async def get_user_db(session: AsyncSession = Depends(get_async_session)):
yield SQLAlchemyUserDatabase(session, User)
async def hash_password(password: str) -> str:
return hashlib.sha256(password.encode()).hexdigest()
# Define context managers for dependencies
get_async_session_context = asynccontextmanager(get_async_session)
get_user_db_context = asynccontextmanager(get_user_db)
get_user_manager_context = asynccontextmanager(get_user_manager)
async def create_user_method(email: str, password: str, is_superuser: bool = False):
try:
async with get_async_session_context() as session:
async with get_user_db_context(session) as user_db:
async with get_user_manager_context(user_db) as user_manager:
user = await user_manager.create(
UserCreate(email=email, password=password, is_superuser=is_superuser)
)
print(f"User created: {user.email}")
except UserAlreadyExists:
print(f"User {email} already exists")
async def authenticate_user_method(email: str, password: str) -> Optional[User]:
try:
async with get_async_session_context() as session:
async with get_user_db_context(session) as user_db:
async with get_user_manager_context(user_db) as user_manager:
credentials = OAuth2PasswordRequestForm(username=email, password=password)
user = await user_manager.authenticate(credentials)
if user is None or not user.is_active:
return None
return user
except Exception as e:
print(f"Error during authentication: {e}")
return None
async def reset_user_password_method(email: str, new_password: str) -> bool:
async with get_async_session_context() as session:
async with get_user_db_context(session) as user_db:
user = await user_db.get_by_email(email)
if not user:
return False
user.hashed_password = await hash_password(new_password)
await user_db.update(user)
return True
# async def generate_verification_token(email: str, tokens_db: dict) -> str:
# async with get_async_session_context() as session:
# async with get_user_db_context(session) as user_db:
# if not await user_db.get_by_email(email):
# raise ValueError("User does not exist")
# token = str(uuid.uuid4())
# tokens_db[token] = email
# return token
# async def verify_user_method(token: str, tokens_db: dict) -> bool:
# async with get_async_session_context() as session:
# async with get_user_db_context(session) as user_db:
# email = tokens_db.get(token)
# if not email or not await user_db.get_by_email(email):
# return False
# user = await user_db.get_by_email(email)
# user.is_verified = True
# await user_db.update(user)
# return True
async def user_create_token(user: User) -> Optional[str]:
try:
async with get_async_session_context() as session:
async with get_user_db_context(session) as user_db:
async with get_user_manager_context(user_db) as user_manager:
if user is None:
return None
strategy = get_jwt_strategy()
token = await strategy.write_token(user)
if token is not None:
return token
else:
return None
except:
return None
async def user_check_token(token: str) -> bool:
try:
async with get_async_session_context() as session:
async with get_user_db_context(session) as user_db:
async with get_user_manager_context(user_db) as user_manager:
if token is None:
return False
strategy = get_jwt_strategy()
user = await strategy.read_token(token, user_manager)
if user is None or not user.is_active:
return False
else:
return True
except:
return False

110
cognee/infrastructure/databases/relational/user_authentication/users.py
View file

@ -1,3 +1,4 @@
import hashlib
import uuid
from typing import Optional
@ -8,10 +9,16 @@ from fastapi_users.authentication import (
BearerTransport,
JWTStrategy,
)
from fastapi_users.exceptions import UserAlreadyExists
from fastapi_users.db import SQLAlchemyUserDatabase
from fastapi import Depends, HTTPException, status
from sqlalchemy.orm import Session
from cognee.infrastructure.databases.relational.user_authentication.authentication_db import User, get_user_db
from cognee.infrastructure.databases.relational.user_authentication.authentication_db import User, get_user_db, \
get_async_session
from fastapi.security import OAuth2PasswordRequestForm
from fastapi_users.authentication import JWTStrategy
from cognee.infrastructure.databases.relational.user_authentication.schemas import UserRead, UserCreate
from contextlib import asynccontextmanager
SECRET = "SECRET"
@ -68,4 +75,103 @@ def has_permission(permission: str):
if permission not in user_permissions:
raise HTTPException(status_code=status.HTTP_403_FORBIDDEN, detail="Not enough permissions")
return True
return Depends(permission_checker)
return Depends(permission_checker)
async def hash_password(password: str) -> str:
return hashlib.sha256(password.encode()).hexdigest()
# Define context managers for dependencies
get_async_session_context = asynccontextmanager(get_async_session)
get_user_db_context = asynccontextmanager(get_user_db)
get_user_manager_context = asynccontextmanager(get_user_manager)
async def create_user_method(email: str, password: str, is_superuser: bool = False):
try:
async with get_async_session_context() as session:
async with get_user_db_context(session) as user_db:
async with get_user_manager_context(user_db) as user_manager:
user = await user_manager.create(
UserCreate(email=email, password=password, is_superuser=is_superuser)
)
print(f"User created: {user.email}")
except UserAlreadyExists:
print(f"User {email} already exists")
async def authenticate_user_method(email: str, password: str) -> Optional[User]:
try:
async with get_async_session_context() as session:
async with get_user_db_context(session) as user_db:
async with get_user_manager_context(user_db) as user_manager:
credentials = OAuth2PasswordRequestForm(username=email, password=password)
user = await user_manager.authenticate(credentials)
if user is None or not user.is_active:
return None
return user
except Exception as e:
print(f"Error during authentication: {e}")
return None
async def reset_user_password_method(email: str, new_password: str) -> bool:
async with get_async_session_context() as session:
async with get_user_db_context(session) as user_db:
user = await user_db.get_by_email(email)
if not user:
return False
user.hashed_password = await hash_password(new_password)
await user_db.update(user)
return True
# async def generate_verification_token(email: str, tokens_db: dict) -> str:
# async with get_async_session_context() as session:
# async with get_user_db_context(session) as user_db:
# if not await user_db.get_by_email(email):
# raise ValueError("User does not exist")
# token = str(uuid.uuid4())
# tokens_db[token] = email
# return token
# async def verify_user_method(token: str, tokens_db: dict) -> bool:
# async with get_async_session_context() as session:
# async with get_user_db_context(session) as user_db:
# email = tokens_db.get(token)
# if not email or not await user_db.get_by_email(email):
# return False
# user = await user_db.get_by_email(email)
# user.is_verified = True
# await user_db.update(user)
# return True
async def user_create_token(user: User) -> Optional[str]:
try:
async with get_async_session_context() as session:
async with get_user_db_context(session) as user_db:
async with get_user_manager_context(user_db) as user_manager:
if user is None:
return None
strategy = get_jwt_strategy()
token = await strategy.write_token(user)
if token is not None:
return token
else:
return None
except:
return None
async def user_check_token(token: str) -> bool:
try:
async with get_async_session_context() as session:
async with get_user_db_context(session) as user_db:
async with get_user_manager_context(user_db) as user_manager:
if token is None:
return False
strategy = get_jwt_strategy()
user = await strategy.read_token(token, user_manager)
if user is None or not user.is_active:
return False
else:
return True
except:
return False