Update for co_locate_with and name change to DocumentDB

Update for co_locate_with and name change to DocumentDB

Author: sandeepsnairms

.gitignore

@@ -0,0 +1 @@
+__pycache__/

README.md

@@ -1,6 +1,6 @@
-# Schema Transformer: Migrating RU-based Azure Cosmos DB for MongoDB to vCore-based
+# Schema Transformer: Migrating RU-based Azure Cosmos DB for MongoDB to Azure DocumentDB
 
-Schema Transformer is a Python script designed to analyze Mongo RU Collection schemas and efficiently transform them into a vCore-optimized structure. This ensures seamless compatibility and enhances query performance.
+Schema Transformer is a Python script designed to analyze Mongo RU Collection schemas and efficiently transform them into a DocumentDB optimized structure. This ensures seamless compatibility and enhances query performance.
 
 With this tool, you can generate index and sharding recommendations tailored specifically to your workload, making your migration smoother and more efficient.
 
@@ -9,16 +9,17 @@ With this tool, you can generate index and sharding recommendations tailored spe
 The tool supports the following versions:
 
 - **Source:** Azure Cosmos DB for MongoDB RU-based (version 4.2 and above)
-- **Target:** Azure Cosmos DB for MongoDB vCore (all versions)
+- **Target:** Azure DocumentDB (all versions)
 
 ## How to Run the Script
 
 ### Prerequisites
 
 Before running the assessment, ensure that the client machine meets the following requirements:
 
-- Access to both source and target MongoDB endpoints, either over a private or public network via the specified IP or hostname.
+- Access to both source MongoDB RU endpoint and target Azure DocumentDb endpoint, either over a private or public network via the specified IP or hostname.
 - Python (version 3.10 or above) must be installed.
+- PyMongo library must be installed (`pip install pymongo`).
 
 ### Steps to Run the Assessment
 
@@ -132,6 +133,27 @@ Before running the assessment, ensure that the client machine meets the followin
         }
         ```
 
+    6. To colocate collections with a reference collection
+    
+        ```json
+        {
+            "sections": [
+                {
+                    "include": [
+                        "db1.coll2",
+                        "db1.coll3"
+                    ],
+                    "migrate_shard_key": "false",
+                    "drop_if_exists": "true",
+                    "optimize_compound_indexes": "true",
+                    "co_locate_with": "coll1"
+                }
+            ]
+        }
+        ```
+        
+        **Note:** The collection specified in `co_locate_with` must already exist in the same database as the collection being processed. If the reference collection is not found, the script will fail with an error.
+
 4. Run the following command, providing the full path of the JSON file created in the previous step:
 
     ```cmd
@@ -148,3 +170,4 @@ This process will generate a vCore-optimized schema with index and sharding reco
 | **migrate_shard_key** | Determines whether the existing shard key definition should be migrated. If set to `True`, the shard key is retained; if `False`, the target collection remains unsharded. Collections that are originally unsharded in the source will remain unsharded in the target, regardless of this setting. **Default:** `False`. |
 | **drop_if_exists** | Specifies whether collections with the same name in the target should be dropped and recreated. If `True`, existing collections are removed before migration; if `False`, they remain unchanged. **Default:** `False`. |
 | **optimize_compound_indexes** | Controls whether compound indexes should be optimized. If `True`, the script identifies redundant indexes and excludes them from migration; if `False`, all indexes are migrated as-is. **Default:** `False`. |
+| **co_locate_with** | Specifies the name of a reference collection from the same database to colocate with. When specified, the target collection will be colocated with the reference collection for improved query performance. The reference collection must exist in the same database before colocation is applied, or an error will be thrown. This option is useful for optimizing queries that join or access related collections together. **Default:** `None`. |

collection_config.py

@@ -1,4 +1,5 @@
 from dataclasses import dataclass
+from typing import Optional
 
 @dataclass
 class CollectionConfig:
@@ -11,3 +12,4 @@ class CollectionConfig:
     migrate_shard_key: bool
     drop_if_exists: bool
     optimize_compound_indexes: bool = False
+    co_locate_with: Optional[str] = None

json_parser.py

@@ -32,6 +32,7 @@ def parse_json(self) -> List[CollectionConfig]:
             migrate_shard_key = section.get("migrate_shard_key", "false").lower() == "true"
             drop_if_exists = section.get("drop_if_exists", "false").lower() == "true"
             optimize_compound_indexes = section.get("optimize_compound_indexes", "false").lower() == "true"
+            co_locate_with = section.get("co_locate_with")
 
             for collection in collections_to_migrate:
                 if collection in collection_configs:
@@ -43,7 +44,8 @@ def parse_json(self) -> List[CollectionConfig]:
                     collection_name=collection_name,
                     migrate_shard_key=migrate_shard_key,
                     drop_if_exists=drop_if_exists,
-                    optimize_compound_indexes=optimize_compound_indexes
+                    optimize_compound_indexes=optimize_compound_indexes,
+                    co_locate_with=co_locate_with
                 )
                 collection_configs[collection] = collection_config
         return collection_configs.values()

schema_migration.py

@@ -54,6 +54,11 @@ def migrate_schema(
             else:
                 print("-- Target collection already exists. Skipping creation.")
 
+            # Handle colocation if specified
+            if collection_config.co_locate_with:
+                print(f"-- Setting up colocation with collection: {collection_config.co_locate_with}")
+                self._setup_colocation(dest_db, collection_name, collection_config.co_locate_with)
+
             # Check if shard key should be created
             if collection_config.migrate_shard_key:
                 source_shard_key = self._get_shard_key_ru(source_db, collection_config)
@@ -176,3 +181,33 @@ def _is_subarray(self, sub: List, main: List) -> bool:
             if main[i:i + sub_len] == sub:
                 return True
         return False
+
+    def _setup_colocation(self, dest_db: Database, collection_name: str, reference_collection: str) -> None:
+        """
+        Set up colocation for a collection with a reference collection.
+
+        :param dest_db: The destination database object.
+        :param collection_name: The name of the collection to colocate.
+        :param reference_collection: The name of the reference collection to colocate with.
+        :raises ValueError: If the reference collection does not exist.
+        """
+        # Check if reference collection exists
+        if reference_collection not in dest_db.list_collection_names():
+            raise ValueError(
+                f"Reference collection '{reference_collection}' not found in database '{dest_db.name}'. "
+                f"Cannot colocate collection '{collection_name}'."
+            )
+
+        # Run collMod command to set up colocation
+        try:
+            dest_db.command({
+                "collMod": collection_name,
+                "colocation": {
+                    "collection": reference_collection
+                }
+            })
+            print(f"---- Successfully colocated '{collection_name}' with '{reference_collection}'")
+        except Exception as e:
+            raise ValueError(
+                f"Failed to colocate collection '{collection_name}' with '{reference_collection}': {str(e)}"
+            )