From 4c45839bc0d734819452e1a73107ad56438dc779 Mon Sep 17 00:00:00 2001 From: dk Date: Thu, 22 Aug 2024 20:29:30 +0700 Subject: [PATCH 1/2] [syft/notebooks] creating `MockBigQueryClient` so that we can run the bigquery API endpoints without using the real API Co-authored-by: Shubham Gupta --- .../scenarios/bigquery/02-configure-api.ipynb | 80 ++++++++++++------- packages/syft/src/syft/util/util.py | 45 +++++++++++ 2 files changed, 96 insertions(+), 29 deletions(-) diff --git a/notebooks/scenarios/bigquery/02-configure-api.ipynb b/notebooks/scenarios/bigquery/02-configure-api.ipynb index c1400401d43..3351c28597d 100644 --- a/notebooks/scenarios/bigquery/02-configure-api.ipynb +++ b/notebooks/scenarios/bigquery/02-configure-api.ipynb @@ -102,7 +102,7 @@ "source": [ "@sy.api_endpoint_method(\n", " settings={\n", - " \"credentials\": secrets[\"service_account_bigquery_private\"],\n", + " \"credentials\": {\"mock\": \"timeout\"}, # success, bigquery_error, timeout\n", " \"region\": secrets[\"region_bigquery\"],\n", " \"project_id\": secrets[\"project_id\"],\n", " }\n", @@ -111,21 +111,29 @@ " context,\n", " sql_query: str,\n", ") -> str:\n", - " # third party\n", - " from google.cloud import bigquery\n", - " from google.oauth2 import service_account\n", + " if \"mock\" in context.settings[\"credentials\"]:\n", + " # syft absolute\n", + " from syft.util.util import MockBigQuery as bigquery\n", + "\n", + " scoped_credentials = bigquery.mock_credentials(\n", + " mock_result=context.settings[\"credentials\"][\"mock\"]\n", + " )\n", + " else:\n", + " # third party\n", + " from google.cloud import bigquery\n", + " from google.oauth2 import service_account\n", + "\n", + " # Auth for Bigquer based on the workload identity\n", + " credentials = service_account.Credentials.from_service_account_info(\n", + " context.settings[\"credentials\"]\n", + " )\n", + " scoped_credentials = credentials.with_scopes(\n", + " [\"https://www.googleapis.com/auth/cloud-platform\"]\n", + " )\n", "\n", " # syft absolute\n", " from syft.service.response import SyftError\n", "\n", - " # Auth for Bigquer based on the workload identity\n", - " credentials = service_account.Credentials.from_service_account_info(\n", - " context.settings[\"credentials\"]\n", - " )\n", - " scoped_credentials = credentials.with_scopes(\n", - " [\"https://www.googleapis.com/auth/cloud-platform\"]\n", - " )\n", - "\n", " client = bigquery.Client(\n", " credentials=scoped_credentials,\n", " location=context.settings[\"region\"],\n", @@ -171,11 +179,13 @@ " \"timeout\",\n", " ]:\n", " return SyftError(\n", - " message=\"Error occured during the call: \" + e._errors[0][\"message\"]\n", + " message=\"A bigquerry error occured during the call: \"\n", + " + e._errors[0][\"message\"]\n", " )\n", " else:\n", " return SyftError(\n", - " message=\"An error occured executing the API call, please contact the domain owner.\"\n", + " message=f\"A bigquerry error with reason '{e._errors[0][\"reason\"]}' \"\n", + " f\"occured executing the API call, please contact the domain owner.\"\n", " )" ] }, @@ -211,7 +221,9 @@ "\n", "@sy.api_endpoint_method(\n", " settings={\n", - " \"credentials\": secrets[\"service_account_bigquery_private\"],\n", + " \"credentials\": {\n", + " \"mock\": \"bigquery_error\"\n", + " }, # success, bigquery_error, error, timeout\n", " \"region\": secrets[\"region_bigquery\"],\n", " \"project_id\": secrets[\"project_id\"],\n", " \"CALLS_PER_MIN\": 10,\n", @@ -225,21 +237,29 @@ " # stdlib\n", " import datetime\n", "\n", - " # third party\n", - " from google.cloud import bigquery\n", - " from google.oauth2 import service_account\n", + " if \"mock\" in context.settings[\"credentials\"]:\n", + " # syft absolute\n", + " from syft.util.util import MockBigQuery as bigquery\n", + "\n", + " scoped_credentials = bigquery.mock_credentials(\n", + " mock_result=context.settings[\"credentials\"][\"mock\"]\n", + " )\n", + " else:\n", + " # third party\n", + " from google.cloud import bigquery\n", + " from google.oauth2 import service_account\n", + "\n", + " # Auth for Bigquer based on the workload identity\n", + " credentials = service_account.Credentials.from_service_account_info(\n", + " context.settings[\"credentials\"]\n", + " )\n", + " scoped_credentials = credentials.with_scopes(\n", + " [\"https://www.googleapis.com/auth/cloud-platform\"]\n", + " )\n", "\n", " # syft absolute\n", " from syft.service.response import SyftError\n", "\n", - " # Auth for Bigquer based on the workload identity\n", - " credentials = service_account.Credentials.from_service_account_info(\n", - " context.settings[\"credentials\"]\n", - " )\n", - " scoped_credentials = credentials.with_scopes(\n", - " [\"https://www.googleapis.com/auth/cloud-platform\"]\n", - " )\n", - "\n", " client = bigquery.Client(\n", " credentials=scoped_credentials,\n", " location=context.settings[\"region\"],\n", @@ -296,11 +316,13 @@ " \"timeout\",\n", " ]:\n", " return SyftError(\n", - " message=\"Error occured during the call: \" + e._errors[0][\"message\"]\n", + " message=\"A bigquerry error occured during the call: \"\n", + " + e._errors[0][\"message\"]\n", " )\n", " else:\n", " return SyftError(\n", - " message=\"An error occured executing the API call, please contact the domain owner.\"\n", + " message=f\"A bigquerry error with reason '{e._errors[0][\"reason\"]}' \"\n", + " f\"occured executing the API call, please contact the domain owner.\"\n", " )" ] }, @@ -673,7 +695,7 @@ "name": "python", "nbconvert_exporter": "python", "pygments_lexer": "ipython3", - "version": "3.12.2" + "version": "3.12.4" } }, "nbformat": 4, diff --git a/packages/syft/src/syft/util/util.py b/packages/syft/src/syft/util/util.py index c9011c82ffc..88f97140985 100644 --- a/packages/syft/src/syft/util/util.py +++ b/packages/syft/src/syft/util/util.py @@ -38,6 +38,7 @@ import types from types import ModuleType from typing import Any +from unittest.mock import Mock # third party from IPython.display import display @@ -1111,3 +1112,47 @@ def repr_truncation(obj: Any, max_elements: int = 10) -> str: r.maxother = 100 # For other objects return r.repr(obj) + + +class MockBigQueryError(Exception): + def __init__(self, errors: list) -> None: + self._errors = errors + super().__init__(self._errors[0]["message"]) + + +class MockBigQueryClient: + def __init__(self, credentials: dict, location: str | None = None) -> None: + self.credentials = credentials + self.location = location + + def query_and_wait(self, sql_query: str, project: str | None = None) -> Mock | None: + if self.credentials["mock_result"] == "timeout": + raise TimeoutError("Simulated query timeout.") + + if self.credentials["mock_result"] == "success": + # Simulate a successful response + rows = Mock() + rows.total_rows = 1 # or any number within acceptable limits + rows.to_dataframe = Mock(return_value="Simulated DataFrame") + return rows + + if self.credentials["mock_result"] == "bigquery_error": + errors = [ + { + "reason": "Simulated BigQuery error.", + "message": "Simulated BigQuery error.", + } + ] + raise MockBigQueryError(errors) + + raise Exception("Simulated non-BigQuery exception.") + + +class MockBigQuery: + @staticmethod + def mock_credentials(mock_result: str = "success") -> dict: + return {"mocked": "credentials", "mock_result": mock_result} + + @staticmethod + def Client(credentials: dict, location: str | None = None) -> MockBigQueryClient: + return MockBigQueryClient(credentials, location) From 1bc610a40b434d1029aa2308df161dce7061aa88 Mon Sep 17 00:00:00 2001 From: dk Date: Tue, 27 Aug 2024 09:50:41 +0700 Subject: [PATCH 2/2] [syft/notebooks] revert second bigquery notebook to sync with dev --- .../scenarios/bigquery/02-configure-api.ipynb | 30 +++++++++++++------ 1 file changed, 21 insertions(+), 9 deletions(-) diff --git a/notebooks/scenarios/bigquery/02-configure-api.ipynb b/notebooks/scenarios/bigquery/02-configure-api.ipynb index 3d5de839d14..40311e68094 100644 --- a/notebooks/scenarios/bigquery/02-configure-api.ipynb +++ b/notebooks/scenarios/bigquery/02-configure-api.ipynb @@ -98,6 +98,14 @@ " # syft absolute\n", " from syft.service.response import SyftError\n", "\n", + " # Auth for Bigquer based on the workload identity\n", + " credentials = service_account.Credentials.from_service_account_info(\n", + " context.settings[\"credentials\"]\n", + " )\n", + " scoped_credentials = credentials.with_scopes(\n", + " [\"https://www.googleapis.com/auth/cloud-platform\"]\n", + " )\n", + "\n", " client = bigquery.Client(\n", " credentials=scoped_credentials,\n", " location=context.settings[\"region\"],\n", @@ -143,13 +151,11 @@ " \"timeout\",\n", " ]:\n", " return SyftError(\n", - " message=\"A bigquerry error occured during the call: \"\n", - " + e._errors[0][\"message\"]\n", + " message=\"Error occured during the call: \" + e._errors[0][\"message\"]\n", " )\n", " else:\n", " return SyftError(\n", - " message=f\"A bigquerry error with reason '{e._errors[0][\"reason\"]}' \"\n", - " f\"occured executing the API call, please contact the domain owner.\"\n", + " message=\"An error occured executing the API call, please contact the domain owner.\"\n", " )" ] }, @@ -206,6 +212,14 @@ " # syft absolute\n", " from syft.service.response import SyftError\n", "\n", + " # Auth for Bigquer based on the workload identity\n", + " credentials = service_account.Credentials.from_service_account_info(\n", + " context.settings[\"credentials\"]\n", + " )\n", + " scoped_credentials = credentials.with_scopes(\n", + " [\"https://www.googleapis.com/auth/cloud-platform\"]\n", + " )\n", + "\n", " client = bigquery.Client(\n", " credentials=scoped_credentials,\n", " location=context.settings[\"region\"],\n", @@ -262,13 +276,11 @@ " \"timeout\",\n", " ]:\n", " return SyftError(\n", - " message=\"A bigquerry error occured during the call: \"\n", - " + e._errors[0][\"message\"]\n", + " message=\"Error occured during the call: \" + e._errors[0][\"message\"]\n", " )\n", " else:\n", " return SyftError(\n", - " message=f\"A bigquerry error with reason '{e._errors[0][\"reason\"]}' \"\n", - " f\"occured executing the API call, please contact the domain owner.\"\n", + " message=\"An error occured executing the API call, please contact the domain owner.\"\n", " )" ] }, @@ -756,7 +768,7 @@ "name": "python", "nbconvert_exporter": "python", "pygments_lexer": "ipython3", - "version": "3.12.3" + "version": "3.12.4" } }, "nbformat": 4,