python - 从 AWS SageMaker 访问 Google BigQuery

Question

在本地运行时，我的 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(或本地)上使用服务帐户或使用环境变量将 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 实例的想法(这意味着我必须将凭据上传到某个私有(private) s3 存储桶，然后将它们存储在 SageMaker 上实例)。不是世界末日，但我宁愿避免这种情况。

有什么想法吗？

Answer 1

正如您所提到的，GCP 目前使用服务帐户、凭据 JSON 和 API token 进行身份验证。您可以考虑使用 AWS Secrets Manager 或 AWS Systems Manager Parameter Store 来存储 GCP 凭证，然后在 Jupyter notebook 中获取它们，而不是将凭证存储在 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 的其他必要路径。我不确定这是否是最好的方法。希望有人有更好的方法。