refactor: Fix mcp migrations (#1897)

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

## Description
- Include alembic migrations with Cognee library
- Use alembic migrations from Cognee library in Cognee MCP

## Type of Change
<!-- Please check the relevant option -->
- [ ] Bug fix (non-breaking change that fixes an issue)
- [ ] New feature (non-breaking change that adds functionality)
- [ ] Breaking change (fix or feature that would cause existing
functionality to change)
- [ ] Documentation update
- [ ] Code refactoring
- [ ] Performance improvement
- [ ] Other (please specify):

## Screenshots/Videos (if applicable)
<!-- Add screenshots or videos to help explain your changes -->

## Pre-submission Checklist
<!-- Please check all boxes that apply before submitting your PR -->
- [ ] **I have tested my changes thoroughly before submitting this PR**
- [ ] **This PR contains minimal changes necessary to address the
issue/feature**
- [ ] My code follows the project's coding standards and style
guidelines
- [ ] I have added tests that prove my fix is effective or that my
feature works
- [ ] I have added necessary documentation (if applicable)
- [ ] All new and existing tests pass
- [ ] I have searched existing PRs to ensure this change hasn't been
submitted already
- [ ] I have linked any relevant issues in the description
- [ ] My commits have clear and descriptive messages

## 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 is an auto-generated comment: release notes by coderabbit.ai
-->
## Summary by CodeRabbit

* **Bug Fixes**
* More robust startup: migration failures now trigger automatic database
initialization fallback, improving recovery and reducing failed starts.
* Improved startup sequencing with environment-aware process startup and
optional debug wait to aid local development.

* **Chores**
* Reorganized where migration configuration is located inside the
image/build.
  * Removed outdated dependency/configuration comments.

<sub>✏️ Tip: You can customize this high-level summary in your review
settings.</sub>
<!-- end of auto-generated comment: release notes by coderabbit.ai -->

Dockerfile

@@ -34,10 +34,6 @@ COPY README.md pyproject.toml uv.lock entrypoint.sh ./
 RUN --mount=type=cache,target=/root/.cache/uv \
     uv sync --extra debug --extra api --extra postgres --extra neo4j --extra llama-index --extra ollama --extra mistral --extra groq --extra anthropic --extra chromadb --frozen --no-install-project --no-dev --no-editable
 
-# Copy Alembic configuration
-COPY alembic.ini /app/alembic.ini
-COPY alembic/ /app/alembic
-
 # Then, add the rest of the project source code and install it
 # Installing separately from its dependencies allows optimal layer caching
 COPY ./cognee /app/cognee

cognee-mcp/Dockerfile

@@ -34,8 +34,8 @@ RUN --mount=type=cache,target=/root/.cache/uv \
     uv sync --frozen --no-install-project --no-dev --no-editable
 
 # Copy Alembic configuration
-COPY alembic.ini /app/alembic.ini
-COPY alembic/ /app/alembic
+COPY cognee/alembic.ini /app/cognee/alembic.ini
+COPY cognee/alembic/ /app/cognee/alembic
 
 # Then, add the rest of the project source code and install it
 # Installing separately from its dependencies allows optimal layer caching

cognee-mcp/entrypoint.sh

@@ -56,24 +56,22 @@ if [ -n "$API_URL" ]; then
     echo "Skipping database migrations (API server handles its own database)"
 else
     echo "Direct mode: Using local cognee instance"
-    # Run Alembic migrations with proper error handling.
-    # Note on UserAlreadyExists error handling:
-    # During database migrations, we attempt to create a default user. If this user
-    # already exists (e.g., from a previous deployment or migration), it's not a
-    # critical error and shouldn't prevent the application from starting. This is
-    # different from other migration errors which could indicate database schema
-    # inconsistencies and should cause the startup to fail. This check allows for
-    # smooth redeployments and container restarts while maintaining data integrity.
     echo "Running database migrations..."
 
-    MIGRATION_OUTPUT=$(alembic upgrade head)
+    set +e # Disable exit on error to handle specific migration errors
+    MIGRATION_OUTPUT=$(cd cognee && alembic upgrade head)
     MIGRATION_EXIT_CODE=$?
+    set -e
 
     if [[ $MIGRATION_EXIT_CODE -ne 0 ]]; then
-        if [[ "$MIGRATION_OUTPUT" == *"UserAlreadyExists"* ]] || [[ "$MIGRATION_OUTPUT" == *"User default_user@example.com already exists"* ]]; then
-            echo "Warning: Default user already exists, continuing startup..."
-        else
-            echo "Migration failed with unexpected error."
+
+        echo "Migration failed with unexpected error. Trying to run Cognee without migrations."
+        echo "Initializing database tables..."
+        python /app/src/run_cognee_database_setup.py
+        INIT_EXIT_CODE=$?
+
+        if [[ $INIT_EXIT_CODE -ne 0 ]]; then
+            echo "Database initialization failed!"
             exit 1
         fi
     fi

cognee-mcp/pyproject.toml

@@ -8,7 +8,6 @@ requires-python = ">=3.10"
 dependencies = [
     # For local cognee repo usage remove comment bellow and add absolute path to cognee. Then run `uv sync --reinstall` in the mcp folder on local cognee changes.
     #"cognee[postgres,codegraph,gemini,huggingface,docs,neo4j] @ file:/Users/igorilic/Desktop/cognee",
-    # TODO: Remove gemini from optional dependecnies for new Cognee version after 0.3.4
     "cognee[postgres,docs,neo4j]==0.5.0",
     "fastmcp>=2.10.0,<3.0.0",
     "mcp>=1.12.0,<2.0.0",

cognee-mcp/src/run_cognee_database_setup.py

@@ -0,0 +1,5 @@
+from cognee.modules.engine.operations.setup import setup
+import asyncio
+
+if __name__ == "__main__":
+    asyncio.run(setup())

cognee-mcp/src/server.py

@@ -892,26 +892,11 @@ async def main():
 
         await setup()
 
-        # Run Alembic migrations from the main cognee directory where alembic.ini is located
+        # Run Cognee migrations
         logger.info("Running database migrations...")
-        migration_result = subprocess.run(
-            ["python", "-m", "alembic", "upgrade", "head"],
-            capture_output=True,
-            text=True,
-            cwd=Path(__file__).resolve().parent.parent.parent,
-        )
+        from cognee.run_migrations import run_migrations
 
-        if migration_result.returncode != 0:
-            migration_output = migration_result.stderr + migration_result.stdout
-            # Check for the expected UserAlreadyExists error (which is not critical)
-            if (
-                "UserAlreadyExists" in migration_output
-                or "User default_user@example.com already exists" in migration_output
-            ):
-                logger.warning("Warning: Default user already exists, continuing startup...")
-            else:
-                logger.error(f"Migration failed with unexpected error: {migration_output}")
-                sys.exit(1)
+        await run_migrations()
 
         logger.info("Database migrations done.")
     elif args.api_url:

cognee/__init__.py

@@ -33,3 +33,5 @@ from .api.v1.ui import start_ui
 
 # Pipelines
 from .modules import pipelines
+
+from cognee.run_migrations import run_migrations

cognee/run_migrations.py

@@ -0,0 +1,48 @@
+import os
+import sys
+import subprocess
+from pathlib import Path
+import importlib.resources as pkg_resources
+
+# Assuming your package is named 'cognee' and the migrations are under 'cognee/alembic'
+# This is a placeholder for the path logic.
+MIGRATIONS_PACKAGE = "cognee"
+MIGRATIONS_DIR_NAME = "alembic"
+
+
+async def run_migrations():
+    """
+    Finds the Alembic configuration within the installed package and
+    programmatically executes 'alembic upgrade head'.
+    """
+    # 1. Locate the base path of the installed package.
+    # This reliably finds the root directory of the installed 'cognee' package.
+    # We look for the parent of the 'migrations' directory.
+    package_root = str(pkg_resources.files(MIGRATIONS_PACKAGE))
+
+    # 2. Define the paths for config and scripts
+    alembic_ini_path = os.path.join(package_root, "alembic.ini")
+    script_location_path = os.path.join(package_root, MIGRATIONS_DIR_NAME)
+
+    if not os.path.exists(alembic_ini_path):
+        raise FileNotFoundError(
+            f"Error: alembic.ini not found at expected locations for package '{MIGRATIONS_PACKAGE}'."
+        )
+    if not os.path.exists(script_location_path):
+        raise FileNotFoundError(
+            f"Error: Migrations directory not found at expected locations for package '{MIGRATIONS_PACKAGE}'."
+        )
+
+    migration_result = subprocess.run(
+        ["python", "-m", "alembic", "upgrade", "head"],
+        capture_output=True,
+        text=True,
+        cwd=Path(package_root),
+    )
+
+    if migration_result.returncode != 0:
+        migration_output = migration_result.stderr + migration_result.stdout
+        print(f"Migration failed with unexpected error: {migration_output}")
+        sys.exit(1)
+
+    print("Migration completed successfully.")

entrypoint.sh

@@ -11,34 +11,24 @@ echo "Debug port: $DEBUG_PORT"
 echo "HTTP port: $HTTP_PORT"
 
 # Run Alembic migrations with proper error handling.
-# Note on UserAlreadyExists error handling:
-# During database migrations, we attempt to create a default user. If this user
-# already exists (e.g., from a previous deployment or migration), it's not a
-# critical error and shouldn't prevent the application from starting. This is
-# different from other migration errors which could indicate database schema
-# inconsistencies and should cause the startup to fail. This check allows for
-# smooth redeployments and container restarts while maintaining data integrity.
 echo "Running database migrations..."
 
+# Move to the cognee directory to run alembic migrations from there
 set +e # Disable exit on error to handle specific migration errors
-MIGRATION_OUTPUT=$(alembic upgrade head)
+MIGRATION_OUTPUT=$(cd cognee && alembic upgrade head)
 MIGRATION_EXIT_CODE=$?
 set -e
 
 if [[ $MIGRATION_EXIT_CODE -ne 0 ]]; then
-    if [[ "$MIGRATION_OUTPUT" == *"UserAlreadyExists"* ]] || [[ "$MIGRATION_OUTPUT" == *"User default_user@example.com already exists"* ]]; then
-        echo "Warning: Default user already exists, continuing startup..."
-    else
-        echo "Migration failed with unexpected error. Trying to run Cognee without migrations."
+    echo "Migration failed with unexpected error. Trying to run Cognee without migrations."
 
-        echo "Initializing database tables..."
-        python /app/cognee/modules/engine/operations/setup.py
-        INIT_EXIT_CODE=$?
+    echo "Initializing database tables..."
+    python /app/cognee/modules/engine/operations/setup.py
+    INIT_EXIT_CODE=$?
 
-        if [[ $INIT_EXIT_CODE -ne 0 ]]; then
-            echo "Database initialization failed!"
-            exit 1
-        fi
+    if [[ $INIT_EXIT_CODE -ne 0 ]]; then
+        echo "Database initialization failed!"
+        exit 1
     fi
 else
     echo "Database migrations done."