from google.cloud.bigtable_v2 import BigtableAsyncClient
import asyncio
import grpc.aio

# Create a custom Interceptor class.
# See https://grpc.io/docs/guides/interceptors
# This example is for unary-unary but gRPC also supports unary-stream and others
# See https://grpc.github.io/grpc/python/grpc.html#grpc.UnaryUnaryClientInterceptor
class MetadataClientAsyncInterceptor(grpc.aio.UnaryUnaryClientInterceptor):
    def __init__(self):
        self._request_metadata = []
        self._response_metadata = []

    async def intercept_unary_unary(self, continuation, client_call_details, request):
        self._request_metadata = client_call_details.metadata
        response = await continuation(client_call_details, request)
        response_metadata = await response.trailing_metadata()
        #print(type(response_metadata))
        # Convert gRPC metadata `<class 'grpc.aio._metadata.Metadata'>` to list of tuples
        metadata = [(k, v) for k, v in response_metadata]
        self._response_metadata = metadata
        return response

async def async_main():
    # Get the transport class
    transport_cls = BigtableAsyncClient.get_transport_class("grpc_asyncio")
    # Create an interceptor object based on user-defined interceptor
    interceptor = MetadataClientAsyncInterceptor()
    # Create a gRPC channel
    channel = transport_cls.create_channel(host=BigtableAsyncClient.DEFAULT_ENDPOINT, interceptors = [interceptor])
    # Create a transport object, using the above channel
    transport = transport_cls(channel=channel)
    # Create a BigTable client using the above transport
    client = BigtableAsyncClient(transport=transport)
    # Call a BigTable API client method `ping_and_warm` and optionally provide additional request metadata.
    await client.ping_and_warm(name="projects/XXXX/instances/XXXX", metadata=[("some_additional_request_metadata_key", "some_additional_request_metadata_value")])
    # Print the request and response metadata
    print(interceptor._request_metadata)
    print(interceptor._response_metadata)

if __name__ == "__main__":
    asyncio.run(async_main())