Skip to content

Commit d515049

Browse files
misoltmisolt
authored
Add a summary YAML for Snowflake (#989)
* Add a DumpMetadataTask clone for Snowflake * Add a simple unit test * Add a test for getDescription * Change from shared YAML to separate file * Rewrite summary yaml generation
1 parent 75a0bef commit d515049

File tree

6 files changed

+210
-0
lines changed

6 files changed

+210
-0
lines changed

‎dumper/app/src/main/java/com/google/edwmigration/dumper/application/dumper/connector/snowflake/SnowflakeLiteConnector.java‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -60,6 +60,7 @@ public String getDescription() {
6060
public final void addTasksTo(List<? super Task<?>> out, ConnectorArguments arguments) {
6161
out.add(new DumpMetadataTask(arguments, FORMAT_NAME));
6262
out.add(new FormatTask(FORMAT_NAME));
63+
out.add(SnowflakeYamlSummaryTask.create(FORMAT_NAME, arguments));
6364
out.addAll(planner.generateLiteSpecificQueries());
6465
}
6566

‎dumper/app/src/main/java/com/google/edwmigration/dumper/application/dumper/connector/snowflake/SnowflakeLogsConnector.java‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -298,6 +298,7 @@ public final void addTasksTo(
298298
throws MetadataDumperUsageException {
299299
out.add(new DumpMetadataTask(arguments, FORMAT_NAME));
300300
out.add(new FormatTask(FORMAT_NAME));
301+
out.add(SnowflakeYamlSummaryTask.create(FORMAT_NAME, arguments));
301302

302303
// (24 * 7) -> 7 trailing days == 168 hours
303304
// Actually, on Snowflake, 7 days ago starts at midnight in an unadvertised time zone. What the

‎dumper/app/src/main/java/com/google/edwmigration/dumper/application/dumper/connector/snowflake/SnowflakeMetadataConnector.java‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -171,6 +171,7 @@ public final void addTasksTo(
171171
@Nonnull List<? super Task<?>> out, @Nonnull ConnectorArguments arguments) {
172172
out.add(new DumpMetadataTask(arguments, FORMAT_NAME));
173173
out.add(new FormatTask(FORMAT_NAME));
174+
out.add(SnowflakeYamlSummaryTask.create(FORMAT_NAME, arguments));
174175

175176
boolean INJECT_IS_FAULT = arguments.isTestFlag('A');
176177
// INFORMATION_SCHEMA queries must be qualified with a database

‎dumper/app/src/main/java/com/google/edwmigration/dumper/application/dumper/connector/snowflake/SnowflakeYamlSummaryTask.java‎

Lines changed: 134 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,134 @@
1+
/*
2+
* Copyright 2022-2025 Google LLC
3+
* Copyright 2013-2021 CompilerWorks
4+
*
5+
* Licensed under the Apache License, Version 2.0 (the "License");
6+
* you may not use this file except in compliance with the License.
7+
* You may obtain a copy of the License at
8+
*
9+
* http://www.apache.org/licenses/LICENSE-2.0
10+
*
11+
* Unless required by applicable law or agreed to in writing, software
12+
* distributed under the License is distributed on an "AS IS" BASIS,
13+
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
14+
* See the License for the specific language governing permissions and
15+
* limitations under the License.
16+
*/
17+
package com.google.edwmigration.dumper.application.dumper.connector.snowflake;
18+
19+
import static com.fasterxml.jackson.dataformat.yaml.YAMLGenerator.Feature.LITERAL_BLOCK_STYLE;
20+
import static com.fasterxml.jackson.dataformat.yaml.YAMLGenerator.Feature.SPLIT_LINES;
21+
import static com.fasterxml.jackson.dataformat.yaml.YAMLGenerator.Feature.WRITE_DOC_START_MARKER;
22+
import static java.lang.Integer.parseInt;
23+
import static java.nio.charset.StandardCharsets.UTF_8;
24+
import static java.util.Objects.requireNonNull;
25+
26+
import com.fasterxml.jackson.databind.ObjectMapper;
27+
import com.fasterxml.jackson.databind.SerializationFeature;
28+
import com.fasterxml.jackson.dataformat.yaml.YAMLFactory;
29+
import com.google.common.collect.ImmutableMap;
30+
import com.google.common.io.ByteSink;
31+
import com.google.edwmigration.dumper.application.dumper.ConnectorArguments;
32+
import com.google.edwmigration.dumper.application.dumper.handle.JdbcHandle;
33+
import com.google.edwmigration.dumper.application.dumper.task.AbstractJdbcTask;
34+
import com.google.edwmigration.dumper.application.dumper.task.TaskRunContext;
35+
import java.io.IOException;
36+
import java.io.UncheckedIOException;
37+
import java.io.Writer;
38+
import java.sql.Connection;
39+
import java.sql.ResultSet;
40+
import java.sql.SQLException;
41+
import java.util.Optional;
42+
import javax.annotation.Nonnull;
43+
import javax.annotation.Nullable;
44+
import javax.annotation.ParametersAreNonnullByDefault;
45+
import org.springframework.dao.DataAccessException;
46+
import org.springframework.jdbc.core.ResultSetExtractor;
47+
48+
/** A {@link Task} that creates YAML with extraction metadata. */
49+
@ParametersAreNonnullByDefault
50+
final class SnowflakeYamlSummaryTask extends AbstractJdbcTask<Void> {
51+
52+
private static final YAMLFactory yamlFactory =
53+
new YAMLFactory()
54+
.enable(LITERAL_BLOCK_STYLE)
55+
.disable(WRITE_DOC_START_MARKER)
56+
.disable(SPLIT_LINES);
57+
private static final ObjectMapper mapper =
58+
new ObjectMapper(yamlFactory).enable(SerializationFeature.INDENT_OUTPUT);
59+
60+
private final boolean isAssessment;
61+
62+
@Override
63+
public final String describeSourceData() {
64+
return "containing dump metadata.";
65+
}
66+
67+
static ImmutableMap<String, ?> createRoot(boolean isAssessment, Optional<String> optionalCount) {
68+
ImmutableMap.Builder<String, Object> builder = ImmutableMap.builder();
69+
builder.put("assessment", isAssessment);
70+
String count = optionalCount.orElse(null);
71+
int parsed;
72+
if (count == null) {
73+
parsed = 0;
74+
} else {
75+
parsed = parseInt(count);
76+
}
77+
builder.put("warehouse_count", parsed);
78+
return ImmutableMap.of("metadata", builder.build());
79+
}
80+
81+
static String rootString(boolean isAssessment, Optional<String> optionalCount) {
82+
try {
83+
ImmutableMap<String, ?> root = createRoot(isAssessment, optionalCount);
84+
return mapper.writeValueAsString(root);
85+
} catch (IOException e) {
86+
throw new UncheckedIOException(e);
87+
}
88+
}
89+
90+
@Nonnull
91+
static SnowflakeYamlSummaryTask create(String zipFormat, ConnectorArguments arguments) {
92+
return new SnowflakeYamlSummaryTask(arguments.isAssessment());
93+
}
94+
95+
private SnowflakeYamlSummaryTask(boolean isAssessment) {
96+
super("snowflake.yaml");
97+
this.isAssessment = isAssessment;
98+
}
99+
100+
@Override
101+
protected Void doInConnection(
102+
TaskRunContext context, JdbcHandle jdbcHandle, ByteSink sink, Connection connection)
103+
throws SQLException {
104+
String sql = "SHOW WAREHOUSES ->> SELECT count(*) FROM $1";
105+
return doSelect(connection, new Extractor(sink, isAssessment), sql);
106+
}
107+
108+
static final class Extractor implements ResultSetExtractor<Void> {
109+
private final ByteSink sink;
110+
private final boolean isAssessment;
111+
112+
Extractor(ByteSink sink, boolean isAssessment) {
113+
this.sink = sink;
114+
this.isAssessment = isAssessment;
115+
}
116+
117+
@Override
118+
public Void extractData(@Nullable ResultSet rs) throws SQLException, DataAccessException {
119+
doExtract(requireNonNull(rs));
120+
return null;
121+
}
122+
123+
void doExtract(ResultSet resultSet) throws SQLException, DataAccessException {
124+
resultSet.next();
125+
Optional<String> count = Optional.ofNullable(resultSet.getString(1));
126+
String yaml = rootString(isAssessment, count);
127+
try (Writer writer = sink.asCharSink(UTF_8).openBufferedStream()) {
128+
mapper.writeValue(writer, yaml);
129+
} catch (IOException e) {
130+
throw new UncheckedIOException(e);
131+
}
132+
}
133+
}
134+
}

‎dumper/app/src/test/java/com/google/edwmigration/dumper/application/dumper/connector/snowflake/SnowflakeLiteConnectorTest.java‎

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -39,6 +39,15 @@ public void getDefaultFileName_success() {
3939
assertTrue(name, name.contains("lite"));
4040
}
4141

42+
@Test
43+
public void getDescription_success() {
44+
45+
String description = connector.getDescription();
46+
47+
assertTrue(description, description.contains("lite"));
48+
assertTrue(description, description.contains("Snowflake"));
49+
}
50+
4251
@Test
4352
public void validate_databaseFlag_throwsException() {
4453
ImmutableList<String> list =

‎dumper/app/src/test/java/com/google/edwmigration/dumper/application/dumper/connector/snowflake/SnowflakeYamlSummaryTaskTest.java‎

Lines changed: 64 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,64 @@
1+
/*
2+
* Copyright 2022-2025 Google LLC
3+
* Copyright 2013-2021 CompilerWorks
4+
*
5+
* Licensed under the Apache License, Version 2.0 (the "License");
6+
* you may not use this file except in compliance with the License.
7+
* You may obtain a copy of the License at
8+
*
9+
* http://www.apache.org/licenses/LICENSE-2.0
10+
*
11+
* Unless required by applicable law or agreed to in writing, software
12+
* distributed under the License is distributed on an "AS IS" BASIS,
13+
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
14+
* See the License for the specific language governing permissions and
15+
* limitations under the License.
16+
*/
17+
package com.google.edwmigration.dumper.application.dumper.connector.snowflake;
18+
19+
import static com.google.edwmigration.dumper.application.dumper.connector.snowflake.SnowflakeYamlSummaryTask.createRoot;
20+
import static com.google.edwmigration.dumper.application.dumper.connector.snowflake.SnowflakeYamlSummaryTask.rootString;
21+
import static org.junit.Assert.assertEquals;
22+
import static org.junit.Assert.assertTrue;
23+
24+
import com.google.common.collect.ImmutableMap;
25+
import java.util.Optional;
26+
import org.junit.Test;
27+
import org.junit.runner.RunWith;
28+
import org.junit.runners.JUnit4;
29+
30+
@RunWith(JUnit4.class)
31+
public class SnowflakeYamlSummaryTaskTest {
32+
33+
@Test
34+
public void rootString_success() {
35+
36+
String rootString = rootString(true, Optional.of("2345"));
37+
38+
assertTrue(rootString, rootString.contains("assessment"));
39+
assertTrue(rootString, rootString.contains("true"));
40+
assertTrue(rootString, rootString.contains("warehouse_count"));
41+
assertTrue(rootString, rootString.contains("2345"));
42+
}
43+
44+
@Test
45+
public void createRoot_countAbsent_setToZero() {
46+
47+
ImmutableMap<String, ?> root = createRoot(true, Optional.empty());
48+
49+
assertTrue(root.toString(), root.containsKey("metadata"));
50+
ImmutableMap<?, ?> value = (ImmutableMap<?, ?>) root.get("metadata");
51+
assertEquals(0, value.get("warehouse_count"));
52+
}
53+
54+
@Test
55+
public void createRoot_noAssessment_success() {
56+
57+
ImmutableMap<String, ?> root = createRoot(false, Optional.of("123"));
58+
59+
assertTrue(root.toString(), root.containsKey("metadata"));
60+
ImmutableMap<?, ?> value = (ImmutableMap<?, ?>) root.get("metadata");
61+
assertEquals(false, value.get("assessment"));
62+
assertEquals(123, value.get("warehouse_count"));
63+
}
64+
}

0 commit comments

Comments
 (0)