[Python] Make cloud storage dependencies optional on import in gcsio and s3io

sdks/python/apache_beam/examples/cookbook/group_with_coder_test.py

@@ -32,6 +32,8 @@
 # Protect against environments where gcsio library is not available.
 try:
   from apache_beam.io.gcp import gcsio
+  if not gcsio.GCS_INSTALLED:
+    gcsio = None
 except ImportError:
   gcsio = None
 

sdks/python/apache_beam/examples/snippets/transforms/elementwise/runinference_test.py

@@ -39,7 +39,10 @@
 
 # pylint: disable=wrong-import-order, wrong-import-position, ungrouped-imports, unused-import
 try:
+  from apache_beam.io.gcp.gcsio import GCS_INSTALLED
   from apache_beam.io.gcp.gcsfilesystem import GCSFileSystem
+  if not GCS_INSTALLED:
+    raise unittest.SkipTest('GCP dependencies are not installed')
 except ImportError:
   raise unittest.SkipTest('GCP dependencies are not installed')
 

sdks/python/apache_beam/io/aws/s3io.py

@@ -36,7 +36,8 @@
 
 try:
   # pylint: disable=wrong-import-order, wrong-import-position
-  # pylint: disable=ungrouped-imports
+  # pylint: disable=ungrouped-imports, unused-import
+  import boto3
   from apache_beam.io.aws.clients.s3 import boto3_client
   BOTO3_INSTALLED = True
 except ImportError:

sdks/python/apache_beam/io/external/xlang_bigqueryio_it_test.py

@@ -44,7 +44,10 @@
 # pylint: disable=wrong-import-order, wrong-import-position
 
 try:
+  from apache_beam.io.gcp.gcsio import GCS_INSTALLED
   from apache_beam.io.gcp.gcsio import GcsIO
+  if not GCS_INSTALLED:
+    GcsIO = None
 except ImportError:
   GcsIO = None
 

sdks/python/apache_beam/io/external/xlang_snowflakeio_it_test.py

@@ -61,7 +61,10 @@
 
 # pylint: disable=wrong-import-order, wrong-import-position, ungrouped-imports
 try:
+  from apache_beam.io.gcp.gcsio import GCS_INSTALLED
   from apache_beam.io.gcp.gcsfilesystem import GCSFileSystem
+  if not GCS_INSTALLED:
+    GCSFileSystem = None
 except ImportError:
   GCSFileSystem = None
 # pylint: enable=wrong-import-order, wrong-import-position, ungrouped-imports

sdks/python/apache_beam/io/filesystems_test.py

@@ -30,8 +30,12 @@
 import mock
 
 from apache_beam.io import localfilesystem
+from apache_beam.io.aws import s3filesystem
+from apache_beam.io.aws import s3io
 from apache_beam.io.filesystem import BeamIOError
 from apache_beam.io.filesystems import FileSystems
+from apache_beam.io.gcp import gcsfilesystem
+from apache_beam.io.gcp import gcsio
 
 
 def _gen_fake_join(separator):
@@ -63,9 +67,27 @@ def test_get_filesystem(self):
         isinstance(
             FileSystems.get_filesystem('c:\\abc\\def'),
             localfilesystem.LocalFileSystem))
+    self.assertTrue(
+        isinstance(
+            FileSystems.get_filesystem('gs://bucket/file'),
+            gcsfilesystem.GCSFileSystem))
+    self.assertTrue(
+        isinstance(
+            FileSystems.get_filesystem('s3://bucket/file'),
+            s3filesystem.S3FileSystem))
     with self.assertRaises(ValueError):
       FileSystems.get_filesystem('error://abc/def')
 
+  def test_missing_cloud_dependencies_raise_runtime_error(self):
+    with mock.patch.object(gcsio, 'GCS_INSTALLED', False):
+      with self.assertRaisesRegex(RuntimeError,
+                                  'GCP dependencies are not installed'):
+        gcsio.GcsIO()
+    with mock.patch.object(s3io, 'BOTO3_INSTALLED', False):
+      with self.assertRaisesRegex(RuntimeError,
+                                  'AWS dependencies are not installed'):
+        s3io.S3IO(options={})
+
   @mock.patch('apache_beam.io.localfilesystem.os')
   def test_unix_path_join(self, *unused_mocks):
     # Test joining of Unix paths.

sdks/python/apache_beam/io/gcp/gcsfilesystem_integration_test.py

@@ -39,8 +39,9 @@
 from apache_beam.testing.test_pipeline import TestPipeline
 
 try:
+  from apache_beam.io.gcp.gcsio import GCS_INSTALLED
   from apache_beam.io.gcp.gcsfilesystem import GCSFileSystem
-  fs_not_available = False
+  fs_not_available = not GCS_INSTALLED
 except ImportError:
   fs_not_available = True  # type: ignore
 

sdks/python/apache_beam/io/gcp/gcsio.py

@@ -35,22 +35,33 @@
 from typing import Optional
 from typing import Union
 
-from google.api_core.exceptions import Conflict
-from google.api_core.exceptions import RetryError
-from google.cloud import storage
-from google.cloud.exceptions import NotFound
-from google.cloud.exceptions import from_http_response
-from google.cloud.storage.fileio import BlobReader
-from google.cloud.storage.fileio import BlobWriter
-from google.cloud.storage.retry import DEFAULT_RETRY
-
 from apache_beam import version as beam_version
 from apache_beam.internal.gcp import auth
-from apache_beam.io.gcp import gcsio_retry
 from apache_beam.metrics.metric import Metrics
 from apache_beam.options.pipeline_options import GoogleCloudOptions
 from apache_beam.options.pipeline_options import PipelineOptions
 
+try:
+  # pylint: disable=wrong-import-order, wrong-import-position
+  # pylint: disable=ungrouped-imports
+  from google.api_core.exceptions import Conflict
+  from google.api_core.exceptions import RetryError
+  from google.cloud import storage
+  from google.cloud.exceptions import NotFound
+  from google.cloud.exceptions import from_http_response
+  from google.cloud.storage.fileio import BlobReader
+  from google.cloud.storage.fileio import BlobWriter
+  from google.cloud.storage.retry import DEFAULT_RETRY
+
+  from apache_beam.io.gcp import gcsio_retry
+  GCS_INSTALLED = True
+except ImportError:
+  GCS_INSTALLED = False
+  storage = None  # type: ignore
+  BlobReader = object  # type: ignore
+  BlobWriter = object  # type: ignore
+  DEFAULT_RETRY = None  # type: ignore
+
 __all__ = ['GcsIO', 'create_storage_client']
 
 _LOGGER = logging.getLogger(__name__)
@@ -202,13 +213,18 @@ class GcsIO(object):
   """Google Cloud Storage I/O client."""
   def __init__(
       self,
-      storage_client: Optional[storage.Client] = None,
+      storage_client: Optional['storage.Client'] = None,
       pipeline_options: Optional[Union[dict, PipelineOptions]] = None) -> None:
     if pipeline_options is None:
       pipeline_options = PipelineOptions()
     elif isinstance(pipeline_options, dict):
       pipeline_options = PipelineOptions.from_dictionary(pipeline_options)
     if storage_client is None:
+      if not GCS_INSTALLED:
+        message = (
+            'GCP dependencies are not installed, and no alternative '
+            'client was provided to GcsIO.')
+        raise RuntimeError(message)
       storage_client = create_storage_client(pipeline_options)
 
     google_cloud_options = pipeline_options.view_as(GoogleCloudOptions)
@@ -220,7 +236,8 @@ def __init__(
     self.client = storage_client
     self._rewrite_cb = None
     self.bucket_to_project_number = {}
-    self._storage_client_retry = gcsio_retry.get_retry(pipeline_options)
+    self._storage_client_retry = (
+        gcsio_retry.get_retry(pipeline_options) if GCS_INSTALLED else None)
     self._use_blob_generation = getattr(
         google_cloud_options, 'enable_gcsio_blob_generation', False)
 

sdks/python/apache_beam/io/gcp/gcsio_integration_test.py

@@ -52,6 +52,8 @@
 
 try:
   from apache_beam.io.gcp import gcsio
+  if not gcsio.GCS_INSTALLED:
+    gcsio = None  # type: ignore
 except ImportError:
   gcsio = None  # type: ignore
 

sdks/python/apache_beam/ml/inference/onnx_inference_test.py

@@ -55,7 +55,10 @@
   raise unittest.SkipTest('Onnx dependencies are not installed')
 
 try:
+  from apache_beam.io.gcp.gcsio import GCS_INSTALLED
   from apache_beam.io.gcp.gcsfilesystem import GCSFileSystem
+  if not GCS_INSTALLED:
+    GCSFileSystem = None  # type: ignore
 except ImportError:
   GCSFileSystem = None  # type: ignore
 

sdks/python/apache_beam/ml/inference/pytorch_inference_test.py

@@ -49,7 +49,10 @@
   raise unittest.SkipTest('PyTorch dependencies are not installed')
 
 try:
+  from apache_beam.io.gcp.gcsio import GCS_INSTALLED
   from apache_beam.io.gcp.gcsfilesystem import GCSFileSystem
+  if not GCS_INSTALLED:
+    GCSFileSystem = None  # type: ignore
 except ImportError:
   GCSFileSystem = None  # type: ignore
 

sdks/python/apache_beam/ml/inference/sklearn_inference_it_test.py

@@ -32,7 +32,10 @@
 
 # pylint: disable=wrong-import-order, wrong-import-position, ungrouped-imports, unused-import
 try:
+  from apache_beam.io.gcp.gcsio import GCS_INSTALLED
   from apache_beam.io.gcp.gcsfilesystem import GCSFileSystem
+  if not GCS_INSTALLED:
+    raise unittest.SkipTest('GCP dependencies are not installed')
 except ImportError:
   raise unittest.SkipTest('GCP dependencies are not installed')
 

sdks/python/apache_beam/ml/inference/tensorrt_inference_test.py

@@ -45,7 +45,10 @@
   raise unittest.SkipTest('TensorRT dependencies are not installed')
 
 try:
+  from apache_beam.io.gcp.gcsio import GCS_INSTALLED
   from apache_beam.io.gcp.gcsfilesystem import GCSFileSystem
+  if not GCS_INSTALLED:
+    GCSFileSystem = None  # type: ignore
 except ImportError:
   GCSFileSystem = None  # type: ignore
 

sdks/python/apache_beam/options/pipeline_options.py

@@ -1212,6 +1212,11 @@ def _create_default_gcs_bucket(self):
     except ImportError:
       _LOGGER.warning('Unable to create default GCS bucket.')
       return None
+    if not gcsio.GCS_INSTALLED:
+      _LOGGER.warning(
+          'Unable to create default GCS bucket because GCP dependencies are '
+          'not installed.')
+      return None
     bucket = gcsio.get_or_create_default_gcs_bucket(self)
     if bucket:
       return 'gs://%s/' % bucket.id
@@ -1229,6 +1234,11 @@ def _warn_if_soft_delete_policy_enabled(self, arg_name):
     gcs_path = getattr(self, arg_name, None)
     try:
       from apache_beam.io.gcp import gcsio
+      if not gcsio.GCS_INSTALLED:
+        _LOGGER.warning(
+            'Unable to check soft delete policy because GCP dependencies are '
+            'not installed.')
+        return
       if gcsio.GcsIO().is_soft_delete_enabled(gcs_path):
         logger.log_first_n(
             logging.WARN,

sdks/python/apache_beam/options/pipeline_options_test.py

@@ -47,8 +47,8 @@
 _LOGGER = logging.getLogger(__name__)
 
 try:
-  import apache_beam.io.gcp.gcsio  # pylint: disable=unused-import
-  has_gcsio = True
+  from apache_beam.io.gcp.gcsio import GCS_INSTALLED
+  has_gcsio = GCS_INSTALLED
 except ImportError:
   has_gcsio = False
 

sdks/python/apache_beam/runners/render.py

@@ -75,6 +75,8 @@
 
 try:
   from apache_beam.io.gcp import gcsio
+  if not gcsio.GCS_INSTALLED:
+    gcsio = None  # type: ignore
 except ImportError:
   gcsio = None  # type: ignore