Implement entry referral tracking for Spanner Advanced Search

- Add support for tracking referral relationships between entries in Spanner
- Create new AdvancedSearchEntryReferral table to store referral connections
- Implement method to convert origin entry IDs to Spanner entry IDs
- Add a property graph to represent referral relationships
- Enhance error handling for batch write operations
- Collect and process referral relationships during entry synchronization

Author: syucream

entry/services.py

@@ -345,6 +345,11 @@ def update_documents(kls, entity: Entity, is_update: bool = False):
 
                 # Process entries in chunks to avoid too large mutation groups
                 entry_chunks = [entry_list[i : i + 100] for i in range(0, len(entry_list), 100)]
+                entry_referrals: set[tuple[str, int]] = (
+                    set()
+                )  # {(spanner_entry_id, referral_origin_entry_id), ...}
+
+                # First, write all entries and their attributes
                 for chunk in entry_chunks:
                     # Process each entry in the chunk with its own mutation group
                     with repo.database.mutation_groups() as mg:
@@ -387,15 +392,28 @@ def update_documents(kls, entity: Entity, is_update: bool = False):
                                 )
 
                                 for attrv in attr.prefetch_values:
-                                    spanner_attribute_values.append(
-                                        AdvancedSearchAttributeValue.create_instance(
-                                            entry_id=spanner_entry_id,
-                                            attribute_id=spanner_attr_id,
-                                            attribute_value_id=str(uuid.uuid4()),
-                                            entity_attr=entity_attr,
-                                            attrv=attrv,
-                                        )
+                                    value = AdvancedSearchAttributeValue.create_instance(
+                                        entry_id=spanner_entry_id,
+                                        attribute_id=spanner_attr_id,
+                                        attribute_value_id=str(uuid.uuid4()),
+                                        entity_attr=entity_attr,
+                                        attrv=attrv,
                                     )
+                                    spanner_attribute_values.append(value)
+
+                                    # Collect referral relationships
+                                    match entity_attr.type:
+                                        case AttrType.OBJECT | AttrType.NAMED_OBJECT:
+                                            if attrv.referral:
+                                                entry_referrals.add(
+                                                    (spanner_entry_id, attrv.referral.id)
+                                                )
+                                        case AttrType.ARRAY_OBJECT | AttrType.ARRAY_NAMED_OBJECT:
+                                            for array_value in attrv.data_array.all():
+                                                if array_value.referral:
+                                                    entry_referrals.add(
+                                                        (spanner_entry_id, array_value.referral.id)
+                                                    )
 
                             # Create a mutation group for this entry and its related data
                             group = mg.group()
@@ -407,8 +425,52 @@ def update_documents(kls, entity: Entity, is_update: bool = False):
 
                         # Batch write all mutation groups for this chunk
                         responses = mg.batch_write()
-                        if not all(response.status.code == 0 for response in responses):
-                            raise Exception(f"Failed to batch write to Spanner: {responses}")
+                        if any(response.status.code != 0 for response in responses):
+                            error_details = [
+                                (
+                                    f"code: {response.status.code}, "
+                                    f"message: {response.status.message}"
+                                )
+                                for response in responses
+                                if response.status.code != 0
+                            ]
+                            raise Exception(f"Failed to batch write to Spanner: {error_details}")
+
+                if entry_referrals:
+                    # Get mapping from OriginEntryId to EntryId for referrals
+                    referral_origin_ids = {ref_id for _, ref_id in entry_referrals}
+                    entry_id_mapping = repo.get_entry_id_mapping(list(referral_origin_ids))
+
+                    # Convert OriginEntryId to EntryId and filter out any missing mappings
+                    converted_referrals = [
+                        (entry_id, entry_id_mapping[ref_id])
+                        for entry_id, ref_id in entry_referrals
+                        if ref_id in entry_id_mapping
+                    ]
+
+                    if converted_referrals:
+                        with repo.database.mutation_groups() as mg:
+                            referral_chunks = [
+                                converted_referrals[i : i + 1000]
+                                for i in range(0, len(converted_referrals), 1000)
+                            ]
+                            for chunk in referral_chunks:
+                                group = mg.group()
+                                repo.insert_entry_referrals(chunk, group)
+
+                            responses = mg.batch_write()
+                            if any(response.status.code != 0 for response in responses):
+                                error_details = [
+                                    (
+                                        f"code: {response.status.code}, "
+                                        f"message: {response.status.message}"
+                                    )
+                                    for response in responses
+                                    if response.status.code != 0
+                                ]
+                                raise Exception(
+                                    f"Failed to batch write referrals to Spanner: {error_details}"
+                                )
 
             except Exception as e:
                 Logger.warning(f"Failed to sync data to Spanner: {e}")

entry/spanner/schema.sql

@@ -26,13 +26,25 @@ CREATE TABLE AdvancedSearchAttributeValue (
     AttributeId STRING(36) NOT NULL,
     AttributeValueId STRING(36) NOT NULL,
     Value STRING(MAX) NOT NULL,
-    RawValue JSON,
-
-    Value_Tokens TOKENLIST AS (TOKENIZE_FULLTEXT(Value)) HIDDEN
+    RawValue JSON
 ) PRIMARY KEY (EntryId, AttributeId, AttributeValueId),
   INTERLEAVE IN PARENT AdvancedSearchAttribute ON DELETE CASCADE;
 
 CREATE INDEX AdvancedSearchAttributeValueParentKeyIndex ON AdvancedSearchAttributeValue(EntryId, AttributeId);
 
-CREATE SEARCH INDEX AdvancedSearchAttributeValueValueIndex
-  ON AdvancedSearchAttributeValue(Value_Tokens);
+
+CREATE TABLE AdvancedSearchEntryReferral (
+    EntryId STRING(36) NOT NULL,
+    ReferralId STRING(36) NOT NULL,
+    FOREIGN KEY (EntryId) REFERENCES AdvancedSearchEntry (EntryId),
+    FOREIGN KEY (ReferralId) REFERENCES AdvancedSearchEntry (EntryId)
+) PRIMARY KEY (EntryId, ReferralId);
+
+CREATE PROPERTY GRAPH AdvancedSearchGraph
+  NODE TABLES (AdvancedSearchEntry)
+  EDGE TABLES (
+    AdvancedSearchEntryReferral
+      SOURCE KEY (EntryId) REFERENCES AdvancedSearchEntry (EntryId)
+      DESTINATION KEY (ReferralId) REFERENCES AdvancedSearchEntry (EntryId)
+      LABEL Referral
+  );

entry/spanner_advanced_search.py

@@ -237,6 +237,25 @@ def insert_attribute_values(
             ],
         )
 
+    def insert_entry_referrals(
+        self, entry_referrals: list[tuple[str, str]], operation: SpannerOperation
+    ) -> None:
+        """Insert entry referrals in a batch or mutation group.
+
+        Args:
+            entry_referrals: List of tuples (spanner_entry_id, referral_entry_id)
+            operation: Spanner operation (batch or mutation group)
+        """
+        if not entry_referrals:
+            return
+
+        # Insert the referral relationships
+        operation.insert(
+            table="AdvancedSearchEntryReferral",
+            columns=("EntryId", "ReferralId"),
+            values=entry_referrals,
+        )
+
     # Keep the single-record methods for backward compatibility
     def insert_entry(self, entry: AdvancedSearchEntry, transaction=None) -> None:
         """Insert a single entry"""
@@ -776,3 +795,29 @@ def get_joined_entries(
                 "ret_count": len(join_entries_dict),
                 "ret_values": list(join_entries_dict.values()),
             }
+
+    def get_entry_id_mapping(self, origin_entry_ids: list[int]) -> dict[int, str]:
+        """Get mapping from OriginEntryId to EntryId.
+
+        Args:
+            origin_entry_ids: List of original entry IDs from Django
+
+        Returns:
+            Dictionary mapping OriginEntryId to EntryId
+        """
+        if not origin_entry_ids:
+            return {}
+
+        query = """
+        SELECT OriginEntryId, EntryId
+        FROM AdvancedSearchEntry
+        WHERE OriginEntryId IN UNNEST(@origin_entry_ids)
+        """
+        params = {"origin_entry_ids": origin_entry_ids}
+        param_types = {
+            "origin_entry_ids": spanner_v1.param_types.Array(spanner_v1.param_types.INT64)
+        }
+
+        with self.database.snapshot() as snapshot:
+            results = snapshot.execute_sql(query, params=params, param_types=param_types)
+            return {row[0]: row[1] for row in results}