Gestion des tâches et des graphiques de tâches Snowflake avec Python¶

Vous pouvez utiliser Python pour gérer les tâches Snowflake, avec lesquelles vous pouvez exécuter des instructionsSQL, des appels de procédure et une logique dans Exécution de scripts Snowflake. Pour une vue d’ensemble des tâches, voir Introduction aux tâches.

Les Snowflake Python APIs représentent des tâches avec deux types distincts :

Task : Expose les propriétés d’une tâche telles que sa planification, ses paramètres et ses prédécesseurs.
TaskResource : Expose des méthodes que vous pouvez utiliser pour récupérer un objet Task correspondant, exécuter la tâche et modifier la tâche.

Conditions préalables¶

Les exemples de cette rubrique supposent que vous ayez ajouté le code nécessaire pour vous connecter à Snowflake et créer un objet Root à partir duquel utiliser les Snowflake Python APIs.

Par exemple, le code suivant utilise les paramètres de connexion définis dans un fichier de configuration pour créer une connexion à Snowflake.

from snowflake.core import Root
from snowflake.snowpark import Session

session = Session.builder.config("connection_name", "myconnection").create()
root = Root(session)

En utilisant l’objet Session obtenu, le code crée un objet Root pour utiliser les types et les méthodes de l’API. Pour plus d’informations, voir Connexion à Snowflake avec Snowflake Python APIs.

Création d’une tâche¶

Pour créer une tâche, commencez par créer un objet Task. Ensuite, en spécifiant la base de données et le schéma dans lesquels créer la tâche, créez un objet TaskCollection. Via TaskCollection.create, ajoutez la nouvelle tâche à Snowflake.

Le code de l’exemple suivant crée un objet Task représentant une tâche nommée my_task qui exécute une requête SQL spécifiée dans le paramètre definition :

from datetime import timedelta
from snowflake.core.task import Task

my_task = Task(name="my_task", definition="<sql query>", schedule=timedelta(hours=1))
tasks = root.databases['my_db'].schemas['my_schema'].tasks
tasks.create(my_task)

Ce code crée une variable TaskCollection tasks à partir de la base de données my_db et du schéma my_schema. Via TaskCollection.create, il crée une nouvelle tâche dans Snowflake.

Cet exemple de code spécifie également une valeur timedelta d’une heure pour la planification de la tâche. Vous pouvez définir la planification d’une tâche via une valeur timedelta ou une expression Cron.

Vous pouvez également créer une tâche qui exécute une fonction Python ou une procédure stockée. Le code de l’exemple suivant crée une tâche nommée my_task2 qui exécute une fonction représentée par un objet StoredProcedureCall :

from snowflake.core.task import StoredProcedureCall, Task

my_task2 = Task(
  "my_task2",
  StoredProcedureCall(
      dosomething, stage_location="@mystage"
  ),
  warehouse="test_warehouse"
)
tasks = root.databases['my_db'].schemas['my_schema'].tasks
tasks.create(my_task2)

Cet objet spécifie une fonction nommée dosomething située dans l’emplacement de zone de préparation @mystage. Vous devez également spécifier un warehouse lorsque vous créez une tâche avec un objet StoredProcedureCall.

Création ou modification d’une tâche¶

Vous pouvez définir les propriétés d’un objet Task et le transmettre à la méthode TaskResource.create_or_alter pour créer une tâche si elle n’existe pas, ou la modifier selon la définition de la tâche si elle existe. Le comportement de create_or_alter est censé être idempotent, ce qui signifie que l’objet de tâche résultant sera le même, que la tâche existe ou non avant l’appel de la méthode.

Note

La méthode create_or_alter utilise des valeurs par défaut pour toute propriété de tâche que vous ne définissez pas explicitement. Par exemple, si vous ne définissez pas schedule, sa valeur par défaut est None même si la tâche existait auparavant avec une valeur différente.

Le code de l’exemple suivant met à jour la définition et la planification de la tâche my_task, puis modifie la tâche sur Snowflake :

from datetime import timedelta
from snowflake.core.task import Task

my_task = root.databases['my_db'].schemas['my_schema'].tasks['my_task'].fetch()
my_task.definition = "<sql query 2>"
my_task.schedule = timedelta(hours=2)

my_task_res = root.databases['my_db'].schemas['my_schema'].tasks['my_task']
my_task_res.create_or_alter(my_task)

Répertorier les tâches¶

Vous pouvez répertorier les tâches via la méthode TaskCollection.iter. La méthode renvoie un itérateur PagedIter d’objets Task.

Le code de l’exemple suivant répertorie les tâches dont le nom commence par my :

from snowflake.core.task import TaskCollection

tasks: TaskCollection = root.databases['my_db'].schemas['my_schema'].tasks
task_iter = tasks.iter(like="my%")  # returns a PagedIter[Task]
for task_obj in task_iter:
  print(task_obj.name)

Effectuer des opérations de tâche¶

Vous pouvez effectuer des opérations de tâche courantes—comme exécuter, suspendre et reprendre des tâches—avec un objet TaskResource.

Le code de l’exemple suivant exécute, suspend, reprend et supprime la tâche my_task :

tasks = root.databases['my_db'].schemas['my_schema'].tasks
task_res = tasks['my_task']

task_res.execute()
task_res.suspend()
task_res.resume()
task_res.drop()

Gestion des tâches dans un graphique de tâches¶

Vous pouvez gérer des tâches rassemblées dans un graphique de tâches. Un graphique de tâches est une série de tâches composée d’une seule tâche racine et de tâches supplémentaires, organisées en fonction de leurs dépendances.

Pour en savoir plus sur les tâches dans un graphique de tâches, voir Créer une séquence de tâches à l’aide d’un graphique de tâches.

Création d’un graphique de tâches¶

Pour créer un graphique de tâches, commencez par créer un objet DAG qui spécifie son nom et d’autres propriétés facultatives telles que sa planification. Vous pouvez définir la planification d’un graphique de tâches via une valeur timedelta ou une expression Cron.

Le code de l’exemple suivant définit une fonction Python dosomething, puis spécifie la fonction sous forme d’objet DAGTask nommé dag_task2 dans le graphique de tâches.

from snowflake.core.task import StoredProcedureCall
from snowflake.core.task.dagv1 import DAG, DAGTask, DAGOperation
from snowflake.snowpark import Session
from snowflake.snowpark.functions import sum as sum_

def dosomething(session: Session) -> None:
  df = session.table("target")
  df.group_by("a").agg(sum_("b")).save_as_table("agg_table")

with DAG("my_dag", schedule=timedelta(days=1)) as dag:
  # Create a task that runs some SQL.
  dag_task1 = DAGTask(
    "dagtask1",
    "MERGE INTO target USING source_stream WHEN MATCHED THEN UPDATE SET target.v = source_stream.v"
  )
  # Create a task that runs a Python function.
  dag_task2 = DAGTask(
    StoredProcedureCall(
      dosomething, stage_location="@mystage",
      packages=["snowflake-snowpark-python"]
    ),
    warehouse="test_warehouse"
  )
# Shift right and left operators can specify task relationships.
dag_task1 >> dag_task2  # dag_task1 is a predecessor of dag_task2
schema = root.databases["my_db"].schemas["my_schema"]
dag_op = DAGOperation(schema)
dag_op.deploy(dag)

Ce code définit également une instruction SQL sous la forme d’un autre objet DAGTask nommé dag_task1, puis spécifie dag_task1 comme prédécesseur de dag_task2. Pour finir, il déploie le graphique de tâches dans Snowflake dans la base de données my_db et le schéma my_schema.

Création d’un graphique de tâches avec une planification cron, des branches de tâches et des valeurs de renvoi de fonctions.¶

Vous pouvez également créer un graphique de tâches avec une planification cron spécifique, des branches de tâches et des valeurs de renvoi de fonctions qui sont utilisées comme valeurs de renvoi de tâches.

Le code de l’exemple suivant crée un objet DAG avec un objet Cron spécifiant sa planification. Il définit un objet DAGTaskBranch nommé task1_branch ainsi que d’autres objets DAGTask et spécifie leurs dépendances les uns par rapport aux autres :

from snowflake.core._common import CreateMode
from snowflake.core.task import Cron
from snowflake.core.task.dagv1 import DAG, DAGTask, DAGOperation, DAGTaskBranch
from snowflake.snowpark import Session

def task_handler(session: Session) -> None:
  pass  # do something

def task_branch_handler(session: Session) -> str:
  # do something
  return "task3"

try:
  with DAG(
    "my_dag",
    schedule=Cron("10 * * * *", "America/Los_Angeles"),
    stage_location="@mystage",
    packages=["snowflake-snowpark-python"],
    use_func_return_value=True,
  ) as dag:
    task1 = DAGTask(
      "task1",
      task_handler,
      warehouse=test_warehouse,
    )
    task1_branch = DAGTaskBranch("task1_branch", task_branch_handler, warehouse=test_warehouse)
    task2 = DAGTask("task2", task_handler, warehouse=test_warehouse)
    task3 = DAGTask("task3", task_handler, warehouse=test_warehouse, condition="1=1")
    task1 >> task1_branch
    task1_branch >> [task2, task3]
  schema = root.databases["my_db"].schemas["my_schema"]
  op = DAGOperation(schema)
  op.deploy(dag, mode=CreateMode.or_replace)
finally:
  session.close()

Cet exemple de code définit également les fonctions de gestion des tâches et crée chaque objet DAGTask et DAGTaskBranch avec un gestionnaire de tâches spécifique affecté à la tâche. Le code définit le paramètre use_func_return_value du DAG sur True, qui indique qu’il convient d’utiliser la valeur de renvoi de la fonction Python comme valeur de renvoi de la tâche correspondante. Sinon, la valeur par défaut de use_func_return_value est False.

Définition et obtention de la valeur de retour d’une tâche dans un graphique de tâches¶

Lorsque la définition d’une tâche est un objet StoredProcedureCall, le gestionnaire de la procédure stockée (ou de la fonction) peut explicitement définir la valeur de renvoi de la tâche via un objet TaskContext.

Pour plus d’informations, voir SYSTEM$SET_RETURN_VALUE.

Le code de l’exemple suivant définit une fonction de gestion des tâches qui crée un objet TaskContext nommé context à partir de la session en cours. Il utilise ensuite la méthode TaskContext.set_return_value pour définir explicitement la valeur de renvoi sur une chaîne spécifiée :

from snowflake.core.task.context import TaskContext
from snowflake.snowpark import Session

def task_handler(session: Session) -> None:
  context = TaskContext(session)
  # this return value can be retrieved by successor Tasks.
  context.set_return_value("predecessor_return_value")

Dans un graphique de tâches, une tâche de successeur immédiate qui identifie la tâche précédente comme son prédécesseur peut alors récupérer la valeur de renvoi explicitement définie par la tâche de prédécesseur.

Pour plus d’informations, voir SYSTEM$GET_PREDECESSOR_RETURN_VALUE.

Le code de l’exemple suivant définit une fonction de gestionnaire de tâches qui utilise la méthode TaskContext.get_predecessor_return_value pour obtenir la valeur de renvoi de la tâche de prédécesseur nommée pred_task_name :

from snowflake.core.task.context import TaskContext
from snowflake.snowpark import Session

def task_handler(session: Session) -> None:
  context = TaskContext(session)
  pred_return_value = context.get_predecessor_return_value("pred_task_name")