cognee 0.5.1__py3-none-any.whl → 0.5.2__py3-none-any.whl

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/alembic/README

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

cognee/alembic/env.py

@@ -0,0 +1,107 @@
+import asyncio
+from alembic import context
+from logging.config import fileConfig
+from sqlalchemy import pool
+from sqlalchemy.engine import Connection
+from sqlalchemy.ext.asyncio import async_engine_from_config
+
+from cognee.infrastructure.databases.relational import get_relational_engine, Base
+
+# this is the Alembic Config object, which provides
+# access to the values within the .ini file in use.
+config = context.config
+
+# Interpret the config file for Python logging.
+# This line sets up loggers basically.
+if config.config_file_name is not None:
+    fileConfig(config.config_file_name)
+
+# add your model's MetaData object here
+# for 'autogenerate' support
+# from myapp import mymodel
+# target_metadata = mymodel.Base.metadata
+
+target_metadata = Base.metadata
+
+# other values from the config, defined by the needs of env.py,
+# can be acquired:
+# my_important_option = config.get_main_option("my_important_option")
+# ... etc.
+
+
+def run_migrations_offline() -> None:
+    """Run migrations in 'offline' mode.
+
+    This configures the context with just a URL
+    and not an Engine, though an Engine is acceptable
+    here as well.  By skipping the Engine creation
+    we don't even need a DBAPI to be available.
+
+    Calls to context.execute() here emit the given string to the
+    script output.
+
+    """
+    url = config.get_main_option("sqlalchemy.url")
+    context.configure(
+        url=url,
+        target_metadata=target_metadata,
+        literal_binds=True,
+        dialect_opts={"paramstyle": "named"},
+    )
+
+    with context.begin_transaction():
+        context.run_migrations()
+
+
+def do_run_migrations(connection: Connection) -> None:
+    context.configure(connection=connection, target_metadata=target_metadata)
+
+    with context.begin_transaction():
+        context.run_migrations()
+
+
+async def run_async_migrations() -> None:
+    """In this scenario we need to create an Engine
+    and associate a connection with the context.
+    """
+
+    connectable = async_engine_from_config(
+        config.get_section(config.config_ini_section, {}),
+        prefix="sqlalchemy.",
+        poolclass=pool.NullPool,
+    )
+
+    async with connectable.connect() as connection:
+        await connection.run_sync(do_run_migrations)
+
+    await connectable.dispose()
+
+
+def run_migrations_online() -> None:
+    """Run migrations in 'online' mode."""
+
+    asyncio.run(run_async_migrations())
+
+
+db_engine = get_relational_engine()
+
+print("Using database:", db_engine.db_uri)
+# In case db_uri is a SQLAlchemy URL object, convert it to string
+db_uri = (
+    db_engine.db_uri
+    if isinstance(db_engine.db_uri, str)
+    else db_engine.db_uri.render_as_string(hide_password=False)
+)
+
+config.set_section_option(
+    config.config_ini_section,
+    "SQLALCHEMY_DATABASE_URI",
+    db_uri,
+)
+
+
+if context.is_offline_mode():
+    print("OFFLINE MODE")
+    run_migrations_offline()
+else:
+    run_migrations_online()

cognee/alembic/script.py.mako

@@ -0,0 +1,26 @@
+"""${message}
+
+Revision ID: ${up_revision}
+Revises: ${down_revision | comma,n}
+Create Date: ${create_date}
+
+"""
+from typing import Sequence, Union
+
+from alembic import op
+import sqlalchemy as sa
+${imports if imports else ""}
+
+# revision identifiers, used by Alembic.
+revision: str = ${repr(up_revision)}
+down_revision: Union[str, None] = ${repr(down_revision)}
+branch_labels: Union[str, Sequence[str], None] = ${repr(branch_labels)}
+depends_on: Union[str, Sequence[str], None] = ${repr(depends_on)}
+
+
+def upgrade() -> None:
+    ${upgrades if upgrades else "pass"}
+
+
+def downgrade() -> None:
+    ${downgrades if downgrades else "pass"}

cognee/alembic/versions/1a58b986e6e1_enable_delete_for_old_tutorial_notebooks.py

@@ -0,0 +1,52 @@
+"""Enable delete for old tutorial notebooks
+
+Revision ID: 1a58b986e6e1
+Revises: 46a6ce2bd2b2
+Create Date: 2025-12-17 11:04:44.414259
+
+"""
+
+from typing import Sequence, Union
+
+from alembic import op
+import sqlalchemy as sa
+
+
+# revision identifiers, used by Alembic.
+revision: str = "1a58b986e6e1"
+down_revision: Union[str, None] = "e1ec1dcb50b6"
+branch_labels: Union[str, Sequence[str], None] = None
+depends_on: Union[str, Sequence[str], None] = None
+
+
+def change_tutorial_deletable_flag(deletable: bool) -> None:
+    bind = op.get_bind()
+    inspector = sa.inspect(bind)
+
+    if "notebooks" not in inspector.get_table_names():
+        return
+
+    columns = {col["name"] for col in inspector.get_columns("notebooks")}
+    required_columns = {"name", "deletable"}
+    if not required_columns.issubset(columns):
+        return
+
+    notebooks = sa.table(
+        "notebooks",
+        sa.Column("name", sa.String()),
+        sa.Column("deletable", sa.Boolean()),
+    )
+
+    tutorial_name = "Python Development with Cognee Tutorial 🧠"
+
+    bind.execute(
+        notebooks.update().where(notebooks.c.name == tutorial_name).values(deletable=deletable)
+    )
+
+
+def upgrade() -> None:
+    change_tutorial_deletable_flag(True)
+
+
+def downgrade() -> None:
+    change_tutorial_deletable_flag(False)

cognee/alembic/versions/1d0bb7fede17_add_pipeline_run_status.py

@@ -0,0 +1,33 @@
+"""Add pipeline run status
+
+Revision ID: 1d0bb7fede17
+Revises: 482cd6517ce4
+Create Date: 2025-05-19 10:58:15.993314
+"""
+
+from typing import Sequence, Union
+
+from alembic import op
+
+from cognee.infrastructure.databases.relational.get_relational_engine import get_relational_engine
+from cognee.modules.pipelines.models.PipelineRun import PipelineRun, PipelineRunStatus
+
+
+# revision identifiers, used by Alembic.
+revision: str = "1d0bb7fede17"
+down_revision: Union[str, None] = "482cd6517ce4"
+branch_labels: Union[str, Sequence[str], None] = None
+depends_on: Union[str, Sequence[str], None] = "482cd6517ce4"
+
+
+def upgrade() -> None:
+    db_engine = get_relational_engine()
+
+    if db_engine.engine.dialect.name == "postgresql":
+        op.execute(
+            "ALTER TYPE pipelinerunstatus ADD VALUE IF NOT EXISTS 'DATASET_PROCESSING_INITIATED'"
+        )
+
+
+def downgrade() -> None:
+    pass

cognee/alembic/versions/1daae0df1866_incremental_loading.py

@@ -0,0 +1,48 @@
+"""incremental_loading
+
+Revision ID: 1daae0df1866
+Revises: b9274c27a25a
+Create Date: 2025-08-12 13:14:12.515935
+
+"""
+
+from typing import Sequence, Union
+
+from alembic import op
+import sqlalchemy as sa
+from sqlalchemy.ext.mutable import MutableDict
+
+# revision identifiers, used by Alembic.
+revision: str = "1daae0df1866"
+down_revision: Union[str, None] = "b9274c27a25a"
+branch_labels: Union[str, Sequence[str], None] = None
+depends_on: Union[str, Sequence[str], None] = None
+
+
+def _get_column(inspector, table, name, schema=None):
+    for col in inspector.get_columns(table, schema=schema):
+        if col["name"] == name:
+            return col
+    return None
+
+
+def upgrade() -> None:
+    conn = op.get_bind()
+    insp = sa.inspect(conn)
+
+    # If column already exists skip migration
+    pipeline_status_column = _get_column(insp, "data", "pipeline_status")
+    if not pipeline_status_column:
+        op.add_column(
+            "data",
+            sa.Column(
+                "pipeline_status",
+                MutableDict.as_mutable(sa.JSON),
+                nullable=False,
+                server_default=sa.text("'{}'"),
+            ),
+        )
+
+
+def downgrade() -> None:
+    op.drop_column("data", "pipeline_status")

cognee/alembic/versions/211ab850ef3d_add_sync_operations_table.py

@@ -0,0 +1,118 @@
+"""Add sync_operations table
+
+Revision ID: 211ab850ef3d
+Revises: 9e7a3cb85175
+Create Date: 2025-09-10 20:11:13.534829
+
+"""
+
+from typing import Sequence, Union
+
+from alembic import op
+import sqlalchemy as sa
+from sqlalchemy.dialects import postgresql
+
+
+# revision identifiers, used by Alembic.
+revision: str = "211ab850ef3d"
+down_revision: Union[str, None] = "45957f0a9849"
+branch_labels: Union[str, Sequence[str], None] = None
+depends_on: Union[str, Sequence[str], None] = None
+
+
+def upgrade() -> None:
+    # ### commands auto generated by Alembic - please adjust! ###
+
+    # Check if table already exists (it might be created by Base.metadata.create_all() in initial migration)
+    connection = op.get_bind()
+    inspector = sa.inspect(connection)
+
+    if op.get_context().dialect.name == "postgresql":
+        syncstatus_enum = postgresql.ENUM(
+            "STARTED", "IN_PROGRESS", "COMPLETED", "FAILED", "CANCELLED", name="syncstatus"
+        )
+        syncstatus_enum.create(op.get_bind(), checkfirst=True)
+
+    if "sync_operations" not in inspector.get_table_names():
+        if op.get_context().dialect.name == "postgresql":
+            syncstatus = postgresql.ENUM(
+                "STARTED",
+                "IN_PROGRESS",
+                "COMPLETED",
+                "FAILED",
+                "CANCELLED",
+                name="syncstatus",
+                create_type=False,
+            )
+        else:
+            syncstatus = sa.Enum(
+                "STARTED",
+                "IN_PROGRESS",
+                "COMPLETED",
+                "FAILED",
+                "CANCELLED",
+                name="syncstatus",
+                create_type=False,
+            )
+
+        # Table doesn't exist, create it normally
+        op.create_table(
+            "sync_operations",
+            sa.Column("id", sa.UUID(), nullable=False),
+            sa.Column("run_id", sa.Text(), nullable=True),
+            sa.Column(
+                "status",
+                syncstatus,
+                nullable=True,
+            ),
+            sa.Column("progress_percentage", sa.Integer(), nullable=True),
+            sa.Column("dataset_ids", sa.JSON(), nullable=True),
+            sa.Column("dataset_names", sa.JSON(), nullable=True),
+            sa.Column("user_id", sa.UUID(), nullable=True),
+            sa.Column("created_at", sa.DateTime(timezone=True), nullable=True),
+            sa.Column("started_at", sa.DateTime(timezone=True), nullable=True),
+            sa.Column("completed_at", sa.DateTime(timezone=True), nullable=True),
+            sa.Column("total_records_to_sync", sa.Integer(), nullable=True),
+            sa.Column("total_records_to_download", sa.Integer(), nullable=True),
+            sa.Column("total_records_to_upload", sa.Integer(), nullable=True),
+            sa.Column("records_downloaded", sa.Integer(), nullable=True),
+            sa.Column("records_uploaded", sa.Integer(), nullable=True),
+            sa.Column("bytes_downloaded", sa.Integer(), nullable=True),
+            sa.Column("bytes_uploaded", sa.Integer(), nullable=True),
+            sa.Column("dataset_sync_hashes", sa.JSON(), nullable=True),
+            sa.Column("error_message", sa.Text(), nullable=True),
+            sa.Column("retry_count", sa.Integer(), nullable=True),
+            sa.PrimaryKeyConstraint("id"),
+        )
+        op.create_index(
+            op.f("ix_sync_operations_run_id"), "sync_operations", ["run_id"], unique=True
+        )
+        op.create_index(
+            op.f("ix_sync_operations_user_id"), "sync_operations", ["user_id"], unique=False
+        )
+    else:
+        # Table already exists, but we might need to add missing columns or indexes
+        # For now, just log that the table already exists
+        print("sync_operations table already exists, skipping creation")
+
+    # ### end Alembic commands ###
+
+
+def downgrade() -> None:
+    # ### commands auto generated by Alembic - please adjust! ###
+
+    # Only drop if table exists (might have been created by Base.metadata.create_all())
+    connection = op.get_bind()
+    inspector = sa.inspect(connection)
+
+    if "sync_operations" in inspector.get_table_names():
+        op.drop_index(op.f("ix_sync_operations_user_id"), table_name="sync_operations")
+        op.drop_index(op.f("ix_sync_operations_run_id"), table_name="sync_operations")
+        op.drop_table("sync_operations")
+
+        # Drop the enum type that was created (only if no other tables are using it)
+        sa.Enum(name="syncstatus").drop(op.get_bind(), checkfirst=True)
+    else:
+        print("sync_operations table doesn't exist, skipping downgrade")
+
+    # ### end Alembic commands ###

cognee/alembic/versions/45957f0a9849_add_notebook_table.py

@@ -0,0 +1,46 @@
+"""Add notebook table
+
+Revision ID: 45957f0a9849
+Revises: 9e7a3cb85175
+Create Date: 2025-09-10 17:47:58.201319
+
+"""
+
+from datetime import datetime, timezone
+from uuid import uuid4
+from typing import Sequence, Union
+
+from alembic import op
+import sqlalchemy as sa
+
+
+# revision identifiers, used by Alembic.
+revision: str = "45957f0a9849"
+down_revision: Union[str, None] = "9e7a3cb85175"
+branch_labels: Union[str, Sequence[str], None] = None
+depends_on: Union[str, Sequence[str], None] = None
+
+
+def upgrade() -> None:
+    bind = op.get_bind()
+    inspector = sa.inspect(bind)
+
+    if "notebooks" not in inspector.get_table_names():
+        # Define table with all necessary columns including primary key
+        op.create_table(
+            "notebooks",
+            sa.Column("id", sa.UUID, primary_key=True, default=uuid4),  # Critical for SQLite
+            sa.Column("owner_id", sa.UUID, index=True),
+            sa.Column("name", sa.String(), nullable=False),
+            sa.Column("cells", sa.JSON(), nullable=False),
+            sa.Column("deletable", sa.Boolean(), default=True),
+            sa.Column("created_at", sa.DateTime(), default=lambda: datetime.now(timezone.utc)),
+        )
+
+
+def downgrade() -> None:
+    bind = op.get_bind()
+    inspector = sa.inspect(bind)
+
+    if "notebooks" in inspector.get_table_names():
+        op.drop_table("notebooks")