Add a md5sum check before overwriting a cached prism file (#34715)

* Check md5 of downloaded/unzipped file and cached before overwriting.

* Allow prism location to be a path from supported filesystems like GCS.

* Apply yapf

Author: shunping

sdks/python/apache_beam/runners/portability/prism_runner.py

@@ -22,6 +22,7 @@
 # sunset it
 from __future__ import annotations
 
+import hashlib
 import logging
 import os
 import platform
@@ -81,6 +82,35 @@ def create_job_service_handle(self, job_service, options):
         job_service, options, retain_unknown_options=True)
 
 
+def _md5sum(filename, block_size=8192) -> str:
+  md5 = hashlib.md5()
+  with open(filename, 'rb') as f:
+    while True:
+      data = f.read(block_size)
+      if not data:
+        break
+      md5.update(data)
+  return md5.hexdigest()
+
+
+def _rename_if_different(src, dst):
+  assert (os.path.isfile(src))
+
+  if os.path.isfile(dst):
+    if _md5sum(src) != _md5sum(dst):
+      # Remove existing binary to prevent exception on Windows during
+      # os.rename.
+      # See: https://docs.python.org/3/library/os.html#os.rename
+      os.remove(dst)
+      os.rename(src, dst)
+    else:
+      _LOGGER.info(
+          'Found %s and %s with the same md5. Skipping overwrite.' % (src, dst))
+      os.remove(src)
+  else:
+    os.rename(src, dst)
+
+
 class PrismJobServer(job_server.SubprocessJobServer):
   PRISM_CACHE = os.path.expanduser("~/.apache_beam/cache/prism")
   BIN_CACHE = os.path.expanduser("~/.apache_beam/cache/prism/bin")
@@ -117,7 +147,13 @@ def maybe_unzip_and_make_executable(
         # True (cache disabled)
         _LOGGER.info("Unzipping prism from %s to %s" % (url, target_url))
         z = zipfile.ZipFile(url)
-        target_url = z.extract(target, path=bin_cache)
+
+        bin_cache_tmp = os.path.join(bin_cache, 'tmp')
+        if not os.path.exists(bin_cache_tmp):
+          os.makedirs(bin_cache_tmp)
+        target_tmp_url = z.extract(target, path=bin_cache_tmp)
+
+        _rename_if_different(target_tmp_url, target_url)
     else:
       target_url = url
 
@@ -154,12 +190,8 @@ def local_bin(
             url_read = urlopen(url)
           with open(cached_file + '.tmp', 'wb') as zip_write:
             shutil.copyfileobj(url_read, zip_write, length=1 << 20)
-          if os.path.isfile(cached_file):
-            # Remove existing binary to prevent exception on Windows during
-            # os.rename.
-            # See: https://docs.python.org/3/library/os.html#os.rename
-            os.remove(cached_file)
-          os.rename(cached_file + '.tmp', cached_file)
+
+          _rename_if_different(cached_file + '.tmp', cached_file)
         except URLError as e:
           raise RuntimeError(
               'Unable to fetch remote prism binary at %s: %s' % (url, e))
@@ -204,6 +236,10 @@ def path_to_binary(self) -> str:
         # The path is local and exists, use directly.
         return self._path
 
+      if FileSystems.exists(self._path):
+        # The path is in one of the supported filesystems.
+        return self._path
+
       # Check if the path is a URL.
       url = urllib.parse.urlparse(self._path)
       if not url.scheme:

sdks/python/apache_beam/runners/portability/prism_runner_test.py

@@ -245,32 +245,37 @@ def tearDown(self) -> None:
     rmtree(self.local_dir)
     pass
 
-  def _make_local_bin(self):
-    with open(self.local_bin_path, 'wb'):
+  def _make_local_bin(self, fn=None):
+    fn = fn or self.local_bin_path
+    with open(fn, 'wb'):
       pass
 
-  def _make_local_zip(self):
-    with zipfile.ZipFile(self.local_zip_path, 'w', zipfile.ZIP_DEFLATED):
+  def _make_local_zip(self, fn=None):
+    fn = fn or self.local_zip_path
+    with zipfile.ZipFile(fn, 'w', zipfile.ZIP_DEFLATED):
       pass
 
-  def _make_cache_bin(self):
-    with open(self.cache_bin_path, 'wb'):
+  def _make_cache_bin(self, fn=None):
+    fn = fn or self.cache_bin_path
+    with open(fn, 'wb'):
       pass
 
-  def _make_cache_zip(self):
-    with zipfile.ZipFile(self.cache_zip_path, 'w', zipfile.ZIP_DEFLATED):
+  def _make_cache_zip(self, fn=None):
+    fn = fn or self.cache_zip_path
+    with zipfile.ZipFile(fn, 'w', zipfile.ZIP_DEFLATED):
       pass
 
-  def _extract_side_effect(self, a, path=None):
+  def _extract_side_effect(self, fn, path=None):
     if path is None:
-      return a
+      return fn
 
+    full_path = os.path.join(str(path), fn)
     if path.startswith(self.cache_dir):
-      self._make_cache_bin()
+      self._make_cache_bin(full_path)
     else:
-      self._make_local_bin()
+      self._make_local_bin(full_path)
 
-    return os.path.join(str(path), a)
+    return full_path
 
   @parameterized.expand([[True, True], [True, False], [False, True],
                          [False, False]])
@@ -352,7 +357,11 @@ def test_with_remote_path(self, has_cache_bin, has_cache_zip, ignore_cache):
     with mock.patch(
         'apache_beam.runners.portability.prism_runner.urlopen') as mock_urlopen:
       mock_response = mock.MagicMock()
-      mock_response.read.return_value = b''
+      if has_cache_zip:
+        with open(self.cache_zip_path, 'rb') as f:
+          mock_response.read.side_effect = [f.read(), b'']
+      else:
+        mock_response.read.return_value = b''
       mock_urlopen.return_value = mock_response
       with mock.patch('zipfile.is_zipfile') as mock_is_zipfile:
         with mock.patch('zipfile.ZipFile') as mock_zipfile_init: