added boilerplate for user management

cognee/api/client.py

@@ -43,6 +43,77 @@
     allow_headers=["*"],
 )
 
+from contextlib import asynccontextmanager
+
+from fastapi import Depends, FastAPI
+
+from cognee.infrastructure.databases.relational.user_authentication.authentication_db import User, create_db_and_tables
+from cognee.infrastructure.databases.relational.user_authentication.schemas import UserCreate, UserRead, UserUpdate
+from cognee.infrastructure.databases.relational.user_authentication.users import auth_backend, current_active_user, fastapi_users
+
+app.include_router(
+    fastapi_users.get_auth_router(auth_backend), prefix="/auth/jwt", tags=["auth"]
+)
+app.include_router(
+    fastapi_users.get_register_router(UserRead, UserCreate),
+    prefix="/auth",
+    tags=["auth"],
+)
+app.include_router(
+    fastapi_users.get_reset_password_router(),
+    prefix="/auth",
+    tags=["auth"],
+)
+app.include_router(
+    fastapi_users.get_verify_router(UserRead),
+    prefix="/auth",
+    tags=["auth"],
+)
+app.include_router(
+    fastapi_users.get_users_router(UserRead, UserUpdate),
+    prefix="/users",
+    tags=["users"],
+)
+
+
+@app.get("/authenticated-route")
+async def authenticated_route(user: User = Depends(current_active_user)):
+    return {"message": f"Hello {user.email}!"}
+
+@asynccontextmanager
+async def lifespan(app: FastAPI):
+    # Not needed if you setup a migration system like Alembic
+    await create_db_and_tables()
+    yield
+app.include_router(
+    fastapi_users.get_auth_router(auth_backend), prefix="/auth/jwt", tags=["auth"]
+)
+app.include_router(
+    fastapi_users.get_register_router(UserRead, UserCreate),
+    prefix="/auth",
+    tags=["auth"],
+)
+app.include_router(
+    fastapi_users.get_reset_password_router(),
+    prefix="/auth",
+    tags=["auth"],
+)
+app.include_router(
+    fastapi_users.get_verify_router(UserRead),
+    prefix="/auth",
+    tags=["auth"],
+)
+app.include_router(
+    fastapi_users.get_users_router(UserRead, UserUpdate),
+    prefix="/users",
+    tags=["users"],
+)
+
+
+@app.get("/authenticated-route")
+async def authenticated_route(user: User = Depends(current_active_user)):
+    return {"message": f"Hello {user.email}!"}
+
 @app.get("/")
 async def root():
     """

cognee/infrastructure/databases/relational/config.py

@@ -11,9 +11,11 @@ class RelationalConfig(BaseSettings):
     db_port: str =  "5432"
     db_user: str = "cognee"
     db_password: str =  "cognee"
-    database_engine: object = create_relational_engine(db_path, db_name)
+    db_provider: str = "duckdb"
+    database_engine: object = create_relational_engine(db_path, db_name, db_provider)
     db_file_path: str = os.path.join(db_path, db_name)
 
+
     model_config = SettingsConfigDict(env_file = ".env", extra = "allow")
 
     def create_engine(self):
@@ -29,6 +31,7 @@ def to_dict(self) -> dict:
             "db_user": self.db_user,
             "db_password": self.db_password,
             "db_engine": self.database_engine,
+            "db_provider": self.db_provider,
         }
 
 @lru_cache

cognee/infrastructure/databases/relational/create_relational_engine.py

@@ -1,10 +1,33 @@
+from enum import Enum
+
+from cognee.infrastructure.databases.relational.sqlalchemy.SqlAlchemyAdapter import SQLAlchemyAdapter
 from cognee.infrastructure.files.storage import LocalStorage
-from cognee.infrastructure.databases.relational import DuckDBAdapter
+from cognee.infrastructure.databases.relational import DuckDBAdapter, get_relationaldb_config
+
+
+class DBProvider(Enum):
+    DUCKDB = "duckdb"
+    POSTGRES = "postgres"
+
+
 
-def create_relational_engine(db_path: str, db_name: str):
+def create_relational_engine(db_path: str, db_name: str, db_provider:str):
     LocalStorage.ensure_directory_exists(db_path)
 
-    return DuckDBAdapter(
-        db_name = db_name,
-        db_path = db_path,
-    )
+    llm_config = get_relationaldb_config()
+
+    provider = DBProvider(llm_config.llm_provider)
+
+
+    if provider == DBProvider.DUCKDB:
+
+        return DuckDBAdapter(
+            db_name = db_name,
+            db_path = db_path,
+        )
+    elif provider == DBProvider.POSTGRES:
+        return SQLAlchemyAdapter(
+            db_name = db_name,
+            db_path = db_path,
+            db_type = db_provider,
+        )

cognee/infrastructure/databases/relational/sqlalchemy/SqlAlchemyAdapter.py

@@ -0,0 +1,125 @@
+import os
+from sqlalchemy import create_engine, MetaData, Table, Column, String, Boolean, TIMESTAMP, text
+from sqlalchemy.orm import sessionmaker
+
+class SQLAlchemyAdapter():
+    def __init__(self, db_type: str, db_path: str, db_name: str):
+        self.db_location = os.path.abspath(os.path.join(db_path, db_name))
+        self.engine = create_engine(f"{db_type}:///{self.db_location}")
+        self.Session = sessionmaker(bind=self.engine)
+
+    def get_datasets(self):
+        with self.engine.connect() as connection:
+            result = connection.execute(text("SELECT DISTINCT schema_name FROM information_schema.tables;"))
+            tables = [row['schema_name'] for row in result]
+        return list(
+            filter(
+                lambda schema_name: not schema_name.endswith("staging") and schema_name != "cognee",
+                tables
+            )
+        )
+
+    def get_files_metadata(self, dataset_name: str):
+        with self.engine.connect() as connection:
+            result = connection.execute(text(f"SELECT id, name, file_path, extension, mime_type FROM {dataset_name}.file_metadata;"))
+            return [dict(row) for row in result]
+
+    def create_table(self, schema_name: str, table_name: str, table_config: list[dict]):
+        fields_query_parts = [f"{item['name']} {item['type']}" for item in table_config]
+        with self.engine.connect() as connection:
+            connection.execute(text(f"CREATE SCHEMA IF NOT EXISTS {schema_name};"))
+            connection.execute(text(f"CREATE TABLE IF NOT EXISTS {schema_name}.{table_name} ({', '.join(fields_query_parts)});"))
+
+    def delete_table(self, table_name: str):
+        with self.engine.connect() as connection:
+            connection.execute(text(f"DROP TABLE IF EXISTS {table_name};"))
+
+    def insert_data(self, schema_name: str, table_name: str, data: list[dict]):
+        columns = ", ".join(data[0].keys())
+        values = ", ".join([f"({', '.join([f':{key}' for key in row.keys()])})" for row in data])
+        insert_query = text(f"INSERT INTO {schema_name}.{table_name} ({columns}) VALUES {values};")
+        with self.engine.connect() as connection:
+            connection.execute(insert_query, data)
+
+    def get_data(self, table_name: str, filters: dict = None):
+        with self.engine.connect() as connection:
+            query = f"SELECT * FROM {table_name}"
+            if filters:
+                filter_conditions = " AND ".join([
+                    f"{key} IN ({', '.join([f':{key}{i}' for i in range(len(value))])})" if isinstance(value, list)
+                    else f"{key} = :{key}" for key, value in filters.items()
+                ])
+                query += f" WHERE {filter_conditions};"
+                query = text(query)
+                results = connection.execute(query, filters)
+            else:
+                query += ";"
+                query = text(query)
+                results = connection.execute(query)
+            return {result["data_id"]: result["status"] for result in results}
+
+    def execute_query(self, query):
+        with self.engine.connect() as connection:
+            result = connection.execute(text(query))
+            return [dict(row) for row in result]
+
+    def load_cognify_data(self, data):
+        metadata = MetaData()
+        cognify_table = Table(
+            'cognify', metadata,
+            Column('document_id', String),
+            Column('layer_id', String),
+            Column('created_at', TIMESTAMP, server_default=text('CURRENT_TIMESTAMP')),
+            Column('updated_at', TIMESTAMP, nullable=True, default=None),
+            Column('processed', Boolean, default=False),
+            Column('document_id_target', String, nullable=True)
+        )
+        metadata.create_all(self.engine)
+        insert_query = cognify_table.insert().values(document_id=text(':document_id'), layer_id=text(':layer_id'))
+        with self.engine.connect() as connection:
+            connection.execute(insert_query, data)
+
+    def fetch_cognify_data(self, excluded_document_id: str):
+        with self.engine.connect() as connection:
+            connection.execute(text("""
+                CREATE TABLE IF NOT EXISTS cognify (
+                    document_id STRING,
+                    layer_id STRING,
+                    created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
+                    updated_at TIMESTAMP DEFAULT NULL,
+                    processed BOOLEAN DEFAULT FALSE,
+                    document_id_target STRING NULL
+                );
+            """))
+            query = text(f"""
+                SELECT document_id, layer_id, created_at, updated_at, processed
+                FROM cognify
+                WHERE document_id != :excluded_document_id AND processed = FALSE;
+            """)
+            records = connection.execute(query, {'excluded_document_id': excluded_document_id}).fetchall()
+            if records:
+                document_ids = tuple(record['document_id'] for record in records)
+                update_query = text(f"UPDATE cognify SET processed = TRUE WHERE document_id IN :document_ids;")
+                connection.execute(update_query, {'document_ids': document_ids})
+            return [dict(record) for record in records]
+
+    def delete_cognify_data(self):
+        with self.engine.connect() as connection:
+            connection.execute(text("""
+                CREATE TABLE IF NOT EXISTS cognify (
+                    document_id STRING,
+                    layer_id STRING,
+                    created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
+                    updated_at TIMESTAMP DEFAULT NULL,
+                    processed BOOLEAN DEFAULT FALSE,
+                    document_id_target STRING NULL
+                );
+            """))
+            connection.execute(text("DELETE FROM cognify;"))
+            connection.execute(text("DROP TABLE cognify;"))
+
+    def delete_database(self):
+        from cognee.infrastructure.files.storage import LocalStorage
+        LocalStorage.remove(self.db_location)
+        if LocalStorage.file_exists(self.db_location + ".wal"):
+            LocalStorage.remove(self.db_location + ".wal")

cognee/infrastructure/databases/relational/user_authentication/authentication_db.py

@@ -0,0 +1,41 @@
+from typing import AsyncGenerator
+
+from fastapi import Depends
+from fastapi_users.db import SQLAlchemyBaseUserTableUUID, SQLAlchemyUserDatabase
+from sqlalchemy.ext.asyncio import AsyncSession, async_sessionmaker, create_async_engine
+from sqlalchemy.orm import DeclarativeBase
+
+from cognee.infrastructure.databases.relational import get_relationaldb_config
+from cognee.infrastructure.databases.relational.create_relational_engine import create_relational_engine
+
+DATABASE_URL = "sqlite+aiosqlite:///./test.db"
+
+
+class Base(DeclarativeBase):
+    pass
+
+
+class User(SQLAlchemyBaseUserTableUUID, Base):
+    pass
+
+
+
+llm_config = get_relationaldb_config()
+
+
+engine = create_relational_engine(llm_config.db_path, llm_config.db_name, llm_config.db_provider)
+async_session_maker = async_sessionmaker(engine, expire_on_commit=False)
+
+
+async def create_db_and_tables():
+    async with engine.begin() as conn:
+        await conn.run_sync(Base.metadata.create_all)
+
+
+async def get_async_session() -> AsyncGenerator[AsyncSession, None]:
+    async with async_session_maker() as session:
+        yield session
+
+
+async def get_user_db(session: AsyncSession = Depends(get_async_session)):
+    yield SQLAlchemyUserDatabase(session, User)

cognee/infrastructure/databases/relational/user_authentication/schemas.py

@@ -0,0 +1,15 @@
+import uuid
+
+from fastapi_users import schemas
+
+
+class UserRead(schemas.BaseUser[uuid.UUID]):
+    pass
+
+
+class UserCreate(schemas.BaseUserCreate):
+    pass
+
+
+class UserUpdate(schemas.BaseUserUpdate):
+    pass

cognee/infrastructure/databases/relational/user_authentication/users.py

@@ -0,0 +1,55 @@
+import uuid
+from typing import Optional
+
+from fastapi import Depends, Request
+from fastapi_users import BaseUserManager, FastAPIUsers, UUIDIDMixin, models
+from fastapi_users.authentication import (
+    AuthenticationBackend,
+    BearerTransport,
+    JWTStrategy,
+)
+from fastapi_users.db import SQLAlchemyUserDatabase
+
+from app.db import User, get_user_db
+
+SECRET = "SECRET"
+
+
+class UserManager(UUIDIDMixin, BaseUserManager[User, uuid.UUID]):
+    reset_password_token_secret = SECRET
+    verification_token_secret = SECRET
+
+    async def on_after_register(self, user: User, request: Optional[Request] = None):
+        print(f"User {user.id} has registered.")
+
+    async def on_after_forgot_password(
+        self, user: User, token: str, request: Optional[Request] = None
+    ):
+        print(f"User {user.id} has forgot their password. Reset token: {token}")
+
+    async def on_after_request_verify(
+        self, user: User, token: str, request: Optional[Request] = None
+    ):
+        print(f"Verification requested for user {user.id}. Verification token: {token}")
+
+
+async def get_user_manager(user_db: SQLAlchemyUserDatabase = Depends(get_user_db)):
+    yield UserManager(user_db)
+
+
+bearer_transport = BearerTransport(tokenUrl="auth/jwt/login")
+
+
+def get_jwt_strategy() -> JWTStrategy[models.UP, models.ID]:
+    return JWTStrategy(secret=SECRET, lifetime_seconds=3600)
+
+
+auth_backend = AuthenticationBackend(
+    name="jwt",
+    transport=bearer_transport,
+    get_strategy=get_jwt_strategy,
+)
+
+fastapi_users = FastAPIUsers[User, uuid.UUID](get_user_manager, [auth_backend])
+
+current_active_user = fastapi_users.current_user(active=True)