diff --git a/cognee/api/v1/authenticate_user/authenticate_user.py b/cognee/api/v1/authenticate_user/authenticate_user.py index d7dd5a22c..7d493adfc 100644 --- a/cognee/api/v1/authenticate_user/authenticate_user.py +++ b/cognee/api/v1/authenticate_user/authenticate_user.py @@ -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(): diff --git a/cognee/api/v1/cognify/cognify_v2.py b/cognee/api/v1/cognify/cognify_v2.py index 20e2624cb..ecb9fd7c4 100644 --- a/cognee/api/v1/cognify/cognify_v2.py +++ b/cognee/api/v1/cognify/cognify_v2.py @@ -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 diff --git a/cognee/api/v1/create_user/create_user.py b/cognee/api/v1/create_user/create_user.py index 41074751d..3901aa593 100644 --- a/cognee/api/v1/create_user/create_user.py +++ b/cognee/api/v1/create_user/create_user.py @@ -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 \ No newline at end of file + 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) \ No newline at end of file diff --git a/cognee/api/v1/verify_user_token/verify_user_token.py b/cognee/api/v1/verify_user_token/verify_user_token.py index bc1e9bcd9..b5a0f2228 100644 --- a/cognee/api/v1/verify_user_token/verify_user_token.py +++ b/cognee/api/v1/verify_user_token/verify_user_token.py @@ -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 diff --git a/cognee/infrastructure/data/models/DatasetData.py b/cognee/infrastructure/data/models/DatasetData.py index 552642686..9e586dcef 100644 --- a/cognee/infrastructure/data/models/DatasetData.py +++ b/cognee/infrastructure/data/models/DatasetData.py @@ -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') diff --git a/cognee/infrastructure/databases/relational/user_authentication/authentication_db.py b/cognee/infrastructure/databases/relational/user_authentication/authentication_db.py index b83327ed7..bbc0a01f7 100644 --- a/cognee/infrastructure/databases/relational/user_authentication/authentication_db.py +++ b/cognee/infrastructure/databases/relational/user_authentication/authentication_db.py @@ -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 \ No newline at end of file diff --git a/cognee/infrastructure/databases/relational/user_authentication/users.py b/cognee/infrastructure/databases/relational/user_authentication/users.py index 685cff7fb..4f7e5dce7 100644 --- a/cognee/infrastructure/databases/relational/user_authentication/users.py +++ b/cognee/infrastructure/databases/relational/user_authentication/users.py @@ -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) \ No newline at end of file + 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 +