airflow_dag_grpc_operator_quickstart_example.py

1import datetime
2
3from airflow import DAG
4from airflow.providers.grpc.operators.grpc import GrpcOperator
5
6# Note: You would typically have your proto-generated classes imported here
7# from my_app import my_grpc_pb2, my_grpc_pb2_grpc
8
9with DAG(
10    dag_id="example_grpc_operator",
11    schedule=None,
12    start_date=datetime.datetime(2021, 1, 1),
13    catchup=False,
14    tags=["example"],
15) as dag:
16
17    # The GrpcOperator allows you to call a gRPC method.
18    # 'grpc_conn_id' should point to an Airflow Connection of type 'grpc'.
19    # 'stub_class' is the generated stub class from your proto file.
20    # 'grpc_method' is the name of the method to call on that stub.
21    # 'data' is a dictionary of arguments passed to the method.
22    
23    call_grpc_service = GrpcOperator(
24        task_id="call_grpc_service",
25        stub_class="temp_pyspark_pb2_grpc.SparkJobServerStub",  # Example stub class string
26        grpc_method="execute_job",
27        grpc_conn_id="grpc_default",
28        data={"name": "test_job"},
29        streaming=False,
30        log_response=True,
31    )
32
33    call_grpc_service