refactor: Package Cognee migrations with library

This commit is contained in:
Igor Ilic 2025-12-15 16:23:32 +01:00
parent 127d9860df
commit 8815c5ee0d
23 changed files with 67 additions and 6657 deletions

View file

@ -1 +0,0 @@
Generic single-database configuration with an async dbapi.

View file

@ -7,9 +7,9 @@ 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",
"cognee[postgres,codegraph,gemini,huggingface,docs,neo4j] @ file:/Users/igorilic/Desktop/cognee/dist/cognee-0.5.0.dev0-py3-none-any.whl",
# TODO: Remove gemini from optional dependecnies for new Cognee version after 0.3.4
"cognee[postgres,docs,neo4j]==0.3.7",
#"cognee[postgres,docs,neo4j]==0.5.0",
"fastmcp>=2.10.0,<3.0.0",
"mcp>=1.12.0,<2.0.0",
"uv>=0.6.3,<1.0.0",

View file

@ -858,26 +858,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:

6635
cognee-mcp/uv.lock generated

File diff suppressed because it is too large Load diff

View file

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

1
cognee/alembic/README Normal file
View file

@ -0,0 +1 @@
Generic single-database configuration with an async dbapi.

View file

@ -79,8 +79,20 @@ async def run_async_migrations() -> None:
def run_migrations_online() -> None:
"""Run migrations in 'online' mode."""
try:
# 1. Check if an event loop is already running (e.g., in your application's main thread).
loop = asyncio.get_running_loop()
asyncio.run(run_async_migrations())
# 2. If a loop is running, use it to synchronously run the async task.
# This prevents the nested loop error.
print("Using existing event loop to run migrations.")
loop.run_until_complete(run_async_migrations())
except RuntimeError:
# 3. If no loop is running (e.g., when calling 'alembic upgrade head' directly),
# start a new one with asyncio.run().
print("Starting new event loop for migrations.")
asyncio.run(run_async_migrations())
db_engine = get_relational_engine()

46
cognee/run_migrations.py Normal file
View file

@ -0,0 +1,46 @@
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/migrations'
# 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(
["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)