From f67073b3d4d0b2c67fe4d9e8db75c8dbe68e4190 Mon Sep 17 00:00:00 2001 From: Marigold Date: Fri, 13 Dec 2024 10:31:41 +0100 Subject: [PATCH] =?UTF-8?q?=E2=9C=A8=20Add=20support=20for=20Google=20Clou?= =?UTF-8?q?d=20service=20account=20authentication=20in=20BigQuery=20read?= =?UTF-8?q?=20function?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- apps/wizard/app_pages/producer_analytics.py | 14 +++++++++++++- etl/config.py | 3 +++ 2 files changed, 16 insertions(+), 1 deletion(-) diff --git a/apps/wizard/app_pages/producer_analytics.py b/apps/wizard/app_pages/producer_analytics.py index 3be62be9165..a6a76707329 100644 --- a/apps/wizard/app_pages/producer_analytics.py +++ b/apps/wizard/app_pages/producer_analytics.py @@ -3,14 +3,16 @@ import owid.catalog.processing as pr import pandas as pd +import pandas_gbq import plotly.express as px import streamlit as st -from pandas_gbq import read_gbq +from google.oauth2 import service_account from st_aggrid import AgGrid, GridUpdateMode, JsCode from st_aggrid.grid_options_builder import GridOptionsBuilder from structlog import get_logger from apps.wizard.utils.components import st_horizontal +from etl.config import GOOGLE_APPLICATION_CREDENTIALS from etl.snapshot import Snapshot from etl.version_tracker import VersionTracker @@ -70,6 +72,16 @@ def columns_producer(min_date, max_date): return cols_prod +def read_gbq(*args, **kwargs) -> pd.DataFrame: + if GOOGLE_APPLICATION_CREDENTIALS: + # Use service account + credentials = service_account.Credentials.from_service_account_file(GOOGLE_APPLICATION_CREDENTIALS) + return pandas_gbq.read_gbq(*args, **kwargs, credentials=credentials) # type: ignore + else: + # Use browser authentication. + return pandas_gbq.read_gbq(*args, **kwargs) # type: ignore + + @st.cache_data(show_spinner=False) def get_grapher_views( date_start: str = MIN_DATE.strftime("%Y-%m-%d"), diff --git a/etl/config.py b/etl/config.py index 59edff4129c..c41382cd8e5 100644 --- a/etl/config.py +++ b/etl/config.py @@ -225,6 +225,9 @@ def variable_metadata_url(variable_id): # Default schema for presentation.grapher_config in metadata. Try to keep it up to date with the latest schema. DEFAULT_GRAPHER_SCHEMA = "https://files.ourworldindata.org/schemas/grapher-schema.006.json" +# Google Cloud service account path (used for BigQuery) +GOOGLE_APPLICATION_CREDENTIALS = env.get("GOOGLE_APPLICATION_CREDENTIALS") + def enable_bugsnag() -> None: if BUGSNAG_API_KEY: