Update BigQuery remote function for vision AI to VPCSC compliant and …

…latest version (#8883) Co-authored-by: Karl Weinmeister <[email protected]> Co-authored-by: Charles Engelke <[email protected]>
GoogleCloudPlatform · Mar 13, 2023 · 7a29c8e · 7a29c8e
1 parent f95bc3a
commit 7a29c8e
Show file tree

Hide file tree

Showing 4 changed files with 28 additions and 21 deletions.
diff --git a/bigquery/remote-function/vision/requirements-test.txt b/bigquery/remote-function/vision/requirements-test.txt
@@ -1,4 +1,4 @@
 Flask==2.2.2
-functions-framework==3.2.0
-google-cloud-vision==3.1.2
-pytest==7.1.3
+functions-framework==3.3.0
+google-cloud-vision==3.2.0
+pytest==7.2.0
diff --git a/bigquery/remote-function/vision/requirements.txt b/bigquery/remote-function/vision/requirements.txt
@@ -1,3 +1,3 @@
 Flask==2.2.2
-functions-framework==3.2.0
-google-cloud-vision==3.1.2
+functions-framework==3.3.0
+google-cloud-vision==3.2.0
diff --git a/bigquery/remote-function/vision/vision_function.py b/bigquery/remote-function/vision/vision_function.py
@@ -13,9 +13,11 @@
 # limitations under the License.
 
 # [START bigquery_remote_function_vision]
+import urllib.request
+
 import flask
 import functions_framework
-from google.cloud import vision_v1
+from google.cloud import vision
 
 
 @functions_framework.http
@@ -29,13 +31,13 @@ def label_detection(request: flask.Request) -> flask.Response:
         https://cloud.google.com/bigquery/docs/reference/standard-sql/remote-functions#output_format
     """
     try:
-        client = vision_v1.ImageAnnotatorClient()
+        client = vision.ImageAnnotatorClient()
         calls = request.get_json()['calls']
         replies = []
         for call in calls:
-            results = client.label_detection(
-                {'source': {'image_uri': call[0]}})
-            replies.append(vision_v1.AnnotateImageResponse.to_dict(results))
+            content = urllib.request.urlopen(call[0]).read()
+            results = client.label_detection({'content': content})
+            replies.append(vision.AnnotateImageResponse.to_dict(results))
         return flask.make_response(flask.jsonify({'replies': replies}))
     except Exception as e:
         return flask.make_response(flask.jsonify({'errorMessage': str(e)}), 400)

diff --git a/bigquery/remote-function/vision/vision_function_test.py b/bigquery/remote-function/vision/vision_function_test.py
@@ -15,7 +15,7 @@
 from unittest import mock
 
 import flask
-from google.cloud import vision_v1
+from google.cloud import vision
 import pytest
 
 import vision_function
@@ -27,16 +27,19 @@ def app() -> flask.Flask:
     return flask.Flask(__name__)
 
 
-@mock.patch('vision_function.vision_v1')
-def test_vision_function(mock_vision_v1: object, app: flask.Flask) -> None:
+@mock.patch('vision_function.urllib.request')
+@mock.patch('vision_function.vision')
+def test_vision_function(mock_vision: object, mock_request: object,
+                         app: flask.Flask) -> None:
+    mock_request.urlopen = mock.Mock(read=mock.Mock(return_value=b'filedata'))
     label_detection_mock = mock.Mock(side_effect=[
-        vision_v1.AnnotateImageResponse(
+        vision.AnnotateImageResponse(
             {'label_annotations': [{'description': 'apple'}]}),
-        vision_v1.AnnotateImageResponse(
+        vision.AnnotateImageResponse(
             {'label_annotations': [{'description': 'banana'}]})])
-    mock_vision_v1.ImageAnnotatorClient = mock.Mock(
+    mock_vision.ImageAnnotatorClient = mock.Mock(
         return_value=mock.Mock(label_detection=label_detection_mock))
-    mock_vision_v1.AnnotateImageResponse = vision_v1.AnnotateImageResponse
+    mock_vision.AnnotateImageResponse = vision.AnnotateImageResponse
     with app.test_request_context(
             json={'calls': [['https://storage.googleapis.com/bucket/apple'],
                             ['https://storage.googleapis.com/bucket/banana']]}):
@@ -47,11 +50,13 @@ def test_vision_function(mock_vision_v1: object, app: flask.Flask) -> None:
         assert 'banana' in str(response.get_json()['replies'][1])
 
 
-@mock.patch('vision_function.vision_v1')
-def test_vision_function_error(
-        mock_vision_v1: object, app: flask.Flask) -> None:
+@mock.patch('vision_function.urllib.request')
+@mock.patch('vision_function.vision')
+def test_vision_function_error(mock_vision: object, mock_request: object,
+                               app: flask.Flask) -> None:
+    mock_request.urlopen = mock.Mock(read=mock.Mock(return_value=b'filedata'))
     label_detection_mock = mock.Mock(side_effect=Exception('API error'))
-    mock_vision_v1.ImageAnnotatorClient = mock.Mock(
+    mock_vision.ImageAnnotatorClient = mock.Mock(
         return_value=mock.Mock(label_detection=label_detection_mock))
     with app.test_request_context(
             json={'calls': [['https://storage.googleapis.com/bucket/apple'],