Add staged package hashes

runners/google-cloud-dataflow-java/src/main/java/org/apache/beam/runners/dataflow/util/PackageUtil.java

@@ -438,6 +438,7 @@ public static PackageAttributes forFileToStage(
               .toString();
       destination.setLocation(resourcePath);
       destination.setName(dest);
+      destination.setSha256(hash);
       return new AutoValue_PackageUtil_PackageAttributes(
           file, null, destination, file.length(), hash);
     }
@@ -456,6 +457,7 @@ public static PackageAttributes forBytesToStage(
       DataflowPackage targetPackage = new DataflowPackage();
       targetPackage.setName(target);
       targetPackage.setLocation(resourcePath);
+      targetPackage.setSha256(hashCode.toString());
 
       return new AutoValue_PackageUtil_PackageAttributes(
           null, bytes, targetPackage, size, hashCode.toString());
@@ -465,6 +467,7 @@ public PackageAttributes withPackageName(String overridePackageName) {
       DataflowPackage newDestination = new DataflowPackage();
       newDestination.setName(overridePackageName);
       newDestination.setLocation(getDestination().getLocation());
+      newDestination.setSha256(getHash());
 
       return new AutoValue_PackageUtil_PackageAttributes(
           getSource(), getBytes(), newDestination, getSize(), getHash());

runners/google-cloud-dataflow-java/src/test/java/org/apache/beam/runners/dataflow/util/PackageUtilTest.java

@@ -191,6 +191,8 @@ public void testFileWithExtensionPackageNamingAndSize() throws Exception {
     assertThat(target.getName(), endsWith(".txt"));
     assertThat(target.getLocation(), equalTo(STAGING_PATH + target.getName()));
     assertThat(attr.getSize(), equalTo((long) contents.length()));
+    String expectedHash = Files.asByteSource(tmpFile).hash(Hashing.sha256()).toString();
+    assertThat(target.getSha256(), equalTo(expectedHash));
   }
 
   @Test
@@ -299,6 +301,8 @@ public void testPackageUploadWithFileSucceeds() throws Exception {
 
     assertThat(target.getName(), endsWith(".txt"));
     assertThat(target.getLocation(), equalTo(STAGING_PATH + target.getName()));
+    String expectedHash = Files.asByteSource(tmpFile).hash(Hashing.sha256()).toString();
+    assertThat(target.getSha256(), equalTo(expectedHash));
     assertThat(
         new LineReader(Channels.newReader(pipe.source(), StandardCharsets.UTF_8.name())).readLine(),
         equalTo(contents));

sdks/python/apache_beam/runners/dataflow/internal/apiclient.py

@@ -171,14 +171,15 @@ def __init__(
         self.proto.experiments.append(experiment)
     # Worker pool(s) information.
     package_descriptors = []
-    for package in packages:
+    for package, sha256 in packages:
       package_descriptors.append(
           dataflow.Package(
               location='%s/%s' % (
                   self.google_cloud_options.staging_location.replace(
                       'gs:/', GoogleCloudOptions.STORAGE_API_SERVICE),
                   package),
-              name=package))
+              name=package,
+              sha256=sha256))
 
     pool = dataflow.WorkerPool(
         kind='local' if self.local else 'harness',
@@ -649,7 +650,9 @@ def _stage_resources(self, pipeline, options):
     resource_stager = _LegacyDataflowStager(self)
     staged_resources = resource_stager.stage_job_resources(
         resources, staging_location=google_cloud_options.staging_location)
-    return staged_resources
+
+    name_to_hash = {remote_name: sha256 for _, remote_name, sha256 in resources}
+    return [(name, name_to_hash.get(name)) for name in staged_resources]
 
   def stage_file(
       self,

sdks/python/apache_beam/runners/dataflow/internal/apiclient_test.py

@@ -96,6 +96,22 @@ def test_pipeline_url(self):
       self.fail('No pipeline_url found in %s' % recovered_options)
 
     self.assertEqual(pipeline_url.string_value, FAKE_PIPELINE_URL)
+  def test_environment_packages_with_hash(self):
+    pipeline_options = PipelineOptions([
+        '--temp_location',
+        'gs://any-location/temp'
+    ])
+    packages = [('package1', 'hash1'), ('package2', 'hash2')]
+    env = apiclient.Environment(
+        packages,
+        pipeline_options,
+        '2.0.0',
+        FAKE_PIPELINE_URL)
+    self.assertEqual(len(env.proto.workerPools[0].packages), 2)
+    self.assertEqual(env.proto.workerPools[0].packages[0].name, 'package1')
+    self.assertEqual(env.proto.workerPools[0].packages[0].sha256, 'hash1')
+    self.assertEqual(env.proto.workerPools[0].packages[1].name, 'package2')
+    self.assertEqual(env.proto.workerPools[0].packages[1].sha256, 'hash2')
 
   def test_set_network(self):
     pipeline_options = PipelineOptions([