apache-airflow-providers-grpc

from __future__ import annotations import os from datetime import datetime from airflow.models.dag import DAG from airflow.providers.grpc.operators.grpc import GrpcOperator from airflow.providers.grpc.hooks.grpc import GrpcHook # NOTE: In a real scenario, you would generate protobuf stub classes # For this example, we assume a 'dummy_pb2_grpc' and 'dummy_pb2' exist # and define a 'DummyService' with a 'Ping' method that takes 'PingRequest' # and returns 'PingResponse'. # Replace with your actual generated proto files and service/method names. try: import grpc from dummy_pb2_grpc import DummyServiceStub from dummy_pb2 import PingRequest, PingResponse except ImportError: # Placeholder for running the quickstart without actual protobufs class DummyServiceStub: def __init__(self, channel): pass def Ping(self, request): print(f"Mock Ping called with: {request.message}") return type('PingResponse', (object,), {'message': 'Pong from mock!'}) class PingRequest: def __init__(self, message=None): self.message = message PingResponse = type('PingResponse', (object,), {'message': None}) print("Warning: Protobuf stubs (dummy_pb2_grpc, dummy_pb2) not found. Using mock classes.") with DAG( dag_id='grpc_example_dag', start_date=datetime(2023, 1, 1), schedule=None, catchup=False, tags=['grpc', 'example'], ) as dag: # Define a gRPC connection in Airflow UI (Admin -> Connections) # Conn Id: 'grpc_default' # Conn Type: 'gRPC Connection' # Host: 'localhost' # Port: '50051' # Auth Type: 'NO_AUTH' for insecure channel, or 'SSL'/'TLS' with 'Credential Pem File' # For JWT_GOOGLE or OATH_GOOGLE, set these and ensure 'google-auth' is installed # Example of using GrpcOperator # This assumes 'grpc_default' connection is configured in Airflow call_grpc_service = GrpcOperator( task_id='call_grpc_service', grpc_conn_id='grpc_default', stub_class=DummyServiceStub, # Your generated gRPC stub class call_func='Ping', # The method on the stub class to call data={'request': PingRequest(message='Hello gRPC from Airflow!')}, # If the gRPC method expects a specific request object, pass it like this. # The 'data' parameter maps to the keyword arguments of the gRPC method. ) # Example of using GrpcHook directly in a PythonOperator (or for custom logic) def _interact_with_grpc_hook(**kwargs): hook = GrpcHook(grpc_conn_id='grpc_default') # For this example, we're using a mock service. In a real scenario, # hook.run() would return a gRPC response or yield streaming responses. response = hook.run(stub_class=DummyServiceStub, call_func='Ping', data={'request': PingRequest(message='Hello from Hook!')}) # For unary calls, response is typically the deserialized message # For streaming, it's an iterator. print(f"gRPC Hook Response: {response.message}") interact_with_hook = GrpcOperator( task_id='interact_with_grpc_hook', grpc_conn_id='grpc_default', stub_class=DummyServiceStub, call_func='Ping', data={'request': PingRequest(message='Hello from Hook via Operator!')} )