From b4b7095ed30b9bf8051fde3d3060eb2c37b7f71a Mon Sep 17 00:00:00 2001 From: Weimin Yu Date: Thu, 30 Jul 2020 14:28:53 -0400 Subject: [PATCH] Run InitSqlPipeline (#727) * Run InitSqlPipeline Added the main() method to InitSqlPipeline. Added a Gradle task to run InitSqlPipeline from command line. This task is meant for testing and experiments. Corrected the file name prefix of Datastore export files. Should be 'output-', defined as 'input-'. --- core/build.gradle | 24 +++++++++++++++++++ .../registry/beam/initsql/BackupPaths.java | 2 +- .../beam/initsql/InitSqlPipeline.java | 7 ++++++ .../beam/initsql/BackupTestStoreTest.java | 14 +++++------ .../initsql/ExportloadingTransformsTest.java | 18 +++++++------- 5 files changed, 48 insertions(+), 17 deletions(-) diff --git a/core/build.gradle b/core/build.gradle index be9d593e1..85bdbbd61 100644 --- a/core/build.gradle +++ b/core/build.gradle @@ -819,6 +819,30 @@ createToolTask( 'google.registry.tools.DevTool', sourceSets.nonprod) +project.tasks.create('initSqlPipeline', JavaExec) { + main = 'google.registry.beam.initsql.InitSqlPipeline' + + doFirst { + getToolArgsList().ifPresent { + args it + } + + def isDirectRunner = + args.contains('DirectRunner') || args.contains('--runner=DirectRunner') + // The dependency containing DirectRunner is intentionally excluded from the + // production binary, so that it won't be chosen by mistake: we definitely do + // not want to use it for the real jobs, yet DirectRunner is the default if + // the user forgets to override it. + // DirectRunner is required for tests and is already on testRuntimeClasspath. + // For simplicity, we add testRuntimeClasspath to this task's classpath instead + // of defining a new configuration just for the DirectRunner dependency. + classpath = + isDirectRunner + ? sourceSets.main.runtimeClasspath.plus(sourceSets.test.runtimeClasspath) + : sourceSets.main.runtimeClasspath + } +} + project.tasks.create('generateSqlSchema', JavaExec) { classpath = sourceSets.nonprod.runtimeClasspath main = 'google.registry.tools.DevTool' diff --git a/core/src/main/java/google/registry/beam/initsql/BackupPaths.java b/core/src/main/java/google/registry/beam/initsql/BackupPaths.java index 33916be85..c5be89061 100644 --- a/core/src/main/java/google/registry/beam/initsql/BackupPaths.java +++ b/core/src/main/java/google/registry/beam/initsql/BackupPaths.java @@ -32,7 +32,7 @@ public final class BackupPaths { private BackupPaths() {} private static final String WILDCARD_CHAR = "*"; - private static final String EXPORT_PATTERN_TEMPLATE = "%s/all_namespaces/kind_%s/input-%s"; + private static final String EXPORT_PATTERN_TEMPLATE = "%s/all_namespaces/kind_%s/output-%s"; public static final String COMMIT_LOG_NAME_PREFIX = "commit_diff_until_"; private static final String COMMIT_LOG_PATTERN_TEMPLATE = "%s/" + COMMIT_LOG_NAME_PREFIX + "*"; diff --git a/core/src/main/java/google/registry/beam/initsql/InitSqlPipeline.java b/core/src/main/java/google/registry/beam/initsql/InitSqlPipeline.java index 81b437e92..546103cf9 100644 --- a/core/src/main/java/google/registry/beam/initsql/InitSqlPipeline.java +++ b/core/src/main/java/google/registry/beam/initsql/InitSqlPipeline.java @@ -40,6 +40,7 @@ import java.util.Collection; import java.util.Optional; import org.apache.beam.sdk.Pipeline; import org.apache.beam.sdk.PipelineResult; +import org.apache.beam.sdk.options.PipelineOptionsFactory; import org.apache.beam.sdk.transforms.ParDo; import org.apache.beam.sdk.transforms.SerializableFunction; import org.apache.beam.sdk.transforms.Wait; @@ -234,4 +235,10 @@ public class InitSqlPipeline implements Serializable { return entityClasses.stream().map(Key::getKind).collect(ImmutableList.toImmutableList()); } } + + public static void main(String[] args) { + InitSqlPipelineOptions options = + PipelineOptionsFactory.fromArgs(args).withValidation().as(InitSqlPipelineOptions.class); + new InitSqlPipeline(options).run(); + } } diff --git a/core/src/test/java/google/registry/beam/initsql/BackupTestStoreTest.java b/core/src/test/java/google/registry/beam/initsql/BackupTestStoreTest.java index bbef30f72..1326812ce 100644 --- a/core/src/test/java/google/registry/beam/initsql/BackupTestStoreTest.java +++ b/core/src/test/java/google/registry/beam/initsql/BackupTestStoreTest.java @@ -109,9 +109,9 @@ public class BackupTestStoreTest { .map(string -> string.substring(exportFolder.getAbsolutePath().length()))) { assertThat(files) .containsExactly( - "/all_namespaces/kind_Registry/input-0", - "/all_namespaces/kind_DomainBase/input-0", - "/all_namespaces/kind_ContactResource/input-0"); + "/all_namespaces/kind_Registry/output-0", + "/all_namespaces/kind_DomainBase/output-0", + "/all_namespaces/kind_ContactResource/output-0"); } } @@ -132,16 +132,16 @@ public class BackupTestStoreTest { String exportRootPath = tempDir.getAbsolutePath(); File exportFolder = export(exportRootPath, Collections.EMPTY_SET); ImmutableList loadedRegistries = - loadExportedEntities(new File(exportFolder, "/all_namespaces/kind_Registry/input-0")); + loadExportedEntities(new File(exportFolder, "/all_namespaces/kind_Registry/output-0")); assertThat(loadedRegistries).containsExactly(registry); ImmutableList loadedDomains = - loadExportedEntities(new File(exportFolder, "/all_namespaces/kind_DomainBase/input-0")); + loadExportedEntities(new File(exportFolder, "/all_namespaces/kind_DomainBase/output-0")); assertThat(loadedDomains).containsExactly(domain); ImmutableList loadedContacts = loadExportedEntities( - new File(exportFolder, "/all_namespaces/kind_ContactResource/input-0")); + new File(exportFolder, "/all_namespaces/kind_ContactResource/output-0")); assertThat(loadedContacts).containsExactly(contact); } @@ -156,7 +156,7 @@ public class BackupTestStoreTest { export( exportRootPath, ImmutableSet.of(Key.create(getCrossTldKey(), Registry.class, "tld1"))); ImmutableList loadedRegistries = - loadExportedEntities(new File(exportFolder, "/all_namespaces/kind_Registry/input-0")); + loadExportedEntities(new File(exportFolder, "/all_namespaces/kind_Registry/output-0")); assertThat(loadedRegistries).containsExactly(newRegistry); } diff --git a/core/src/test/java/google/registry/beam/initsql/ExportloadingTransformsTest.java b/core/src/test/java/google/registry/beam/initsql/ExportloadingTransformsTest.java index a3ac3e1db..0fa2486b5 100644 --- a/core/src/test/java/google/registry/beam/initsql/ExportloadingTransformsTest.java +++ b/core/src/test/java/google/registry/beam/initsql/ExportloadingTransformsTest.java @@ -119,9 +119,9 @@ class ExportloadingTransformsTest implements Serializable { ImmutableList expectedPatterns = ImmutableList.of( - exportDir.getAbsolutePath() + "/all_namespaces/kind_Registry/input-*", - exportDir.getAbsolutePath() + "/all_namespaces/kind_DomainBase/input-*", - exportDir.getAbsolutePath() + "/all_namespaces/kind_ContactResource/input-*"); + exportDir.getAbsolutePath() + "/all_namespaces/kind_Registry/output-*", + exportDir.getAbsolutePath() + "/all_namespaces/kind_DomainBase/output-*", + exportDir.getAbsolutePath() + "/all_namespaces/kind_ContactResource/output-*"); PAssert.that(patterns).containsInAnyOrder(expectedPatterns); @@ -135,10 +135,10 @@ class ExportloadingTransformsTest implements Serializable { .apply( "File patterns to metadata", Create.of( - exportDir.getAbsolutePath() + "/all_namespaces/kind_Registry/input-*", - exportDir.getAbsolutePath() + "/all_namespaces/kind_DomainBase/input-*", + exportDir.getAbsolutePath() + "/all_namespaces/kind_Registry/output-*", + exportDir.getAbsolutePath() + "/all_namespaces/kind_DomainBase/output-*", exportDir.getAbsolutePath() - + "/all_namespaces/kind_ContactResource/input-*") + + "/all_namespaces/kind_ContactResource/output-*") .withCoder(StringUtf8Coder.of())) .apply(Transforms.getFilesByPatterns()); @@ -157,9 +157,9 @@ class ExportloadingTransformsTest implements Serializable { ImmutableList expectedFilenames = ImmutableList.of( - exportDir.getAbsolutePath() + "/all_namespaces/kind_Registry/input-0", - exportDir.getAbsolutePath() + "/all_namespaces/kind_DomainBase/input-0", - exportDir.getAbsolutePath() + "/all_namespaces/kind_ContactResource/input-0"); + exportDir.getAbsolutePath() + "/all_namespaces/kind_Registry/output-0", + exportDir.getAbsolutePath() + "/all_namespaces/kind_DomainBase/output-0", + exportDir.getAbsolutePath() + "/all_namespaces/kind_ContactResource/output-0"); PAssert.that(fileNames).containsInAnyOrder(expectedFilenames);