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

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.

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.

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)

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()

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

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