Automatize reservas e atribuições do BigQuery com Dataform

Um dos nossos clientes mantém um projeto de Pesquisa & Desenvolvimento no BigQuery, usado tanto pelo time de Desenvolvimento quanto pelo de Data Science. Os times têm um padrão de uso bem definido, com as seguintes características:

1. As consultas interativas acontecem principalmente em dias úteis, no horário comercial.
2. Alta concorrência de consultas, com tempos de resposta rápidos.
3. Algumas tarefas específicas são agendadas para rodar à noite ou nos fins de semana.

Para atender a essas demandas, o projeto vai usar o plano 'On-Demand' no horário comercial, garantindo baixa latência nas consultas e alta concorrência. Fora do horário de trabalho, o projeto passa a usar a 'Enterprise Edition' com no máximo 100 slots.

Um fluxo de automação alterna o plano de preços conforme o momento.

Hoje, o Google não oferece uma solução pronta para alternar automaticamente entre os planos de preços do BigQuery. Em compensação, disponibiliza uma interface para gerenciar reservas via CLI ou SQL.

A seguir, vamos usar o Dataform e sua linguagem SQLX para deixar a atualização dos planos de preços mais eficiente. Essa abordagem permite gerenciar qualquer mudança em reservas e atribuições dentro de um workflow do Dataform, com versionamento no Git.

Outra vantagem do Dataform é a simplicidade: tudo é feito em SQL, dentro do próprio BigQuery, sem precisar de serviços externos. O Dataform é gratuito, então não há custo adicional envolvido.

Como preparação, vamos criar manualmente uma reserva 'Enterprise Edition' chamada test_rd, com o limite máximo de 100 slots.

Em seguida, para gerenciar a configuração da reserva e as mudanças de atribuição, vamos usar a BigQuery Reservation API (SQL DDL).

Veja um exemplo de código SQL para criar uma atribuição na região US, usando uma reserva chamada ndv-rd e o assignment_id df_rd:

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

Se você ainda não conhece o Dataform, recomendamos a leitura do post Build SQL pipelines to BigQuery with Dataform e a execução do 'Quick start'.

Todo o código deste post está disponível neste repositório do GitHub.

São seis etapas para montar um workflow automatizado:

[1] Crie um repositório e atribua as roles

Para criar um repositório do Dataform no projeto admin do BigQuery, acesse a seção Dataform e inicie o processo de criação. (Criar um repositório do Dataform)

Para que a execução funcione corretamente, tanto o principal do desenvolvedor (usuário) quanto a SA do Dataform (service account) precisam das seguintes roles durante a execução do workflow: bigquery.resourceAdmin, bigquery.user

Para atribuir as roles à SA (código abaixo), use o ID da SA exibido na mensagem que aparece após a criação do repositório.

A SA do 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] Crie um workspace de desenvolvimento

Para criar um novo workspace de desenvolvimento do Dataform, siga estas etapas: Criar um workspace do Dataform.

Depois de inicializado, o workspace já vem com alguns métodos SQLX de exemplo e um arquivo de configuração chamado "dataform.json". Esses arquivos de exemplo (first_view.sqlx, second_view.sqlx) podem ser removidos.

Para configurar os métodos do Dataform, adicione a seção "vars", que contém a lista de parâmetros mostrada abaixo. Essa seção deve aparecer depois do último parâmetro existente e incluir os parâmetros listados a seguir. Nela, você define os valores padrão de cada parâmetro, que podem ser sobrescritos no workflow mais adiante.

1. Não esqueça de adicionar uma vírgula "," antes da seção "vars".
2. Substitua os parâmetros marcados com "<>" (e remova os "<>").

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

Veja um exemplo da nossa implementação:

"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] Desenvolva os métodos SQLX

Nesta etapa, vamos criar o código SQLX que vai cuidar das reservas e atribuições.

• Crie um arquivo chamado drop_and_create_assignment.sqlx dentro da pasta 'definitions'.
• Inclua o código SQLX abaixo para excluir todas as atribuições do projeto e criar a atribuição necessária.
• Confirme se a compilação foi bem-sucedida pelo ✅:

Este é o código SQLX do 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 este é o código gerado esperado, conforme a nossa configuração:

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] Teste e valide os métodos SQLX

Antes de rodar o código para testar, lembre-se de que a configuração dos projetos atribuídos será alterada de acordo com o valor definido no arquivo dataform.json. Vale a pena testar primeiro em um projeto de 'teste'.

Para executar e testar o código, você tem duas opções: clicar em 'Run' usando suas próprias credenciais ou em 'Start Execution' para executar com a SA do Dataform, conforme indicado na mensagem abaixo.

Execute todas as ações ou selecione um subconjunto delas. A service account [email protected] será utilizada.

Confirme pela saída da execução: This statement created a new assignment. Em seguida, verifique as alterações na visualização 'Capacity management'.

A reserva está atribuída ao projeto

Com a verificação da SA concluída com sucesso, é hora de fazer o commit das alterações e o push para a branch padrão. A partir daí, o desenvolvimento está pronto para o release.

[5] Crie a configuração de release do Dataform

Para gerar um snapshot do código e executá-lo com parâmetros sobrescritos, use a release configuration. Para simplificar, vamos configurar apenas um parâmetro e definir o refresh como Never.

Lembre-se de compilar após cada alteração.

O projeto de R&D vai ter dois releases: um para o 'On-Demand', com reservation_name= `none``` e outro para a 'Enterprise edition', reservation_name=ndv-rd``` (preste atenção nos ``""``).

Teste a "release configuration" com a ação selecionada:

ndv-playground-bq-mgmt.dataform.drop_and_create_assignment

Release para a reserva Enterprise

Release para a reserva On-demand

[6] Crie a configuração de workflow do Dataform

A configuração do workflow define como as ações SQLX são executadas automaticamente. Cada workflow tem como base uma "release configuration" e seus parâmetros. As ações selecionadas são, então, executadas a partir de um trigger específico.

O projeto de R&D tem dois workflows que utilizam release configurations diferentes e são acionados em horários distintos. O horário de início é definido no formato 'cron'.

Workflow On-demand

Workflow Enterprise edition