From 60c328270eda2964bb9ef5081e0bfe1c4e37e916 Mon Sep 17 00:00:00 2001
From: prasar-ashutosh <ashutosh.prasar@hk.email.gs.com>
Date: Tue, 29 Oct 2024 14:57:38 +0530
Subject: [PATCH] Fix dataset Reference to have dataset optional properties

---
 .../datasets/DatasetReferenceImplAbstract.java       |  3 +++
 .../RelationalSchemaEvolutionServiceTest.java        | 12 ++++++++----
 2 files changed, 11 insertions(+), 4 deletions(-)

diff --git a/legend-engine-xts-persistence/legend-engine-xt-persistence-component/legend-engine-xt-persistence-component-logical-plan/src/main/java/org/finos/legend/engine/persistence/components/logicalplan/datasets/DatasetReferenceImplAbstract.java b/legend-engine-xts-persistence/legend-engine-xt-persistence-component/legend-engine-xt-persistence-component-logical-plan/src/main/java/org/finos/legend/engine/persistence/components/logicalplan/datasets/DatasetReferenceImplAbstract.java
index 6ad30972cc6..2c96c932ab0 100644
--- a/legend-engine-xts-persistence/legend-engine-xt-persistence-component/legend-engine-xt-persistence-component-logical-plan/src/main/java/org/finos/legend/engine/persistence/components/logicalplan/datasets/DatasetReferenceImplAbstract.java
+++ b/legend-engine-xts-persistence/legend-engine-xt-persistence-component/legend-engine-xt-persistence-component-logical-plan/src/main/java/org/finos/legend/engine/persistence/components/logicalplan/datasets/DatasetReferenceImplAbstract.java
@@ -17,6 +17,8 @@
 import org.immutables.value.Value.Immutable;
 import org.immutables.value.Value.Style;
 
+import java.util.Optional;
+
 @Immutable
 @Style(
     typeAbstract = "*Abstract",
@@ -27,4 +29,5 @@
 )
 public interface DatasetReferenceImplAbstract extends DatasetReference
 {
+    Optional<DatasetAdditionalProperties> datasetAdditionalProperties();
 }
diff --git a/legend-engine-xts-persistence/legend-engine-xt-persistence-component/legend-engine-xt-persistence-component-relational-h2/src/test/java/org/finos/legend/engine/persistence/components/RelationalSchemaEvolutionServiceTest.java b/legend-engine-xts-persistence/legend-engine-xt-persistence-component/legend-engine-xt-persistence-component-relational-h2/src/test/java/org/finos/legend/engine/persistence/components/RelationalSchemaEvolutionServiceTest.java
index 17d64d089c7..21bc3dcd310 100644
--- a/legend-engine-xts-persistence/legend-engine-xt-persistence-component/legend-engine-xt-persistence-component-relational-h2/src/test/java/org/finos/legend/engine/persistence/components/RelationalSchemaEvolutionServiceTest.java
+++ b/legend-engine-xts-persistence/legend-engine-xt-persistence-component/legend-engine-xt-persistence-component-relational-h2/src/test/java/org/finos/legend/engine/persistence/components/RelationalSchemaEvolutionServiceTest.java
@@ -18,9 +18,7 @@
 import org.finos.legend.engine.persistence.components.ingestmode.audit.DateTimeAuditing;
 import org.finos.legend.engine.persistence.components.ingestmode.deduplication.FilterDuplicates;
 import org.finos.legend.engine.persistence.components.ingestmode.digest.UserProvidedDigestGenStrategy;
-import org.finos.legend.engine.persistence.components.logicalplan.datasets.DatasetDefinition;
-import org.finos.legend.engine.persistence.components.logicalplan.datasets.DatasetReference;
-import org.finos.legend.engine.persistence.components.logicalplan.datasets.DatasetReferenceImpl;
+import org.finos.legend.engine.persistence.components.logicalplan.datasets.*;
 import org.finos.legend.engine.persistence.components.relational.CaseConversion;
 import org.finos.legend.engine.persistence.components.relational.api.RelationalSchemaEvolutionService;
 import org.finos.legend.engine.persistence.components.relational.api.SchemaEvolutionServiceResult;
@@ -98,7 +96,13 @@ void testAddColumnUpperCase() throws Exception
     {
         DatasetDefinition mainTable = TestUtils.getSchemaEvolutionAddColumnMainTableUpperCase(); // This is only used to create a database table in upper case
         DatasetDefinition stagingTable = TestUtils.getBasicStagingTable();
-        DatasetReference mainTableDatasetReference = DatasetReferenceImpl.builder().group(testSchemaName).name(mainTableName).build(); // This is the model user has
+        DatasetReference mainTableDatasetReference = DatasetReferenceImpl.builder().group(testSchemaName)
+                .name(mainTableName)
+                .datasetAdditionalProperties(DatasetAdditionalProperties.builder().tableOrigin(TableOrigin.ICEBERG).build())
+                .build(); // This is the model user has
+
+        Assertions.assertEquals(DatasetAdditionalProperties.builder().tableOrigin(TableOrigin.ICEBERG).build(),
+                mainTableDatasetReference.datasetAdditionalProperties().get());
 
         // Create staging table
         createStagingTable(stagingTable);