Images sdk changes

setup.py

@@ -18,6 +18,8 @@
         "attrs>=21.2.0",
         "frozendict>=1.2",
         "google-auth>=1.31.0",
+        "grpcio>=1.50.0",
+        "grpcio-tools>=1.50.0",
         "mock>=4.0.3",
         "Pillow>=8.3.1",
         "protobuf>=3.19.0",

src/google/appengine/api/images/__init__.py

@@ -33,17 +33,50 @@
 
 
 import json
+import logging
+import os
 import struct
 
+# import google.auth
+# import google.auth.transport.grpc
+# import google.auth.transport.requests
+from google.oauth2 import id_token
+import grpc
 from google.appengine.api import apiproxy_stub_map
 from google.appengine.api import blobstore
 from google.appengine.api import datastore_types
-from google.appengine.api.images import images_service_pb2
+from . import images_service_pb2
+from . import images_service_rpc_pb2_grpc
 from google.appengine.runtime import apiproxy_errors
 
 import six
 from six.moves import range
 
+
+def _make_grpc_call(method_name, request):
+    """Creates an authenticated gRPC channel and makes an RPC call."""
+    target_host = os.environ.get('IMAGES_GRPC_SERVICE_TARGET')
+    target_audience = os.environ.get('IMAGES_GRPC_SERVICE_TARGET_URL')
+
+    credentials, project_id = google.auth.default()
+    auth_req = google.auth.transport.requests.Request()
+    credentials.refresh(auth_req)
+    token = id_token.fetch_id_token(auth_req, target_audience)
+
+    channel_credentials = grpc.ssl_channel_credentials()
+    call_credentials = grpc.access_token_call_credentials(token)
+    composite_credentials = grpc.composite_channel_credentials(
+        channel_credentials,
+        call_credentials,
+    )
+
+    authed_channel = grpc.secure_channel(target_host, composite_credentials)
+    stub = images_service_rpc_pb2_grpc.ImagesServiceStub(authed_channel)
+
+    grpc_method = getattr(stub, method_name)
+    return grpc_method(request)
+
+
 BlobKey = datastore_types.BlobKey
 
 
@@ -851,6 +884,29 @@ def execute_transforms_async(self,
       ValueError: When `transparent_substitution_rgb` is not an integer.
       Error: All other error conditions.
     """
+    env_value = os.environ.get('USE_CUSTOM_IMAGES_GRPC_SERVICE')
+    logging.warning(f"USE_CUSTOM_IMAGES_GRPC_SERVICE raw value: {env_value}")
+    logging.warning(f"USE_CUSTOM_IMAGES_GRPC_SERVICE boolean evaluation: {bool(env_value)}")
+    if env_value:
+        logging.warning('Using custom gRPC image service.')
+        if not self._transforms:
+            raise BadRequestError('Must specify at least one transformation.')
+
+        request = images_service_pb2.ImagesTransformRequest()
+        self._set_imagedata(request.image)
+        for transform in self._transforms:
+            request.transform.add().CopyFrom(transform)
+        request.output.mime_type = output_encoding
+        if quality is not None:
+            request.output.quality = quality
+
+        grpc_response = _make_grpc_call("Transform", request)
+
+        class FakeRpc:
+            def get_result(self):
+                return grpc_response.image.content
+        return FakeRpc()
+
     if output_encoding not in OUTPUT_ENCODING_TYPES:
       raise BadRequestError("Output encoding type not in recognized set "
                             "%s" % OUTPUT_ENCODING_TYPES)
@@ -983,6 +1039,19 @@ def histogram_async(self, rpc=None):
       LargeImageError: When the image data supplied is too large to process.
       Error: All other error conditions.
     """
+    if os.environ.get('USE_CUSTOM_IMAGES_GRPC_SERVICE'):
+        logging.warning('Using custom gRPC image service for histogram.')
+        request = images_service_pb2.ImagesHistogramRequest()
+        self._set_imagedata(request.image)
+
+        grpc_response = _make_grpc_call("Histogram", request)
+
+        class FakeRpc:
+            def get_result(self):
+                histogram_ = grpc_response.histogram
+                return [histogram_.red, histogram_.green, histogram_.blue]
+        return FakeRpc()
+
     request = images_service_pb2.ImagesHistogramRequest()
     response = images_service_pb2.ImagesHistogramResponse()
 
@@ -1733,7 +1802,6 @@ def composite_async(inputs,
   image_map = {}
 
   request = images_service_pb2.ImagesCompositeRequest()
-  response = images_service_pb2.ImagesTransformResponse()
   for (image, x, y, opacity, anchor) in inputs:
     if not image:
       raise BadRequestError("Each input must include an image")
@@ -1775,26 +1843,36 @@ def composite_async(inputs,
       (quality is not None)):
     request.canvas.output.quality = quality
 
-  def composite_hook(rpc):
-    """Checks success, handles exceptions, and returns the converted RPC result.
-
-    Args:
-      rpc: A UserRPC object.
-
-    Returns:
-      Images bytes of the composite image.
+  if os.environ.get('USE_CUSTOM_IMAGES_GRPC_SERVICE'):
+      logging.warning('Using custom gRPC image service for composite.')
+      grpc_response = _make_grpc_call("Composite", request)
 
-    Raises:
-      See `composite_async` for more details.
-    """
-    try:
-      rpc.check_success()
-    except apiproxy_errors.ApplicationError as e:
-      raise _ToImagesError(e)
-    return rpc.response.image.content
-
-  return _make_async_call(rpc, "Composite", request, response, composite_hook,
-                          None)
+      class FakeRpc:
+          def get_result(self):
+              return grpc_response.image.content
+      return FakeRpc()
+  else:
+      response = images_service_pb2.ImagesTransformResponse()
+      def composite_hook(rpc):
+        """Checks success, handles exceptions, and returns the converted RPC result.
+
+        Args:
+          rpc: A UserRPC object.
+
+        Returns:
+          Images bytes of the composite image.
+
+        Raises:
+          See `composite_async` for more details.
+        """
+        try:
+          rpc.check_success()
+        except apiproxy_errors.ApplicationError as e:
+          raise _ToImagesError(e)
+        return rpc.response.image.content
+
+      return _make_async_call(rpc, "Composite", request, response, composite_hook,
+                              None)
 
 
 def histogram(image_data, rpc=None):