From 3a8dbce14b3adb8a31c1fd97b40d6ee64d0ae08f Mon Sep 17 00:00:00 2001
From: Andrew Sisley <a.j.sisley@hotmail.com>
Date: Tue, 19 Nov 2024 13:31:22 -0500
Subject: [PATCH] WIP - Add support for col-level tt queries

---
 internal/db/fetcher/versioned.go              |  18 +-
 .../query/simple/with_cid_branchable_test.go  | 161 ++++++++++++++++++
 2 files changed, 177 insertions(+), 2 deletions(-)
 create mode 100644 tests/integration/query/simple/with_cid_branchable_test.go

diff --git a/internal/db/fetcher/versioned.go b/internal/db/fetcher/versioned.go
index a1b46d3c5f..ef43ce38e9 100644
--- a/internal/db/fetcher/versioned.go
+++ b/internal/db/fetcher/versioned.go
@@ -318,11 +318,21 @@ func (vf *VersionedFetcher) merge(c cid.Cid) error {
 
 	var mcrdt merklecrdt.MerkleCRDT
 	switch {
+	case block.Delta.IsCollection():
+		mcrdt = merklecrdt.NewMerkleCollection(
+			vf.store,
+			keys.NewCollectionSchemaVersionKey(vf.col.Description().SchemaVersionID, vf.col.Description().ID),
+			keys.NewHeadstoreColKey(vf.col.Description().RootID),
+		)
 	case block.Delta.IsComposite():
 		mcrdt = merklecrdt.NewMerkleCompositeDAG(
 			vf.store,
 			keys.NewCollectionSchemaVersionKey(block.Delta.GetSchemaVersionID(), vf.col.Description().RootID),
-			vf.dsKey.WithFieldID(core.COMPOSITE_NAMESPACE),
+			keys.DataStoreKey{
+				CollectionRootID: vf.col.Description().RootID,
+				DocID:            string(block.Delta.GetDocID()),
+				FieldID:          fmt.Sprint(core.COMPOSITE_NAMESPACE),
+			},
 		)
 	default:
 		field, ok := vf.col.Definition().GetFieldByName(block.Delta.GetFieldName())
@@ -335,7 +345,11 @@ func (vf *VersionedFetcher) merge(c cid.Cid) error {
 			keys.NewCollectionSchemaVersionKey(block.Delta.GetSchemaVersionID(), vf.col.Description().RootID),
 			field.Typ,
 			field.Kind,
-			vf.dsKey.WithFieldID(fmt.Sprint(field.ID)),
+			keys.DataStoreKey{
+				CollectionRootID: vf.col.Description().RootID,
+				DocID:            string(block.Delta.GetDocID()),
+				FieldID:          fmt.Sprint(field.ID),
+			},
 			field.Name,
 		)
 		if err != nil {
diff --git a/tests/integration/query/simple/with_cid_branchable_test.go b/tests/integration/query/simple/with_cid_branchable_test.go
new file mode 100644
index 0000000000..58590a5d5a
--- /dev/null
+++ b/tests/integration/query/simple/with_cid_branchable_test.go
@@ -0,0 +1,161 @@
+// Copyright 2024 Democratized Data Foundation
+//
+// Use of this software is governed by the Business Source License
+// included in the file licenses/BSL.txt.
+//
+// As of the Change Date specified in that file, in accordance with
+// the Business Source License, use of this software will be governed
+// by the Apache License, Version 2.0, included in the file
+// licenses/APL.txt.
+
+package simple
+
+import (
+	"testing"
+
+	testUtils "github.com/sourcenetwork/defradb/tests/integration"
+)
+
+func TestQuerySimpleWithCidOfBranchableCollection_FirstCid(t *testing.T) {
+	test := testUtils.TestCase{
+		Actions: []any{
+			testUtils.SchemaUpdate{
+				Schema: `
+					type Users @branchable {
+						name: String
+					}
+				`,
+			},
+			testUtils.CreateDoc{
+				Doc: `{
+					"name": "Fred"
+				}`,
+			},
+			testUtils.UpdateDoc{
+				Doc: `{
+					"name": "Freddddd"
+				}`,
+			},
+			testUtils.CreateDoc{
+				Doc: `{
+					"name": "John"
+				}`,
+			},
+			testUtils.Request{
+				Request: `query {
+					Users (
+							cid: "bafyreiewwsnu2ld5qlntamdm77ayb7xtmxz3p5difvaaakaome7zbtpo4u"
+						) {
+						name
+					}
+				}`,
+				Results: map[string]any{
+					"Users": []map[string]any{
+						{
+							"name": "Fred",
+						},
+					},
+				},
+			},
+		},
+	}
+
+	testUtils.ExecuteTestCase(t, test)
+}
+
+func TestQuerySimpleWithCidOfBranchableCollection_MiddleCid(t *testing.T) {
+	test := testUtils.TestCase{
+		Actions: []any{
+			testUtils.SchemaUpdate{
+				Schema: `
+					type Users @branchable {
+						name: String
+					}
+				`,
+			},
+			testUtils.CreateDoc{
+				Doc: `{
+					"name": "Fred"
+				}`,
+			},
+			testUtils.UpdateDoc{
+				Doc: `{
+					"name": "Freddddd"
+				}`,
+			},
+			testUtils.CreateDoc{
+				Doc: `{
+					"name": "John"
+				}`,
+			},
+			testUtils.Request{
+				Request: `query {
+					Users (
+							cid: "bafyreifpamlyhcbriztgbhds5ctgi5rm6w5wcar2py7246lo6j5v7iusxm"
+						) {
+						name
+					}
+				}`,
+				Results: map[string]any{
+					"Users": []map[string]any{
+						{
+							"name": "Freddddd",
+						},
+					},
+				},
+			},
+		},
+	}
+
+	testUtils.ExecuteTestCase(t, test)
+}
+
+func TestQuerySimpleWithCidOfBranchableCollection_LastCid(t *testing.T) {
+	test := testUtils.TestCase{
+		Actions: []any{
+			testUtils.SchemaUpdate{
+				Schema: `
+					type Users @branchable {
+						name: String
+					}
+				`,
+			},
+			testUtils.CreateDoc{
+				Doc: `{
+					"name": "Fred"
+				}`,
+			},
+			testUtils.UpdateDoc{
+				Doc: `{
+					"name": "Freddddd"
+				}`,
+			},
+			testUtils.CreateDoc{
+				Doc: `{
+					"name": "John"
+				}`,
+			},
+			testUtils.Request{
+				Request: `query {
+					Users (
+							cid: "bafyreigmt6ytph32jjxts2bij7fkne5ntionldsnklp35vcamvvl2x3a5i"
+						) {
+						name
+					}
+				}`,
+				Results: map[string]any{
+					"Users": []map[string]any{
+						{
+							"name": "Freddddd",
+						},
+						{
+							"name": "John",
+						},
+					},
+				},
+			},
+		},
+	}
+
+	testUtils.ExecuteTestCase(t, test)
+}