Automatizzare prenotazioni e assegnazioni BigQuery con Dataform

Uno dei nostri clienti gestisce un progetto di Ricerca e Sviluppo su BigQuery, utilizzato sia dal team di sviluppo sia da quello di Data Science. I team mostrano un pattern di utilizzo ben preciso, con queste caratteristiche:

1. Le query interattive vengono eseguite soprattutto nei giorni feriali, in orario lavorativo.
2. Elevata concorrenza tra le query, con tempi di risposta rapidi.
3. Alcuni task specifici sono pianificati di notte o nel fine settimana.

Per rispondere a queste esigenze, il progetto utilizzerà il piano 'On-Demand' durante l'orario lavorativo, così da garantire bassa latenza ed elevata concorrenza. Nelle ore non lavorative passerà alla 'Enterprise Edition' con un massimo di 100 slot.

Sarà un flusso di automazione a cambiare il piano tariffario di volta in volta.

Ad oggi Google non offre una soluzione pronta all'uso per passare automaticamente da un piano tariffario di BigQuery a un altro. Mette però a disposizione un'interfaccia per gestire le prenotazioni tramite CLI o SQL.

Useremo quindi Dataform e il suo linguaggio SQLX per rendere più efficiente l'aggiornamento dei piani tariffari. Questo approccio consente di gestire in modo fluido qualsiasi modifica a prenotazioni e assegnazioni all'interno di un workflow Dataform, peraltro versionato su Git.

Un altro vantaggio di Dataform è la semplicità: tutto avviene in SQL all'interno di BigQuery, senza ricorrere a servizi esterni. Dataform è gratuito e non comporta costi aggiuntivi.

Come operazione preliminare creeremo (manualmente) una prenotazione 'Enterprise Edition' chiamata test_rd, configurata con un massimo di 100 slot.

Per gestire la configurazione delle prenotazioni e le modifiche di assegnazione utilizzeremo poi la BigQuery Reservation API (SQL DDL).

Ecco un esempio di codice SQL per creare un'assegnazione nella regione US, usando una prenotazione chiamata ndv-rd e un assignment_id pari a df_rd.

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

Se non ha familiarità con Dataform, le consigliamo di leggere l'articolo Build SQL pipelines to BigQuery with Dataform e di seguire la Quick start.

Tutto il codice citato nell'articolo è disponibile in questo repository GitHub.

I passaggi per costruire un workflow automatizzato sono sei:

[1] Creare un repository e assegnare i ruoli

Per creare un repository Dataform nel progetto admin di BigQuery, vada alla sezione Dataform e avvii la procedura di creazione. (Create a Dataform repository)

Per garantire la corretta esecuzione, sia l'utente principal sviluppatore sia il SA (service account) di Dataform devono disporre dei seguenti ruoli durante l'esecuzione del workflow: bigquery.resourceAdmin, bigquery.user

Per assegnare i ruoli al SA (codice qui sotto), utilizzi l'ID del SA mostrato nel messaggio che compare al termine della creazione del repository.

Il SA di 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] Creare un workspace di sviluppo

Per creare un nuovo workspace di sviluppo Dataform segua questi passaggi: Create a Dataform workspace.

Una volta inizializzato, il workspace conterrà alcuni metodi SQLX di esempio e un file di configurazione chiamato "dataform.json". I file di esempio (first_view.sqlx, second_view.sqlx) possono essere rimossi.

Per configurare i metodi Dataform, aggiunga la sezione "vars" con l'elenco di parametri riportato sotto. Questa sezione va inserita dopo l'ultimo parametro esistente e deve includere i parametri indicati. Qui si impostano i valori di default di ciascun parametro, che potranno essere sovrascritti nel workflow in una fase successiva.

1. Si ricordi di aggiungere una virgola "," prima della sezione "vars".
2. Sostituisca i parametri contrassegnati con "<>" (i caratteri "<>" vanno rimossi).

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

Ecco un esempio della nostra implementazione:

"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] Sviluppare i metodi SQLX

In questo passaggio scriveremo il codice SQLX per gestire prenotazioni e assegnazioni.

• Crei un file chiamato drop_and_create_assignment.sqlx nella cartella 'definitions'.
• Inserisca il codice SQLX seguente per eliminare tutte le assegnazioni del progetto e creare quella necessaria.
• Verifichi che la compilazione sia andata a buon fine controllando la presenza del segno ✅:

Ecco il codice SQLX di 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;

E questo è il codice che, in base alla nostra configurazione, ci aspettiamo venga generato

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] Testare e verificare i metodi SQLX

Prima di lanciare il test, tenga presente che la configurazione dei progetti assegnati verrà modificata in base al valore presente nel file dataform.json. Le suggeriamo di valutare se eseguire il test su un progetto di 'test'.

Per eseguire e testare il codice ha due opzioni: usare il pulsante 'Run' con le proprie credenziali, oppure scegliere 'Start Execution' per eseguirlo con il SA di Dataform, come indicato nel messaggio sottostante.

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

Verifichi l'output dell'esecuzione: This statement created a new assignment e controlli quindi le modifiche nella vista 'Capacity management'.

La prenotazione è assegnata al progetto

Una volta superata la verifica con il SA, può procedere con il commit delle modifiche e il push sul branch predefinito. A questo punto lo sviluppo è completato e tutto è pronto per il rilascio.

[5] Creare la release configuration di Dataform

Per creare uno snapshot del codice ed eseguirlo con parametri sovrascritti utilizzi la release configuration. Per semplicità configureremo un solo parametro e imposteremo il refresh su Never.

Si ricordi di compilare dopo ogni modifica.

Il progetto R&D avrà due release: una per l'"On-Demand", con reservation_name= `none``` e un'altra per la 'Enterprise edition', reservation_name=ndv-rd``` (faccia attenzione ai ``""``).

Testi la "release configuration" con l'azione selezionata:

ndv-playground-bq-mgmt.dataform.drop_and_create_assignment

Release per la prenotazione Enterprise

Release per la prenotazione On-demand

[6] Creare la workflow configuration di Dataform

La workflow configuration definisce come le azioni SQLX vengono eseguite in automatico. Alla base di ogni workflow ci sono una "release configuration" e i suoi parametri. Le azioni selezionate vengono poi eseguite in risposta a un trigger specifico.

Il progetto R&D ha due workflow che si appoggiano a release configuration diverse e si attivano in momenti diversi. L'orario di avvio è definito con il formato 'cron'.

Workflow On-demand

Workflow Enterprise edition