Automatiza reservas y asignaciones de BigQuery con Dataform

Uno de nuestros clientes está ejecutando un proyecto de Investigación y Desarrollo en BigQuery, que utilizan tanto el equipo de Desarrollo como el de Data Science. Los equipos tienen un patrón de uso bien definido, con las siguientes características:

1. Las consultas interactivas se hacen principalmente en días laborables, dentro del horario de oficina.
2. Alta concurrencia de consultas con tiempos de respuesta rápidos.
3. Hay tareas puntuales que se programan para ejecutarse de noche o los fines de semana.

Para cubrir estas necesidades, el proyecto usará el plan ‘On-Demand’ durante el horario laboral, de modo que se logre una latencia de consulta eficiente y alta concurrencia. Fuera del horario laboral, el proyecto usará la ‘Enterprise Edition’ con un máximo de 100 slots.

Un flujo de automatización se encarga de cambiar el plan de precios según corresponda.

Hoy por hoy, Google no ofrece una solución lista para usar que permita alternar de forma automática entre los planes de precios de BigQuery. Lo que sí proporciona es una interfaz para administrar reservas mediante CLI o SQL.

A continuación, vamos a usar Dataform y su lenguaje SQLX para mejorar la eficiencia al actualizar los planes de precios. Con este enfoque se puede gestionar de forma fluida cualquier modificación en reservas y asignaciones dentro de un workflow de Dataform, con todo controlado por versiones en Git.

Otra ventaja de Dataform es su simplicidad: todo el trabajo se hace en SQL dentro de BigQuery, sin necesidad de servicios externos. Dataform es gratuito y no implica ningún costo adicional.

Como paso previo, vamos a crear (manualmente) una reserva ‘Enterprise Edition’ llamada test_rd con una configuración máxima de 100 slots.

Después, para administrar la configuración de la reserva y los cambios en las asignaciones, usaremos la BigQuery Reservation API (SQL DDL).

Ejemplo de código SQL para crear una asignación en la región US, usando una reserva llamada ndv-rd, con assignment_id df_rd.

CREATE ASSIGNMENT `ndv-playground-bq-mgmt.region-us.ndv-rd.df_rd`
OPTIONS( assignee="projects/ndv-playground", job_type="QUERY");

Si todavía no conoces Dataform, te recomiendo leer este blog, Build SQL pipelines to BigQuery with Dataform, y completar el ‘Quick start’.

Puedes acceder a todo el código asociado a este blog desde este repositorio de GitHub.

El proceso para construir un workflow automatizado consta de seis pasos:

[1] Crear un repositorio y otorgar roles

Para crear un repositorio de Dataform en el proyecto admin de BigQuery, dirígete a la sección de Dataform e inicia el proceso de creación. (Create a Dataform repository)

Para asegurar una ejecución correcta, tanto el principal del desarrollador (usuario) como la SA (service account) de Dataform necesitan los siguientes roles durante la ejecución del workflow: bigquery.resourceAdmin, bigquery.user

Para asignar los roles a la SA (código abajo), utiliza el ID de la SA que aparece en el mensaje al crear el repositorio.

La SA de Dataform

gcloud projects add-iam-policy-binding <PROJECT-ID> \
--member="<service-account-dataform>" \
--role=roles/bigquery.resourceAdmin --condition=None
gcloud projects add-iam-policy-binding <PROJECT-ID> \
--member="<service-account-dataform>" \
--role=roles/bigquery.user --condition=None

[2] Crear un workspace de desarrollo

Para crear un nuevo workspace de desarrollo en Dataform, sigue estos pasos: Create a Dataform workspace.

Una vez inicializado, el workspace incluirá varios métodos SQLX de ejemplo y un archivo de configuración llamado "dataform.json". Esos archivos de ejemplo (first_view.sqlx, second_view.sqlx) se pueden eliminar.

Para configurar los métodos de Dataform, agrega la sección "vars", que contiene la lista de parámetros que se muestra a continuación. Esta sección debe ir después del último parámetro existente e incluir los parámetros indicados. Aquí defines los valores por defecto de cada parámetro, que luego podrán sobrescribirse en el workflow.

1. Asegúrate de agregar una coma "," antes de la sección "vars".
2. Reemplaza los parámetros marcados con "<>" (los "<>" deben eliminarse).

"vars": {
"assigned_project":"'<R&D project>'",
"Assignment_id":"`df_rd_assignment`",
"reservation_prefix":"<BQ admin project>.region-<reservation location>",
"reservation_name":"`<R&D reservation name>`",
"region":"<reservation location>"
}

Este es un ejemplo de nuestra implementación:

"vars": {
"assigned_project":"'ndv-playground'",
"Assignment_id":"`df_rd_assignment`",
"reservation_prefix":"`ndv-playground-bq-mgmt.region-us`",
"reservation_name":"`ndv-rd`",
"region":"region-us"
}

[3] Desarrollar los métodos SQLX

En este paso vamos a crear el código SQLX para gestionar reservas y asignaciones.

• Crea un archivo llamado drop_and_create_assignment.sqlx dentro de la carpeta 'definitions'.
• Incluye el siguiente código SQLX para eliminar todas las asignaciones del proyecto y crear la asignación necesaria.
• Confirma que se compile correctamente; debe aparecer un ✅:

Este es el código SQLX de drop_and_create_assignment.sqlx

config {type: "operations"}
BEGIN

CREATE TEMP TABLE to_drop (drop_string STRING) AS (
SELECT
CONCAT("DROP ASSIGNMENT IF EXISTS `",project_id,".",SPLIT(ddl,".")[OFFSET(1)],
       ".",reservation_name,".",assignment_id,"`")
FROM
${dataform.projectConfig.vars.region}.INFORMATION_SCHEMA.ASSIGNMENTS_BY_PROJECT
WHERE
assignee_id = ${dataform.projectConfig.vars.assigned_project}
AND job_type = 'QUERY' );

FOR drop_statement IN (
SELECT drop_string FROM to_drop)
DO EXECUTE IMMEDIATE drop_statement.drop_string;
END FOR;

END;

Este es el código que se debería generar según nuestra configuración:

BEGIN
CREATE TEMP TABLE to_drop (drop_string STRING) AS (
SELECT
CONCAT("DROP ASSIGNMENT IF EXISTS `",project_id,".",SPLIT(ddl,".")[OFFSET(1)],".",reservation_name,".",assignment_id,"`")
FROM
region-us.INFORMATION_SCHEMA.ASSIGNMENTS_BY_PROJECT
WHERE
assignee_id = 'ndv-playground' AND job_type = 'QUERY' );
FOR drop_statement IN (
SELECT drop_string FROM to_drop)
DO EXECUTE IMMEDIATE drop_statement.drop_string;
END FOR;
END;

CREATE ASSIGNMENT
`ndv-playground-bq-mgmt.region-us`.`ndv-rd`.`df_rd_assignment`
OPTIONS(
assignee=CONCAT("projects/",'ndv-playground'),
job_type="QUERY");

[4] Probar y verificar los métodos SQLX

Antes de probar el código, ten en cuenta que la configuración de los proyectos asignados se modificará según el valor del archivo dataform.json. Te recomendamos probarlo en un proyecto de ‘test’.

Para ejecutar y probar el código tienes dos opciones: usar el botón ‘Run’ con tus credenciales o elegir ‘Start Execution’ para ejecutarlo con la SA de Dataform, como se indica en el mensaje a continuación.

Execute all actions, or select a subset of actions. Service account [email protected] will be used.

Verifica con el resultado de la ejecución: This statement created a new assignment, y luego confirma los cambios en la vista ‘Capacity management’.

La reserva queda asignada al proyecto

Una vez que la verificación con la SA sea exitosa, puedes hacer commit de los cambios y enviarlos a la rama por defecto. Llegado este punto, tu desarrollo está completo y listo para release.

[5] Crear la configuración de release de Dataform

Para crear un snapshot del código y ejecutarlo con parámetros sobrescritos, utiliza la release configuration. Para simplificar, vamos a configurar un solo parámetro y dejaremos el refresh en Never.

Por favor, recuerda compilar después de cada cambio.

El proyecto de I+D tendrá dos releases: uno para ‘On-Demand’, con reservation_name= `none``` y otro para ‘Enterprise edition’, reservation_name=ndv-rd``` (presta atención a las ``""``).

Prueba la "release configuration" con la acción seleccionada:

ndv-playground-bq-mgmt.dataform.drop_and_create_assignment

Release para la reserva Enterprise

Release para la reserva On-demand

[6] Crear la configuración del workflow de Dataform

La configuración del workflow define cómo se ejecutan automáticamente las acciones SQLX. Cada workflow se basa en una "release configuration" y sus parámetros. Las acciones seleccionadas se ejecutan a partir de un trigger específico.

El proyecto de I+D tiene dos workflows que utilizan distintas release configurations y se disparan en horarios diferentes. La hora de inicio se define con el formato 'cron'.

Workflow On-demand

Workflow Enterprise edition