cognee 0.2.3.dev1__py3-none-any.whl → 0.3.0__py3-none-any.whl

cognee/modules/sync/methods/get_sync_operation.py

@@ -0,0 +1,107 @@
+from uuid import UUID
+from typing import List, Optional
+from sqlalchemy import select, desc, and_
+from cognee.modules.sync.models import SyncOperation, SyncStatus
+from cognee.infrastructure.databases.relational import get_relational_engine
+
+
+async def get_sync_operation(run_id: str) -> Optional[SyncOperation]:
+    """
+    Get a sync operation by its run_id.
+
+    Args:
+        run_id: The public run_id of the sync operation
+
+    Returns:
+        SyncOperation: The sync operation record, or None if not found
+    """
+    db_engine = get_relational_engine()
+
+    async with db_engine.get_async_session() as session:
+        query = select(SyncOperation).where(SyncOperation.run_id == run_id)
+        result = await session.execute(query)
+        return result.scalars().first()
+
+
+async def get_user_sync_operations(
+    user_id: UUID, limit: int = 50, offset: int = 0
+) -> List[SyncOperation]:
+    """
+    Get sync operations for a specific user, ordered by most recent first.
+
+    Args:
+        user_id: UUID of the user
+        limit: Maximum number of records to return
+        offset: Number of records to skip
+
+    Returns:
+        List[SyncOperation]: List of sync operations for the user
+    """
+    db_engine = get_relational_engine()
+
+    async with db_engine.get_async_session() as session:
+        query = (
+            select(SyncOperation)
+            .where(SyncOperation.user_id == user_id)
+            .order_by(desc(SyncOperation.created_at))
+            .limit(limit)
+            .offset(offset)
+        )
+        result = await session.execute(query)
+        return list(result.scalars().all())
+
+
+async def get_sync_operations_by_dataset(
+    dataset_id: UUID, limit: int = 50, offset: int = 0
+) -> List[SyncOperation]:
+    """
+    Get sync operations for a specific dataset.
+
+    Args:
+        dataset_id: UUID of the dataset
+        limit: Maximum number of records to return
+        offset: Number of records to skip
+
+    Returns:
+        List[SyncOperation]: List of sync operations for the dataset
+    """
+    db_engine = get_relational_engine()
+
+    async with db_engine.get_async_session() as session:
+        query = (
+            select(SyncOperation)
+            .where(SyncOperation.dataset_id == dataset_id)
+            .order_by(desc(SyncOperation.created_at))
+            .limit(limit)
+            .offset(offset)
+        )
+        result = await session.execute(query)
+        return list(result.scalars().all())
+
+
+async def get_running_sync_operations_for_user(user_id: UUID) -> List[SyncOperation]:
+    """
+    Get all currently running sync operations for a specific user.
+    Checks for operations with STARTED or IN_PROGRESS status.
+
+    Args:
+        user_id: UUID of the user
+
+    Returns:
+        List[SyncOperation]: List of running sync operations for the user
+    """
+    db_engine = get_relational_engine()
+
+    async with db_engine.get_async_session() as session:
+        query = (
+            select(SyncOperation)
+            .where(
+                and_(
+                    SyncOperation.user_id == user_id,
+                    SyncOperation.status.in_([SyncStatus.STARTED, SyncStatus.IN_PROGRESS]),
+                )
+            )
+            .order_by(desc(SyncOperation.created_at))
+        )
+        result = await session.execute(query)
+        return list(result.scalars().all())

cognee/modules/sync/methods/update_sync_operation.py

@@ -0,0 +1,248 @@
+import asyncio
+from typing import Optional, List
+from datetime import datetime, timezone
+from sqlalchemy import select
+from sqlalchemy.exc import SQLAlchemyError, DisconnectionError, OperationalError, TimeoutError
+from cognee.modules.sync.models import SyncOperation, SyncStatus
+from cognee.infrastructure.databases.relational import get_relational_engine
+from cognee.shared.logging_utils import get_logger
+from cognee.infrastructure.utils.calculate_backoff import calculate_backoff
+
+logger = get_logger("sync.db_operations")
+
+
+async def _retry_db_operation(operation_func, run_id: str, max_retries: int = 3):
+    """
+    Retry database operations with exponential backoff for transient failures.
+
+    Args:
+        operation_func: Async function to retry
+        run_id: Run ID for logging context
+        max_retries: Maximum number of retry attempts
+
+    Returns:
+        Result of the operation function
+
+    Raises:
+        Exception: Re-raises the last exception if all retries fail
+    """
+    attempt = 0
+    last_exception = None
+
+    while attempt < max_retries:
+        try:
+            return await operation_func()
+        except (DisconnectionError, OperationalError, TimeoutError) as e:
+            attempt += 1
+            last_exception = e
+
+            if attempt >= max_retries:
+                logger.error(
+                    f"Database operation failed after {max_retries} attempts for run_id {run_id}: {str(e)}"
+                )
+                break
+
+            backoff_time = calculate_backoff(attempt - 1)  # calculate_backoff is 0-indexed
+            logger.warning(
+                f"Database operation failed for run_id {run_id}, retrying in {backoff_time:.2f}s (attempt {attempt}/{max_retries}): {str(e)}"
+            )
+            await asyncio.sleep(backoff_time)
+
+        except Exception as e:
+            # Non-transient errors should not be retried
+            logger.error(f"Non-retryable database error for run_id {run_id}: {str(e)}")
+            raise
+
+    # If we get here, all retries failed
+    raise last_exception
+
+
+async def update_sync_operation(
+    run_id: str,
+    status: Optional[SyncStatus] = None,
+    progress_percentage: Optional[int] = None,
+    records_downloaded: Optional[int] = None,
+    total_records_to_sync: Optional[int] = None,
+    total_records_to_download: Optional[int] = None,
+    total_records_to_upload: Optional[int] = None,
+    records_uploaded: Optional[int] = None,
+    bytes_downloaded: Optional[int] = None,
+    bytes_uploaded: Optional[int] = None,
+    dataset_sync_hashes: Optional[dict] = None,
+    error_message: Optional[str] = None,
+    retry_count: Optional[int] = None,
+    started_at: Optional[datetime] = None,
+    completed_at: Optional[datetime] = None,
+) -> Optional[SyncOperation]:
+    """
+    Update a sync operation record with new status/progress information.
+
+    Args:
+        run_id: The public run_id of the sync operation to update
+        status: New status for the operation
+        progress_percentage: Progress percentage (0-100)
+        records_downloaded: Number of records downloaded so far
+        total_records_to_sync: Total number of records that need to be synced
+        total_records_to_download: Total number of records to download from cloud
+        total_records_to_upload: Total number of records to upload to cloud
+        records_uploaded: Number of records uploaded so far
+        bytes_downloaded: Total bytes downloaded from cloud
+        bytes_uploaded: Total bytes uploaded to cloud
+        dataset_sync_hashes: Dict mapping dataset_id -> {uploaded: [hashes], downloaded: [hashes]}
+        error_message: Error message if operation failed
+        retry_count: Number of retry attempts
+        started_at: When the actual processing started
+        completed_at: When the operation completed (success or failure)
+
+    Returns:
+        SyncOperation: The updated sync operation record, or None if not found
+    """
+
+    async def _perform_update():
+        db_engine = get_relational_engine()
+
+        async with db_engine.get_async_session() as session:
+            try:
+                # Find the sync operation
+                query = select(SyncOperation).where(SyncOperation.run_id == run_id)
+                result = await session.execute(query)
+                sync_operation = result.scalars().first()
+
+                if not sync_operation:
+                    logger.warning(f"Sync operation not found for run_id: {run_id}")
+                    return None
+
+                # Log what we're updating for debugging
+                updates = []
+                if status is not None:
+                    updates.append(f"status={status.value}")
+                if progress_percentage is not None:
+                    updates.append(f"progress={progress_percentage}%")
+                if records_downloaded is not None:
+                    updates.append(f"downloaded={records_downloaded}")
+                if records_uploaded is not None:
+                    updates.append(f"uploaded={records_uploaded}")
+                if total_records_to_sync is not None:
+                    updates.append(f"total_sync={total_records_to_sync}")
+                if total_records_to_download is not None:
+                    updates.append(f"total_download={total_records_to_download}")
+                if total_records_to_upload is not None:
+                    updates.append(f"total_upload={total_records_to_upload}")
+
+                if updates:
+                    logger.debug(f"Updating sync operation {run_id}: {', '.join(updates)}")
+
+                # Update fields that were provided
+                if status is not None:
+                    sync_operation.status = status
+
+                if progress_percentage is not None:
+                    sync_operation.progress_percentage = max(0, min(100, progress_percentage))
+
+                if records_downloaded is not None:
+                    sync_operation.records_downloaded = records_downloaded
+
+                if records_uploaded is not None:
+                    sync_operation.records_uploaded = records_uploaded
+
+                if total_records_to_sync is not None:
+                    sync_operation.total_records_to_sync = total_records_to_sync
+
+                if total_records_to_download is not None:
+                    sync_operation.total_records_to_download = total_records_to_download
+
+                if total_records_to_upload is not None:
+                    sync_operation.total_records_to_upload = total_records_to_upload
+
+                if bytes_downloaded is not None:
+                    sync_operation.bytes_downloaded = bytes_downloaded
+
+                if bytes_uploaded is not None:
+                    sync_operation.bytes_uploaded = bytes_uploaded
+
+                if dataset_sync_hashes is not None:
+                    sync_operation.dataset_sync_hashes = dataset_sync_hashes
+
+                if error_message is not None:
+                    sync_operation.error_message = error_message
+
+                if retry_count is not None:
+                    sync_operation.retry_count = retry_count
+
+                if started_at is not None:
+                    sync_operation.started_at = started_at
+
+                if completed_at is not None:
+                    sync_operation.completed_at = completed_at
+
+                # Auto-set completion timestamp for terminal statuses
+                if (
+                    status in [SyncStatus.COMPLETED, SyncStatus.FAILED, SyncStatus.CANCELLED]
+                    and completed_at is None
+                ):
+                    sync_operation.completed_at = datetime.now(timezone.utc)
+
+                # Auto-set started timestamp when moving to IN_PROGRESS
+                if status == SyncStatus.IN_PROGRESS and sync_operation.started_at is None:
+                    sync_operation.started_at = datetime.now(timezone.utc)
+
+                await session.commit()
+                await session.refresh(sync_operation)
+
+                logger.debug(f"Successfully updated sync operation {run_id}")
+                return sync_operation
+
+            except SQLAlchemyError as e:
+                logger.error(
+                    f"Database error updating sync operation {run_id}: {str(e)}", exc_info=True
+                )
+                await session.rollback()
+                raise
+            except Exception as e:
+                logger.error(
+                    f"Unexpected error updating sync operation {run_id}: {str(e)}", exc_info=True
+                )
+                await session.rollback()
+                raise
+
+    # Use retry logic for the database operation
+    return await _retry_db_operation(_perform_update, run_id)
+
+
+async def mark_sync_started(run_id: str) -> Optional[SyncOperation]:
+    """Convenience method to mark a sync operation as started."""
+    return await update_sync_operation(
+        run_id=run_id, status=SyncStatus.IN_PROGRESS, started_at=datetime.now(timezone.utc)
+    )
+
+
+async def mark_sync_completed(
+    run_id: str,
+    records_downloaded: int = 0,
+    records_uploaded: int = 0,
+    bytes_downloaded: int = 0,
+    bytes_uploaded: int = 0,
+    dataset_sync_hashes: Optional[dict] = None,
+) -> Optional[SyncOperation]:
+    """Convenience method to mark a sync operation as completed successfully."""
+    return await update_sync_operation(
+        run_id=run_id,
+        status=SyncStatus.COMPLETED,
+        progress_percentage=100,
+        records_downloaded=records_downloaded,
+        records_uploaded=records_uploaded,
+        bytes_downloaded=bytes_downloaded,
+        bytes_uploaded=bytes_uploaded,
+        dataset_sync_hashes=dataset_sync_hashes,
+        completed_at=datetime.now(timezone.utc),
+    )
+
+
+async def mark_sync_failed(run_id: str, error_message: str) -> Optional[SyncOperation]:
+    """Convenience method to mark a sync operation as failed."""
+    return await update_sync_operation(
+        run_id=run_id,
+        status=SyncStatus.FAILED,
+        error_message=error_message,
+        completed_at=datetime.now(timezone.utc),
+    )

cognee/modules/sync/models/SyncOperation.py

@@ -0,0 +1,142 @@
+from uuid import uuid4
+from enum import Enum
+from typing import Optional, List
+from datetime import datetime, timezone
+from sqlalchemy import (
+    Column,
+    Text,
+    DateTime,
+    UUID as SQLAlchemy_UUID,
+    Integer,
+    Enum as SQLEnum,
+    JSON,
+)
+
+from cognee.infrastructure.databases.relational import Base
+
+
+class SyncStatus(str, Enum):
+    """Enumeration of possible sync operation statuses."""
+
+    STARTED = "started"
+    IN_PROGRESS = "in_progress"
+    COMPLETED = "completed"
+    FAILED = "failed"
+    CANCELLED = "cancelled"
+
+
+class SyncOperation(Base):
+    """
+    Database model for tracking sync operations.
+
+    This model stores information about background sync operations,
+    allowing users to monitor progress and query the status of their sync requests.
+    """
+
+    __tablename__ = "sync_operations"
+
+    # Primary identifiers
+    id = Column(SQLAlchemy_UUID, primary_key=True, default=uuid4, doc="Database primary key")
+    run_id = Column(Text, unique=True, index=True, doc="Public run ID returned to users")
+
+    # Status and progress tracking
+    status = Column(
+        SQLEnum(SyncStatus), default=SyncStatus.STARTED, doc="Current status of the sync operation"
+    )
+    progress_percentage = Column(Integer, default=0, doc="Progress percentage (0-100)")
+
+    # Operation metadata
+    dataset_ids = Column(JSON, doc="Array of dataset IDs being synced")
+    dataset_names = Column(JSON, doc="Array of dataset names being synced")
+    user_id = Column(SQLAlchemy_UUID, index=True, doc="ID of the user who initiated the sync")
+
+    # Timing information
+    created_at = Column(
+        DateTime(timezone=True),
+        default=lambda: datetime.now(timezone.utc),
+        doc="When the sync was initiated",
+    )
+    started_at = Column(DateTime(timezone=True), doc="When the actual sync processing began")
+    completed_at = Column(
+        DateTime(timezone=True), doc="When the sync finished (success or failure)"
+    )
+
+    # Operation details
+    total_records_to_sync = Column(Integer, doc="Total number of records to sync")
+    total_records_to_download = Column(Integer, doc="Total number of records to download")
+    total_records_to_upload = Column(Integer, doc="Total number of records to upload")
+
+    records_downloaded = Column(Integer, default=0, doc="Number of records successfully downloaded")
+    records_uploaded = Column(Integer, default=0, doc="Number of records successfully uploaded")
+    bytes_downloaded = Column(Integer, default=0, doc="Total bytes downloaded from cloud")
+    bytes_uploaded = Column(Integer, default=0, doc="Total bytes uploaded to cloud")
+
+    # Data lineage tracking per dataset
+    dataset_sync_hashes = Column(
+        JSON, doc="Mapping of dataset_id -> {uploaded: [hashes], downloaded: [hashes]}"
+    )
+
+    # Error handling
+    error_message = Column(Text, doc="Error message if sync failed")
+    retry_count = Column(Integer, default=0, doc="Number of retry attempts")
+
+    def get_duration_seconds(self) -> Optional[float]:
+        """Get the duration of the sync operation in seconds."""
+        if not self.created_at:
+            return None
+
+        end_time = self.completed_at or datetime.now(timezone.utc)
+        return (end_time - self.created_at).total_seconds()
+
+    def get_progress_info(self) -> dict:
+        """Get comprehensive progress information."""
+        total_records_processed = (self.records_downloaded or 0) + (self.records_uploaded or 0)
+        total_bytes_transferred = (self.bytes_downloaded or 0) + (self.bytes_uploaded or 0)
+
+        return {
+            "status": self.status.value,
+            "progress_percentage": self.progress_percentage,
+            "records_processed": f"{total_records_processed}/{self.total_records_to_sync or 'unknown'}",
+            "records_downloaded": self.records_downloaded or 0,
+            "records_uploaded": self.records_uploaded or 0,
+            "bytes_transferred": total_bytes_transferred,
+            "bytes_downloaded": self.bytes_downloaded or 0,
+            "bytes_uploaded": self.bytes_uploaded or 0,
+            "duration_seconds": self.get_duration_seconds(),
+            "error_message": self.error_message,
+            "dataset_sync_hashes": self.dataset_sync_hashes or {},
+        }
+
+    def _get_all_sync_hashes(self) -> List[str]:
+        """Get all content hashes for data created/modified during this sync operation."""
+        all_hashes = set()
+        dataset_hashes = self.dataset_sync_hashes or {}
+
+        for dataset_id, operations in dataset_hashes.items():
+            if isinstance(operations, dict):
+                all_hashes.update(operations.get("uploaded", []))
+                all_hashes.update(operations.get("downloaded", []))
+
+        return list(all_hashes)
+
+    def _get_dataset_sync_hashes(self, dataset_id: str) -> dict:
+        """Get uploaded/downloaded hashes for a specific dataset."""
+        dataset_hashes = self.dataset_sync_hashes or {}
+        return dataset_hashes.get(dataset_id, {"uploaded": [], "downloaded": []})
+
+    def was_data_synced(self, content_hash: str, dataset_id: str = None) -> bool:
+        """
+        Check if a specific piece of data was part of this sync operation.
+
+        Args:
+            content_hash: The content hash to check for
+            dataset_id: Optional - check only within this dataset
+        """
+        if dataset_id:
+            dataset_hashes = self.get_dataset_sync_hashes(dataset_id)
+            return content_hash in dataset_hashes.get(
+                "uploaded", []
+            ) or content_hash in dataset_hashes.get("downloaded", [])
+
+        all_hashes = self.get_all_sync_hashes()
+        return content_hash in all_hashes

cognee/modules/sync/models/__init__.py

@@ -0,0 +1,3 @@
+from .SyncOperation import SyncOperation, SyncStatus
+
+__all__ = ["SyncOperation", "SyncStatus"]

cognee/modules/users/__init__.py

@@ -1,2 +1 @@
 from .get_user_db import get_user_db
-from .get_user_db import get_async_session

cognee/modules/users/methods/__init__.py

@@ -4,4 +4,7 @@ from .delete_user import delete_user
 from .get_default_user import get_default_user
 from .get_user_by_email import get_user_by_email
 from .create_default_user import create_default_user
-from .get_authenticated_user import get_authenticated_user
+from .get_authenticated_user import (
+    get_authenticated_user,
+    REQUIRE_AUTHENTICATION,
+)

cognee/modules/users/methods/create_user.py

@@ -1,6 +1,10 @@
+from uuid import uuid4
 from fastapi_users.exceptions import UserAlreadyExists
-from cognee.modules.users.exceptions import TenantNotFoundError
+
 from cognee.infrastructure.databases.relational import get_relational_engine
+from cognee.modules.notebooks.methods import create_notebook
+from cognee.modules.notebooks.models.Notebook import NotebookCell
+from cognee.modules.users.exceptions import TenantNotFoundError
 from cognee.modules.users.get_user_manager import get_user_manager_context
 from cognee.modules.users.get_user_db import get_user_db_context
 from cognee.modules.users.models.User import UserCreate
@@ -56,6 +60,27 @@ async def create_user(
                     if auto_login:
                         await session.refresh(user)
 
+                    await create_notebook(
+                        user_id=user.id,
+                        notebook_name="Welcome to cognee 🧠",
+                        cells=[
+                            NotebookCell(
+                                id=uuid4(),
+                                name="Welcome",
+                                content="Cognee is your toolkit for turning text into a structured knowledge graph, optionally enhanced by ontologies, and then querying it with advanced retrieval techniques. This notebook will guide you through a simple example.",
+                                type="markdown",
+                            ),
+                            NotebookCell(
+                                id=uuid4(),
+                                name="Example",
+                                content="",
+                                type="code",
+                            ),
+                        ],
+                        deletable=False,
+                        session=session,
+                    )
+
                     return user
     except UserAlreadyExists as error:
         print(f"User {email} already exists")

cognee/modules/users/methods/get_authenticated_user.py

@@ -1,48 +1,42 @@
+import os
+from typing import Optional
+from fastapi import Depends, HTTPException
+from ..models import User
 from ..get_fastapi_users import get_fastapi_users
+from .get_default_user import get_default_user
+from cognee.shared.logging_utils import get_logger
 
 
-fastapi_users = get_fastapi_users()
-
-get_authenticated_user = fastapi_users.current_user(active=True)
-
-# from types import SimpleNamespace
-
-# from ..get_fastapi_users import get_fastapi_users
-# from fastapi import HTTPException, Security
-# from fastapi.security import HTTPBearer, HTTPAuthorizationCredentials
-# import os
-# import jwt
-
-# from uuid import UUID
-
-# fastapi_users = get_fastapi_users()
+logger = get_logger("get_authenticated_user")
 
-# # Allows Swagger to understand authorization type and allow single sign on for the Swagger docs to test backend
-# bearer_scheme = HTTPBearer(scheme_name="BearerAuth", description="Paste **Bearer <JWT>**")
+# Check environment variable to determine authentication requirement
+REQUIRE_AUTHENTICATION = (
+    os.getenv("REQUIRE_AUTHENTICATION", "false").lower() == "true"
+    or os.getenv("ENABLE_BACKEND_ACCESS_CONTROL", "false").lower() == "true"
+)
 
+fastapi_users = get_fastapi_users()
 
-# async def get_authenticated_user(
-#     creds: HTTPAuthorizationCredentials = Security(bearer_scheme),
-# ) -> SimpleNamespace:
-#     """
-#     Extract and validate the JWT presented in the Authorization header.
-#     """
-#     if creds is None:  # header missing
-#         raise HTTPException(status_code=401, detail="Not authenticated")
-
-#     if creds.scheme.lower() != "bearer":  # shouldn't happen extra guard
-#         raise HTTPException(status_code=401, detail="Invalid authentication scheme")
-
-#     token = creds.credentials
-#     try:
-#         payload = jwt.decode(
-#             token, os.getenv("FASTAPI_USERS_JWT_SECRET", "super_secret"), algorithms=["HS256"]
-#         )
-
-#         auth_data = SimpleNamespace(id=UUID(payload["user_id"]))
-#         return auth_data
-
-#     except jwt.ExpiredSignatureError:
-#         raise HTTPException(status_code=401, detail="Token has expired")
-#     except jwt.InvalidTokenError:
-#         raise HTTPException(status_code=401, detail="Invalid token")
+_auth_dependency = fastapi_users.current_user(active=True, optional=not REQUIRE_AUTHENTICATION)
+
+
+async def get_authenticated_user(
+    user: Optional[User] = Depends(_auth_dependency),
+) -> User:
+    """
+    Get authenticated user with environment-controlled behavior:
+    - If REQUIRE_AUTHENTICATION=true: Enforces authentication (raises 401 if not authenticated)
+    - If REQUIRE_AUTHENTICATION=false: Falls back to default user if not authenticated
+
+    Always returns a User object for consistent typing.
+    """
+    if user is None:
+        # When authentication is optional and user is None, use default user
+        try:
+            user = await get_default_user()
+        except Exception as e:
+            # Convert any get_default_user failure into a proper HTTP 500 error
+            logger.error(f"Failed to create default user: {str(e)}")
+            raise HTTPException(status_code=500, detail=f"Failed to create default user: {str(e)}")
+
+    return user

cognee/modules/users/methods/get_default_user.py

@@ -29,7 +29,9 @@ async def get_default_user() -> SimpleNamespace:
 
             # We return a SimpleNamespace to have the same user type as our SaaS
             # SimpleNamespace is just a dictionary which can be accessed through attributes
-            auth_data = SimpleNamespace(id=user.id, tenant_id=user.tenant_id, roles=[])
+            auth_data = SimpleNamespace(
+                id=user.id, email=user.email, tenant_id=user.tenant_id, roles=[]
+            )
             return auth_data
     except Exception as error:
         if "principals" in str(error.args):

cognee/modules/users/permissions/methods/get_specific_user_permission_datasets.py

@@ -1,4 +1,5 @@
 from uuid import UUID
+from typing import Optional
 from cognee.modules.data.models.Dataset import Dataset
 from cognee.modules.users.permissions.methods.get_all_user_permission_datasets import (
     get_all_user_permission_datasets,
@@ -8,7 +9,7 @@ from cognee.modules.users.methods import get_user
 
 
 async def get_specific_user_permission_datasets(
-    user_id: UUID, permission_type: str, dataset_ids: list[UUID] = None
+    user_id: UUID, permission_type: str, dataset_ids: Optional[list[UUID]] = None
 ) -> list[Dataset]:
     """
         Return a list of datasets user has given permission for. If a list of datasets is provided,