refactor: Package Cognee migrations with library
This commit is contained in:
parent
127d9860df
commit
8815c5ee0d
23 changed files with 67 additions and 6657 deletions
|
|
@ -1 +0,0 @@
|
|||
Generic single-database configuration with an async dbapi.
|
||||
|
|
@ -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",
|
||||
|
|
|
|||
|
|
@ -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
6635
cognee-mcp/uv.lock
generated
File diff suppressed because it is too large
Load diff
|
|
@ -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
1
cognee/alembic/README
Normal file
|
|
@ -0,0 +1 @@
|
|||
Generic single-database configuration with an async dbapi.
|
||||
|
|
@ -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
46
cognee/run_migrations.py
Normal 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)
|
||||
Loading…
Add table
Reference in a new issue