Add migration script for user limits overrides

js/sdk/package.json

@@ -1,6 +1,6 @@
 {
   "name": "r2r-js",
-  "version": "0.4.7",
+  "version": "0.4.8",
   "description": "",
   "main": "dist/index.js",
   "browser": "dist/index.browser.js",

py/cli/commands/system.py

@@ -149,6 +149,12 @@ async def serve(
 
     click.echo(f"Running on {host}:{port}, with docker={docker}")
 
+    # TODO: Remove after the next couple of releases
+    click.secho(
+        "Warning: if you are migrating from R2R version 3.3.18 or earlier, you must run `r2r db upgrade` before starting the server.",
+        fg="red",
+    )
+
     if full:
         click.echo(
             "Running the full R2R setup which includes `Hatchet` and `Unstructured.io`."

py/compose.yaml

@@ -37,7 +37,8 @@ services:
       -c max_connections=${R2R_POSTGRES_MAX_CONNECTIONS:-1024}
 
   r2r:
-    image: ${R2R_IMAGE:-ragtoriches/prod:latest}
+    # image: ${R2R_IMAGE:-ragtoriches/prod:latest}
+    image: r2r/test
     build:
       context: .
       args:

py/migrations/alembic.ini

@@ -0,0 +1,37 @@
+[alembic]
+script_location = .
+sqlalchemy.url = postgresql://postgres:postgres@localhost:5432/postgres
+
+[loggers]
+keys = root,sqlalchemy,alembic
+
+[handlers]
+keys = console
+
+[formatters]
+keys = generic
+
+[logger_root]
+level = WARN
+handlers = console
+qualname =
+
+[logger_sqlalchemy]
+level = WARN
+handlers =
+qualname = sqlalchemy.engine
+
+[logger_alembic]
+level = INFO
+handlers =
+qualname = alembic
+
+[handler_console]
+class = StreamHandler
+args = (sys.stderr,)
+level = NOTSET
+formatter = generic
+
+[formatter_generic]
+format = %(levelname)-5.5s [%(name)s] %(message)s
+datefmt = %H:%M:%S

py/migrations/versions/2fac23e4d91b_migrate_to_document_search.py

@@ -188,12 +188,19 @@ def generate_all_summaries():
 
 
 def check_if_upgrade_needed():
-    """Check if the upgrade has already been applied by checking for summary column"""
+    """Check if the upgrade has already been applied or is needed"""
     # Get database connection
     connection = op.get_bind()
     inspector = inspect(connection)
 
-    # Check if the columns exist
+    # First check if the document_info table exists
+    if not inspector.has_table("document_info", schema=project_name):
+        print(
+            f"Migration not needed: '{project_name}.document_info' table doesn't exist yet"
+        )
+        return False
+
+    # Then check if the columns exist
     existing_columns = [
         col["name"]
         for col in inspector.get_columns(f"document_info", schema=project_name)

py/migrations/versions/7eb70560f406_add_limits_overrides_to_users.py

@@ -0,0 +1,66 @@
+"""add_limits_overrides_to_users
+
+Revision ID: 7eb70560f406
+Revises: c45a9cf6a8a4
+Create Date: 2025-01-03 20:27:16.139511
+
+"""
+
+import os
+from typing import Sequence, Union
+
+import sqlalchemy as sa
+from alembic import op
+from sqlalchemy import inspect
+
+# revision identifiers, used by Alembic.
+revision: str = "7eb70560f406"
+down_revision: Union[str, None] = "c45a9cf6a8a4"
+branch_labels: Union[str, Sequence[str], None] = None
+depends_on: Union[str, Sequence[str], None] = None
+
+project_name = os.getenv("R2R_PROJECT_NAME", "r2r_default")
+
+
+def check_if_upgrade_needed():
+    """Check if the upgrade has already been applied"""
+    connection = op.get_bind()
+    inspector = inspect(connection)
+
+    # Check if users table exists
+    if not inspector.has_table("users", schema=project_name):
+        print(
+            f"Migration not needed: '{project_name}.users' table doesn't exist"
+        )
+        return False
+
+    users_columns = {
+        col["name"]
+        for col in inspector.get_columns("users", schema=project_name)
+    }
+
+    if "limits_overrides" in users_columns:
+        print(
+            "Migration not needed: users table already has limits_overrides column"
+        )
+        return False
+    else:
+        print("Migration needed: users table needs limits_overrides column")
+        return True
+
+
+def upgrade() -> None:
+    if not check_if_upgrade_needed():
+        return
+
+    # Add the limits_overrides column as JSONB with default NULL
+    op.add_column(
+        "users",
+        sa.Column("limits_overrides", sa.JSON(), nullable=True),
+        schema=project_name,
+    )
+
+
+def downgrade() -> None:
+    # Remove the limits_overrides column
+    op.drop_column("users", "limits_overrides", schema=project_name)

py/migrations/versions/8077140e1e99_v3_api_database_revision.py

@@ -11,6 +11,7 @@
 
 import sqlalchemy as sa
 from alembic import op
+from sqlalchemy import inspect
 
 # revision identifiers, used by Alembic.
 revision: str = "8077140e1e99"
@@ -24,16 +25,39 @@
         "Environment variable `R2R_PROJECT_NAME` must be provided migrate, it should be set equal to the value of `project_name` in your `r2r.toml`."
     )
 
-if (
-    input(
-        "WARNING: This migration will delete all graph data. Are you sure you want to continue? (yes/no) "
-    ).lower()
-    != "yes"
-):
-    raise ValueError("Migration aborted.")
+
+def check_if_upgrade_needed():
+    """Check if the upgrade has already been applied or is needed"""
+    connection = op.get_bind()
+    inspector = inspect(connection)
+
+    # Check collections table column names
+    collections_columns = {
+        col["name"]
+        for col in inspector.get_columns("collections", schema=project_name)
+    }
+
+    # If we find a new column name, we don't need to migrate
+    # If we find an old column name, we do need to migrate
+    if "id" in collections_columns:
+        print(
+            "Migration not needed: collections table already has 'id' column"
+        )
+        return False
+    elif "collection_id" in collections_columns:
+        print("Migration needed: collections table has old column names")
+        return True
+    else:
+        print(
+            "Migration not needed: collections table doesn't exist or has different structure"
+        )
+        return False
 
 
 def upgrade() -> None:
+    if not check_if_upgrade_needed():
+        return
+
     # Collections table migration
     op.alter_column(
         "collections",

py/migrations/versions/c45a9cf6a8a4_add_user_and_document_count_to_.py

@@ -1,7 +1,7 @@
 """Add user and document count to collection
 
 Revision ID: c45a9cf6a8a4
-Revises:
+Revises: 8077140e1e99
 Create Date: 2024-12-10 13:28:07.798167
 
 """
@@ -11,10 +11,11 @@
 
 import sqlalchemy as sa
 from alembic import op
+from sqlalchemy import inspect
 
 # revision identifiers, used by Alembic.
 revision: str = "c45a9cf6a8a4"
-down_revision: Union[str, None] = None
+down_revision: Union[str, None] = "8077140e1e99"
 branch_labels: Union[str, Sequence[str], None] = None
 depends_on: Union[str, Sequence[str], None] = None
 
@@ -25,7 +26,30 @@
     )
 
 
+def check_if_upgrade_needed():
+    """Check if the upgrade has already been applied"""
+    connection = op.get_bind()
+    inspector = inspect(connection)
+
+    collections_columns = {
+        col["name"]
+        for col in inspector.get_columns("collections", schema=project_name)
+    }
+
+    if "user_count" in collections_columns:
+        print(
+            "Migration not needed: collections table already has count columns"
+        )
+        return False
+    else:
+        print("Migration needed: collections table needs count columns")
+        return True
+
+
 def upgrade():
+    if not check_if_upgrade_needed():
+        return
+
     # Add the new columns with default value of 0
     op.add_column(
         "collections",

py/migrations/versions/d342e632358a_migrate_to_asyncpg.py

@@ -34,23 +34,34 @@ def get_col_spec(self, **kw):
 
 
 def check_if_upgrade_needed():
-    """Check if the upgrade has already been applied"""
-    # Get database connection
+    """Check if the upgrade has already been applied or is needed"""
     connection = op.get_bind()
     inspector = inspect(connection)
 
-    # Check if the new vectors table exists
+    # First check if the old table exists - if it doesn't, we don't need this migration
+    has_old_table = inspector.has_table(
+        old_vector_table_name, schema=project_name
+    )
+    if not has_old_table:
+        print(
+            f"Migration not needed: Original '{old_vector_table_name}' table doesn't exist"
+        )
+        # Skip this migration since we're starting from a newer state
+        return False
+
+    # Only if the old table exists, check if we need to migrate it
     has_new_table = inspector.has_table(
         new_vector_table_name, schema=project_name
     )
-
     if has_new_table:
         print(
             f"Migration not needed: '{new_vector_table_name}' table already exists"
         )
         return False
 
-    print(f"Migration needed: '{new_vector_table_name}' table does not exist")
+    print(
+        f"Migration needed: Need to migrate from '{old_vector_table_name}' to '{new_vector_table_name}'"
+    )
     return True
 
 

py/pyproject.toml

@@ -5,7 +5,7 @@ build-backend = "poetry.core.masonry.api"
 [tool.poetry]
 name = "r2r"
 readme = "README.md"
-version = "3.3.18"
+version = "3.3.19"
 
 description = "SciPhi R2R"
 authors = ["Owen Colegrove <owen@sciphi.ai>"]