Automatiser les réservations et affectations BigQuery avec Dataform

L'un de nos clients mène un projet de R&D sur BigQuery, exploité à la fois par les équipes Développement et Data Science. Leur usage présente les caractéristiques suivantes :

1. Les requêtes interactives ont lieu principalement en semaine, durant les heures ouvrables.
2. Forte concurrence des requêtes et temps de réponse rapides.
3. Certaines tâches sont planifiées la nuit ou le week-end.

Pour répondre à ces besoins, le projet s'appuiera sur le plan On-Demand durant les heures ouvrables, afin de garantir une faible latence et une forte concurrence des requêtes. En dehors de ces plages, il basculera sur l'Enterprise Edition avec un maximum de 100 slots.

Un flux d'automatisation modifiera le plan tarifaire en conséquence.

À ce jour, Google ne propose aucune solution clé en main pour basculer automatiquement entre les modes de tarification BigQuery. En revanche, une interface permet de gérer les réservations en CLI ou en SQL.

Nous allons utiliser Dataform et son langage SQLX pour fluidifier la mise à jour des plans tarifaires. Cette approche permet de gérer toutes les modifications de réservations et d'affectations au sein d'un workflow Dataform, lui-même versionné dans Git.

Autre atout de Dataform : sa simplicité. Tout se passe en SQL, directement dans BigQuery, sans service externe. Dataform est gratuit et n'engendre aucun coût supplémentaire.

En préambule, nous créerons (manuellement) une réservation Enterprise Edition nommée test_rd, configurée pour un maximum de 100 slots.

Pour gérer ensuite la configuration des réservations et leurs affectations, nous utiliserons l'API BigQuery Reservation (SQL DDL).

Exemple de code SQL pour créer une affectation dans la région US à partir de la réservation ndv-rd, avec l'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 Dataform vous est encore peu familier, je vous recommande la lecture du billet Build SQL pipelines to BigQuery with Dataform ainsi que le Quick start.

L'ensemble du code associé à ce billet est disponible sur ce dépôt GitHub.

La mise en place d'un workflow automatisé se déroule en six étapes :

[1] Créer un dépôt et attribuer les rôles

Pour créer un dépôt Dataform dans le projet d'administration BigQuery, rendez-vous dans la section Dataform et lancez la création du dépôt (Créer un dépôt Dataform).

Pour garantir une exécution sans accroc, le développeur (utilisateur principal) ainsi que le SA de Dataform (compte de service) doivent disposer des rôles suivants pendant l'exécution du workflow : bigquery.resourceAdmin, bigquery.user.

Pour attribuer les rôles au SA (code ci-dessous), utilisez l'ID du SA affiché dans le message qui apparaît après la création du dépôt.

Le 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] Créer un workspace de développement

Pour créer un nouveau workspace de développement Dataform, suivez ces étapes : Créer un workspace Dataform.

Une fois initialisé, le workspace contient plusieurs exemples de méthodes SQLX et un fichier de configuration nommé dataform.json. Les fichiers d'exemple (first_view.sqlx, second_view.sqlx) peuvent être supprimés.

Pour configurer les méthodes Dataform, ajoutez la section vars contenant la liste de paramètres ci-dessous. Cette section doit apparaître après le dernier paramètre existant et inclure les paramètres indiqués. Vous y définissez les valeurs par défaut, qui pourront être surchargées plus tard dans le workflow.

1. Pensez à ajouter une virgule "," avant la section vars.
2. Remplacez les paramètres marqués par "<>" (les "<>" doivent être supprimés).

"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>"
}

Voici un exemple de notre implémentation :

"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] Développer les méthodes SQLX

À cette étape, nous allons écrire le code SQLX qui prendra en charge les réservations et les affectations.

• Créez un fichier nommé drop_and_create_assignment.sqlx dans le dossier 'definitions'.
• Ajoutez le code SQLX suivant pour supprimer toutes les affectations du projet et créer celle qui vous est nécessaire.
• Vérifiez la bonne compilation grâce au ✅ :

Voici le code 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;

Et voici le code généré attendu d'après notre configuration :

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] Tester et vérifier les méthodes SQLX

Avant de lancer le test, notez que la configuration des projets affectés sera modifiée selon la valeur définie dans le fichier dataform.json. Mieux vaut tester d'abord sur un projet de test.

Pour exécuter et tester le code, deux options s'offrent à vous : utiliser le bouton Run avec vos identifiants, ou choisir Start Execution pour lancer l'exécution avec le SA de Dataform, comme indiqué dans le message ci-dessous.

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

Vérifiez la sortie d'exécution : This statement created a new assignment, puis confirmez les changements dans la vue Capacity management.

La réservation est affectée au projet

Une fois la vérification du SA réussie, committez les changements et poussez-les sur la branche par défaut. Votre développement est alors terminé et prêt pour la mise en production.

[5] Créer une release configuration Dataform

Pour produire un instantané du code et l'exécuter avec des paramètres surchargés, utilisez la release configuration. Pour rester simple, nous ne configurerons qu'un seul paramètre et fixerons le refresh sur Never.

Pensez à recompiler après chaque modification.

Le projet R&D disposera de deux releases : l'une pour le mode On-Demand avec reservation_name= `none``` et une autre pour l'Enterprise edition avec reservation_name=ndv-rd``` (attention aux ``""``).

Testez la release configuration avec l'action sélectionnée :

ndv-playground-bq-mgmt.dataform.drop_and_create_assignment

Release pour la réservation Enterprise

Release pour la réservation On-demand

[6] Créer la workflow configuration Dataform

La workflow configuration définit la manière dont les actions SQLX sont exécutées automatiquement. Chaque workflow s'appuie sur une release configuration et ses paramètres. Les actions sélectionnées sont ensuite déclenchées par un trigger spécifique.

Le projet R&D dispose de deux workflows qui s'appuient sur des release configurations distinctes et se déclenchent à des moments différents. L'heure de démarrage est définie au format 'cron'.

Workflow On-demand

Workflow Enterprise edition