Initial commit of migration qdrant

SubhrajitSW · SubhrajitSW · commit 024e8076ea33 · 2024-07-05T16:24:44.000+05:30
diff --git a/README.md b/README.md
@@ -0,0 +1,44 @@
+Qdrant Config Migration Script-
+
+1. In the folder create a file in the following format qdrant*migrator/migrations/<index>*<message>.py . Example qdrant_migrator/migrations/001_initial_migration.py . Now for the next version create 002
+2. Now these file executions are version controlled and you can see the versions in GUI ->
+   a. Go to collections there you will find migrations with a single point, with a metadata version. As of now it is in version 1 both dev and prod.
+   b. Once you run the next version 002, it will change.
+   c. You can't rerun versions forward. You have to create a new file with a new script to run next version.
+3. There is a rollback feature also, (backward)-
+   a. Your task will be to update the file with the backward function, for example, if you are writing something in v2, like create collection, hence it's rollback will be delete collection. Every file should consist of rollbacks also.
+   b. There can be some non rollback files also, like I was not able to find methods to remove payload*indexes in v1. Hence my rollback function is empty. You can anytime come in and re-write the rollback method in future before rolling back.
+   c. Ref code will be written below for file structure of qdrant_migrator/migrations/<index>*<message>.py
+4. How to run the code? There are 2 clis -
+   a. python cli.py migrate --url http://your-qdrant-url --api-key your-api-key --migration-folder migrations
+   b. python cli.py rollback --url http://your-qdrant-url --api-key your-api-key --migration-folder migrations --target-version <target_version>
+
+Summing it up, steps to go forward ->
+
+1. Create a version file (ref pt.1).
+2. Write the code for migration according to 3.c.
+3. Execute 4.a.
+4. Your Migration will be successfully completed.
+
+For Backward/Rollback ->
+
+1. Check its backward() function is written as expected.
+2. Execute 4.b.
+3. Your Rollback will be successfully completed.
+4. You can track the current version from migration collections
+
+Eg- code for point 3.c.-
+
+```python
+#Ensure each migration file has both `forward` and `backward` functions
+def forward(client):
+    # Your migration logic here
+    client.create_collection('new_collection')
+def backward(client):
+    # Logic to undo the migration if needed
+    client.delete_collection('new_collection')
+```
+
+Why to use this?
+
+Now we can manage the versions and executions of scripts that we do run over. We don't have a GUI to change configs for qdrant, we have to do it via notebook. Now the thing was, we were not able to manage it across that who and what is getting updated in the config. This script will make it easy, we don't anymore need to create notebooks and, you need API keys and URL to execute any changes on collection, adding up to access control. Plus, we can commit the config changes over org repo, so, we will be able to control the versions too.
diff --git a/cli.py b/cli.py
@@ -0,0 +1,47 @@
+import click
+from migrator import migrate, rollback
+
+
+@click.group()
+def cli():
+    pass
+
+
+@click.command()
+@click.option("--url", prompt="Qdrant URL", help="The URL of the Qdrant instance.")
+@click.option(
+    "--api-key", prompt="API Key", help="The API key for the Qdrant instance."
+)
+@click.option(
+    "--migration-folder",
+    prompt="Migration Folder",
+    help="The folder containing migration scripts.",
+)
+def migrate_cmd(url, api_key, migration_folder):
+    migrate(url, api_key, migration_folder)
+
+
+@click.command()
+@click.option("--url", prompt="Qdrant URL", help="The URL of the Qdrant instance.")
+@click.option(
+    "--api-key", prompt="API Key", help="The API key for the Qdrant instance."
+)
+@click.option(
+    "--migration-folder",
+    prompt="Migration Folder",
+    help="The folder containing migration scripts.",
+)
+@click.option(
+    "--target-version",
+    prompt="Target Version",
+    type=int,
+    help="The target version to rollback to.",
+)
+def rollback_cmd(url, api_key, migration_folder, target_version):
+    rollback(url, api_key, migration_folder, target_version)
+
+
+cli.add_command(migrate_cmd, name="migrate")
+cli.add_command(rollback_cmd, name="rollback")
+if __name__ == "__main__":
+    cli()
diff --git a/migrations/001_initial_migration.py b/migrations/001_initial_migration.py
@@ -0,0 +1,42 @@
+from qdrant_client import models
+
+
+def forward(client):
+    # Your migration logic here
+    client.update_collection(
+        collection_name="collection-name",
+        # vectors_config=models.VectorParams(
+        #     size=3072, distance=models.Distance.COSINE, on_disk=True
+        # ),
+        hnsw_config=models.HnswConfigDiff(
+            payload_m=16,
+            m=0,
+            on_disk=True,
+        ),
+        quantization_config=models.BinaryQuantization(
+            binary=models.BinaryQuantizationConfig(
+                always_ram=False,  # This mode allows to achieve the smallest memory footprint, but at the cost of the search speed.
+            ),
+        ),
+    )
+    indices = [
+        "organisation_id",
+        "document_id",
+        "ref_doc_id",
+        "doc_id",
+        "provider",
+        "category",
+    ]
+
+    for index in indices:
+        client.create_payload_index(
+            collection_name="collection-name",
+            field_name=index,
+            field_schema=models.PayloadSchemaType.KEYWORD,
+        )
+
+
+def backward(client):
+    # Logic to undo the migration if needed
+    # client.delete_collection("new_collection")
+    print("Your code to rollback the migration here")
diff --git a/migrator.py b/migrator.py
@@ -0,0 +1,126 @@
+import os
+import importlib.util
+from qdrant_client import QdrantClient, models
+import logging
+
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
+logging.getLogger("httpx").setLevel(logging.WARNING)
+
+
+def initialize_qdrant_client(url, api_key):
+    client = QdrantClient(url=url, api_key=api_key)
+    return client
+
+
+def check_and_create_migrations_collection(client):
+    collections_response = client.get_collections()
+    collections = collections_response.collections
+    collection_names = [collection.name for collection in collections]
+    if "migrations" not in collection_names:
+        client.create_collection(
+            collection_name="migrations",
+            vectors_config=models.VectorParams(size=2, distance=models.Distance.COSINE),
+        )
+        client.upsert(
+            collection_name="migrations",
+            points=[
+                models.PointStruct(
+                    id="5c56c793-69f3-4fbf-87e6-c4bf54c28c26",
+                    payload={
+                        "version": 0,
+                    },
+                    vector=[0.0, 0.1],
+                ),
+            ],
+        )
+
+
+def get_current_version(client):
+    points = client.scroll(
+        "migrations",
+        with_payload=True,
+    )
+    if points:
+        version = points[0][0].payload["version"]
+        return version
+    return 0
+
+
+def set_current_version(client, version):
+    client.upsert(
+        collection_name="migrations",
+        points=[
+            models.PointStruct(
+                id="5c56c793-69f3-4fbf-87e6-c4bf54c28c26",
+                payload={
+                    "version": version,
+                },
+                vector=[0.0, 0.1],
+            ),
+        ],
+    )
+
+
+def get_migration_files(migration_folder):
+    files = []
+    for file in os.listdir(migration_folder):
+        if file.endswith(".py"):
+            index = int(file.split("_")[0])
+            files.append((index, file))
+    files.sort(key=lambda x: x[0])
+    return files
+
+
+def run_migrations(client, migration_folder, current_version, target_version=None):
+    migration_files = get_migration_files(migration_folder)
+    for index, file in migration_files:
+        if index > current_version and (
+            target_version is None or index <= target_version
+        ):
+            module_name = file.replace(".py", "")
+            spec = importlib.util.spec_from_file_location(
+                module_name, os.path.join(migration_folder, file)
+            )
+            module = importlib.util.module_from_spec(spec)
+            spec.loader.exec_module(module)
+            module.forward(client)
+            set_current_version(client, index)
+            logger.info(
+                f"Migration completed successfully for {file}! Enjoy your migration :D"
+            )
+        else:
+            logger.info(f"Skipping migration {file}")
+
+
+def rollback_migrations(client, migration_folder, current_version, target_version):
+    migration_files = get_migration_files(migration_folder)
+    for index, file in reversed(migration_files):
+        if index <= current_version and index > target_version:
+            module_name = file.replace(".py", "")
+            spec = importlib.util.spec_from_file_location(
+                module_name, os.path.join(migration_folder, file)
+            )
+            module = importlib.util.module_from_spec(spec)
+            spec.loader.exec_module(module)
+            module.backward(client)
+            set_current_version(client, index - 1)
+            logger.info(
+                f"Rollback completed successfully for {file}! Enjoy your migration :D"
+            )
+        else:
+            logger.info(f"Skipping rollback {file}")
+
+
+def migrate(url, api_key, migration_folder):
+    client = initialize_qdrant_client(url, api_key)
+    check_and_create_migrations_collection(client)
+    current_version = get_current_version(client)
+    run_migrations(client, migration_folder, current_version)
+
+
+def rollback(url, api_key, migration_folder, target_version):
+    client = initialize_qdrant_client(url, api_key)
+    check_and_create_migrations_collection(client)
+    current_version = get_current_version(client)
+    rollback_migrations(client, migration_folder, current_version, target_version)
