PCSM-200: Shard collections with non-default collation (#139)

Author: inelpandzic

plm/catalog.go

@@ -378,6 +378,22 @@ func (c *Catalog) CreateIndexes(
 			lg.Info("Create hidden index as unhidden: " + index.Name)
 		}
 
+		if index.Collation == nil {
+			lg.Info("Create index with missing collation, setting simple collation: " + index.Name)
+
+			d := bson.D{{"locale", "simple"}}
+
+			collation, err := bson.Marshal(d)
+			if err != nil {
+				return errors.Wrapf(err, "marshal simple collation for index %s.%s.%s",
+					db, coll, index.Name)
+			}
+
+			idxCopy := *index
+			idxCopy.Collation = collation
+			index = &idxCopy
+		}
+
 		idxs = append(idxs, index)
 	}
 
@@ -1147,6 +1163,7 @@ func (c *Catalog) ShardCollection(
 	cmd := bson.D{
 		{"shardCollection", db + "." + coll},
 		{"key", shardKey},
+		{"collation", bson.D{{"locale", "simple"}}},
 	}
 
 	if unique {

plm/clone.go

@@ -417,12 +417,7 @@ func (c *Clone) doCollectionClone(
 		return ErrTimeseriesUnsupported
 	}
 
-	shInfo, err := topo.GetCollectionShardingInfo(ctx, c.source, ns.Database, ns.Collection)
-	if err != nil && !errors.Is(err, topo.ErrNotFound) {
-		return errors.Wrap(err, "get sharding info")
-	}
-
-	err = c.createCollection(ctx, ns, spec, shInfo)
+	err = c.createCollection(ctx, ns, spec)
 	if err != nil {
 		if !errors.Is(err, context.Canceled) {
 			lg.Errorf(err, "Failed to create %q collection", ns.String())
@@ -438,7 +433,21 @@ func (c *Clone) doCollectionClone(
 		}
 	}
 
-	lg.Infof("Collection %q has been created", ns.String())
+	lg.Infof("Collection %q created", ns.String())
+
+	shInfo, err := topo.GetCollectionShardingInfo(ctx, c.source, ns.Database, ns.Collection)
+	if err != nil && !errors.Is(err, topo.ErrNotFound) {
+		return errors.Wrap(err, "get sharding info")
+	}
+
+	if shInfo != nil && shInfo.IsSharded() {
+		err := c.catalog.ShardCollection(ctx, ns.Database, ns.Collection, shInfo.ShardKey, shInfo.Unique)
+		if err != nil {
+			return errors.Wrap(err, "shard collection")
+		}
+	}
+
+	lg.Infof("Collection %q sharded", ns.String())
 
 	c.catalog.SetCollectionTimestamp(ctx, ns.Database, ns.Collection, capturedAt)
 
@@ -694,7 +703,6 @@ func (c *Clone) createCollection(
 	ctx context.Context,
 	ns Namespace,
 	spec *topo.CollectionSpecification,
-	shInfo *topo.ShardingInfo,
 ) error {
 	if spec.Type == topo.TypeTimeseries {
 		return ErrTimeseriesUnsupported
@@ -717,13 +725,6 @@ func (c *Clone) createCollection(
 		return errors.Wrap(err, "create collection")
 	}
 
-	if shInfo != nil && shInfo.IsSharded() {
-		err := c.catalog.ShardCollection(ctx, ns.Database, ns.Collection, shInfo.ShardKey, shInfo.Unique)
-		if err != nil {
-			return errors.Wrap(err, "shard collection")
-		}
-	}
-
 	return nil
 }
 

tests/test_collections_sharded.py

@@ -38,3 +38,33 @@ def test_rename_sharded(t: Testing, phase: Runner.Phase):
         t.source["db_1"]["coll_1"].rename("coll_2")
 
     t.compare_all()
+
+
+@pytest.mark.parametrize("phase", [Runner.Phase.APPLY, Runner.Phase.CLONE])
+def test_create_collection_with_collation(t: Testing, phase: Runner.Phase):
+    with t.run(phase):
+        t.source["db_1"].create_collection("coll_1", collation={"locale": "en", "strength": 2})
+        t.source.admin.command(
+            "shardCollection", "db_1.coll_1", key={"name": 1}, collation={"locale": "simple"}
+        )
+
+    t.compare_all_sharded()
+
+
+@pytest.mark.parametrize("phase", [Runner.Phase.APPLY, Runner.Phase.CLONE])
+def test_create_collection_with_collation_with_shard_key_index_prefix(
+    t: Testing, phase: Runner.Phase
+):
+    with t.run(phase):
+        t.source["db_1"].create_collection("coll_2", collation={"locale": "en", "strength": 2})
+        t.source["db_1"]["coll_2"].create_index(
+            [("name", 1), ("date", 1), ("age", 1)], collation={"locale": "simple"}
+        )
+        t.source.admin.command(
+            "shardCollection",
+            "db_1.coll_1",
+            key={"name": 1, "date": 1},
+            collation={"locale": "simple"},
+        )
+
+    t.compare_all_sharded()

tests/test_indexes.py

@@ -24,6 +24,16 @@ def test_create_with_collation(t: Testing, phase: Runner.Phase):
     t.compare_all()
 
 
+@pytest.mark.parametrize("phase", [Runner.Phase.APPLY, Runner.Phase.CLONE])
+def test_create_with_inherited_collation(t: Testing, phase: Runner.Phase):
+    with t.run(phase):
+        t.source["db_1"].create_collection("coll_1_collation", collation={"locale": "en_US"})
+        t.source["db_1"]["coll_1_collation"].create_index({"i": 1})
+        t.source["db_1"]["coll_1_collation"].create_index({"j": 1}, collation={"locale": "simple"})
+
+    t.compare_all()
+
+
 @pytest.mark.parametrize("phase", [Runner.Phase.APPLY, Runner.Phase.CLONE])
 def test_create_unique(t: Testing, phase: Runner.Phase):
     with t.run(phase):