From 91fda6174284740d96a66d9544c44c1a0383e5b7 Mon Sep 17 00:00:00 2001
From: Sergei Zaychenko <szaychenko@kamu.dev>
Date: Thu, 14 Dec 2023 08:47:13 -0800
Subject: [PATCH] Merge corrections

---
 .../tests/tests/test_gql_metadata_chain.rs    |  2 +-
 .../core/src/repos/dataset_repository.rs      |  1 +
 .../src/repos/dataset_repository_helpers.rs   | 34 +++++++++++--------
 .../tests/engine/test_engine_transform.rs     |  2 +-
 4 files changed, 23 insertions(+), 16 deletions(-)

diff --git a/src/adapter/graphql/tests/tests/test_gql_metadata_chain.rs b/src/adapter/graphql/tests/tests/test_gql_metadata_chain.rs
index 05f6a2d8bf..abdaa8804b 100644
--- a/src/adapter/graphql/tests/tests/test_gql_metadata_chain.rs
+++ b/src/adapter/graphql/tests/tests/test_gql_metadata_chain.rs
@@ -29,7 +29,7 @@ async fn test_metadata_chain_events() {
 
     let base_catalog = dill::CatalogBuilder::new()
         .add_builder(
-            dill::builder_for::<DatasetRepositoryLocalFs>()
+            DatasetRepositoryLocalFs::builder()
                 .with_root(tempdir.path().join("datasets"))
                 .with_multi_tenant(false),
         )
diff --git a/src/domain/core/src/repos/dataset_repository.rs b/src/domain/core/src/repos/dataset_repository.rs
index d88f16fe53..fc0f026d80 100644
--- a/src/domain/core/src/repos/dataset_repository.rs
+++ b/src/domain/core/src/repos/dataset_repository.rs
@@ -12,6 +12,7 @@ use std::pin::Pin;
 use std::sync::Arc;
 
 use async_trait::async_trait;
+use chrono::Utc;
 use internal_error::InternalError;
 use opendatafabric::*;
 use thiserror::Error;
diff --git a/src/infra/core/src/repos/dataset_repository_helpers.rs b/src/infra/core/src/repos/dataset_repository_helpers.rs
index de04129d99..c431e10a41 100644
--- a/src/infra/core/src/repos/dataset_repository_helpers.rs
+++ b/src/infra/core/src/repos/dataset_repository_helpers.rs
@@ -54,14 +54,14 @@ pub async fn create_dataset_from_snapshot_impl(
     // Validate / resolve events
     for event in snapshot.metadata.iter_mut() {
         match event {
-            MetadataEvent::Seed(_) => Err(InvalidSnapshotError {
-                reason: "Seed event is generated and cannot be specified explicitly".to_owned(),
-            }
+            MetadataEvent::Seed(_) => Err(InvalidSnapshotError::new(
+                "Seed event is generated and cannot be specified explicitly",
+            )
             .into()),
-            MetadataEvent::SetPollingSource(_) => {
+            MetadataEvent::SetPollingSource(_) | MetadataEvent::AddPushSource(_) => {
                 if snapshot.kind != DatasetKind::Root {
                     Err(InvalidSnapshotError {
-                        reason: "SetPollingSource is only allowed on root datasets".to_owned(),
+                        reason: format!("Event is only allowed on root datasets: {:?}", event),
                     }
                     .into())
                 } else {
@@ -70,26 +70,32 @@ pub async fn create_dataset_from_snapshot_impl(
             }
             MetadataEvent::SetTransform(e) => {
                 if snapshot.kind != DatasetKind::Derivative {
-                    Err(InvalidSnapshotError {
-                        reason: "SetTransform is only allowed on derivative datasets".to_owned(),
-                    }
+                    Err(InvalidSnapshotError::new(
+                        "SetTransform is only allowed on derivative datasets",
+                    )
                     .into())
                 } else {
                     resolve_transform_inputs(dataset_repo, &snapshot.name, &mut e.inputs).await
                 }
             }
+            MetadataEvent::SetDataSchema(_) => {
+                // It shouldn't be common to provide schema as part of the snapshot. In most
+                // cases it will inferred upon first ingest/transform. But no reason not to
+                // allow it.
+                Ok(())
+            }
             MetadataEvent::SetAttachments(_)
             | MetadataEvent::SetInfo(_)
             | MetadataEvent::SetLicense(_)
             | MetadataEvent::SetVocab(_) => Ok(()),
             MetadataEvent::AddData(_)
             | MetadataEvent::ExecuteQuery(_)
-            | MetadataEvent::SetWatermark(_) => Err(InvalidSnapshotError {
-                reason: format!(
-                    "Event is not allowed to appear in a DatasetSnapshot: {:?}",
-                    event
-                ),
-            }
+            | MetadataEvent::SetWatermark(_)
+            | MetadataEvent::DisablePollingSource(_)
+            | MetadataEvent::DisablePushSource(_) => Err(InvalidSnapshotError::new(format!(
+                "Event is not allowed to appear in a DatasetSnapshot: {:?}",
+                event
+            ))
             .into()),
         }?;
     }
diff --git a/src/infra/core/tests/tests/engine/test_engine_transform.rs b/src/infra/core/tests/tests/engine/test_engine_transform.rs
index 86798d3432..49ba94394a 100644
--- a/src/infra/core/tests/tests/engine/test_engine_transform.rs
+++ b/src/infra/core/tests/tests/engine/test_engine_transform.rs
@@ -243,7 +243,6 @@ async fn test_transform_common(transform: Transform) {
                 .with_data_format_registry(Arc::new(DataFormatRegistryImpl::new())),
         )
         .bind::<dyn PollingIngestService, PollingIngestServiceImpl>()
-        .add::<SystemTimeSourceDefault>()
         .add::<TransformServiceImpl>()
         .add_value(SystemTimeSourceStub::new_set(
             Utc.with_ymd_and_hms(2050, 1, 1, 12, 0, 0).unwrap(),
@@ -341,6 +340,7 @@ async fn test_transform_common(transform: Transform) {
     let deriv_helper = DatasetHelper::new(dataset.clone(), tempdir.path());
     let deriv_data_helper = DatasetDataHelper::new(dataset);
 
+    let time_source = catalog.get_one::<SystemTimeSourceStub>().unwrap();
     time_source.set(Utc.with_ymd_and_hms(2050, 1, 2, 12, 0, 0).unwrap());
 
     let res = transform_svc