From da17df0a33796fd8aa7dd511b89879371f4f3ab5 Mon Sep 17 00:00:00 2001
From: tokebe <43009413+tokebe@users.noreply.github.com>
Date: Thu, 11 Apr 2024 16:55:26 -0400
Subject: [PATCH] feat: support source_record_urls

---
 src/graph/kg_edge.ts         | 14 +++++++++++---
 src/graph/knowledge_graph.ts |  1 +
 2 files changed, 12 insertions(+), 3 deletions(-)

diff --git a/src/graph/kg_edge.ts b/src/graph/kg_edge.ts
index b70ede02..797c82f1 100644
--- a/src/graph/kg_edge.ts
+++ b/src/graph/kg_edge.ts
@@ -20,6 +20,7 @@ export default class KGEdge {
         resource_id: string;
         resource_role: string;
         upstream_resource_ids?: Set<string>;
+        source_record_urls?: Set<string>;
       };
     };
   };
@@ -77,19 +78,26 @@ export default class KGEdge {
     }
     source.forEach((item) => {
       if (!this.sources[item.resource_id]) this.sources[item.resource_id] = {};
+      if (item.upstream_resource_ids && !Array.isArray(item.upstream_resource_ids)) {
+        item.upstream_resource_ids = [item.upstream_resource_ids];
+      }
+      if (item.source_record_urls && !Array.isArray(item.source_record_urls)) {
+        item.source_record_urls = [item.source_record_urls];
+      }
       if (!this.sources[item.resource_id][item.resource_role]) {
         this.sources[item.resource_id][item.resource_role] = {
           resource_id: item.resource_id,
           resource_role: item.resource_role,
           upstream_resource_ids: item.upstream_resource_ids ? new Set(item.upstream_resource_ids) : undefined,
+          source_record_urls: item.source_record_urls ? new Set(item.source_record_urls) : undefined,
         };
       }
-      if (item.upstream_resource_ids && !Array.isArray(item.upstream_resource_ids)) {
-        item.upstream_resource_ids = [item.upstream_resource_ids];
-      }
       item.upstream_resource_ids?.forEach((upstream) =>
         this.sources[item.resource_id][item.resource_role].upstream_resource_ids.add(upstream),
       );
+      item.source_record_urls?.forEach((url) =>
+        this.sources[item.resource_id][item.resource_role].source_record_urls.add(url),
+      );
     });
   }
 
diff --git a/src/graph/knowledge_graph.ts b/src/graph/knowledge_graph.ts
index c495c5a2..12a53476 100644
--- a/src/graph/knowledge_graph.ts
+++ b/src/graph/knowledge_graph.ts
@@ -137,6 +137,7 @@ export default class KnowledgeGraph {
         const trapiSource: TrapiSource = {
           ...sourceObj,
           upstream_resource_ids: sourceObj.upstream_resource_ids ? [...sourceObj.upstream_resource_ids] : undefined,
+          source_record_urls: sourceObj.source_record_urls ? [...sourceObj.source_record_urls] : undefined,
         };
         sources.push(trapiSource);
       });