从 AWS SageMaker 访问 Google BigQuery

Posted 2023-03-25

【中文标题】从 AWS SageMaker 访问 Google BigQuery

【英文标题】：Accessing Google BigQuery from AWS SageMaker

【发布时间】：2019-04-05 08:52:01

【问题描述】：

在本地运行时，我的 Jupyter 笔记本能够像这样引用 Google BigQuery：

%%bigquery some_bq_table

SELECT *
FROM
  `some_bq_dataset.some_bq_table` 

以便稍后在我的笔记本中我可以引用 some_bq_table 作为 pandas 数据框，如下所示：https://cloud.google.com/bigquery/docs/visualize-jupyter

我想在 AWS SageMaker 上运行我的笔记本来测试一些东西。要使用 BigQuery 进行身份验证，似乎只有两种方法是在 GCP（或本地）上使用服务帐户，或者使用 env var 将 SDK 指向凭据 JSON（如下所述：https://cloud.google.com/docs/authentication/getting-started）。

例如

export GOOGLE_APPLICATION_CREDENTIALS="/home/user/Downloads/[FILE_NAME].json"

是否有从 SageMaker 连接到 bigquery 的简单方法？我现在最好的想法是将 JSON 从某处下载到 SageMaker 实例，然后从 python 代码设置环境变量。

例如，我会这样做：

os.environ["GOOGLE_APPLICATION_CREDENTIALS"] = "/home/user/Downloads/[FILE_NAME].json"

但是，这不是很安全 - 我不喜欢将我的凭据 JSON 下载到 SageMaker 实例的想法（这意味着我必须将凭据上传到一些私有 s3 存储桶，然后将它们存储在 SageMaker实例）。不是世界末日，但我宁愿避免这种情况。

有什么想法吗？

【参考方案1】：

正如您提到的，GCP 目前使用服务帐户、凭据 JSON 和 API 令牌进行身份验证。您可以考虑使用 AWS Secrets Manager 或 AWS Systems Manager Parameter Store 来存储 GCP 凭证，然后在 Jupyter 笔记本中获取它们，而不是将凭证存储在 S3 存储桶中。这样可以保护凭证，并且仅在需要时才从 Secrets Manager 创建凭证文件。

这是我之前用于从 SageMaker 实例连接到 BigQuery 的示例代码。

import os
import json
import boto3
from google.cloud.bigquery import magics
from google.oauth2 import service_account

def get_gcp_credentials_from_ssm(param_name):
    # read credentials from SSM parameter store
    ssm = boto3.client('ssm')
    # Get the requested parameter
    response = ssm.get_parameters(Names=[param_name], WithDecryption=True)
    # Store the credentials in a variable
    gcp_credentials = response['Parameters'][0]['Value']
    # save credentials temporarily to a file
    credentials_file = '/tmp/.gcp/service_credentials.json'
    with open(credentials_file, 'w') as outfile:  
        json.dump(json.loads(gcp_credentials), outfile)
    # create google.auth.credentials.Credentials to use for queries 
    credentials = service_account.Credentials.from_service_account_file(credentials_file)
    # remove temporary file
    if os.path.exists(credentials_file):
        os.remove(credentials_file)
    return credentials

# this will set the context credentials to use for queries performed in jupyter 
# using bigquery cell magic
magics.context.credentials = get_gcp_credentials_from_ssm('my_gcp_credentials')

请注意，SageMaker 执行角色应该有权访问 SSM，当然还有其他必要的路由来连接到 GCP。我不确定这是否是最好的方法。希望有人有更好的方法。