# Copyright 2023 Google LLC

#

# Licensed under the Apache License, Version 2.0 (the "License");

# you may not use this file except in compliance with the License.

# You may obtain a copy of the License at

#

# http://www.apache.org/licenses/LICENSE-2.0

#

# Unless required by applicable law or agreed to in writing, software

# distributed under the License is distributed on an "AS IS" BASIS,

# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.

# See the License for the specific language governing permissions and

# limitations under the License.

"""

The Python implementation of the `cloud-bigtable-clients-test` proxy server.

https://github.com/googleapis/cloud-bigtable-clients-test

This server is intended to be used to test the correctness of Bigtable

clients across languages.

Contributor Note: the proxy implementation is split across TestProxyClientHandler

and TestProxyGrpcServer. This is due to the fact that generated protos and proto-plus

objects cannot be used in the same process, so we had to make use of the

multiprocessing module to allow them to work together.

"""

import multiprocessing

import argparse

import sys

import os

sys.path.append("handlers")

def grpc_server_process(request_q, queue_pool, port=50055):

"""

Defines a process that hosts a grpc server

proxies requests to a client_handler_process

"""

sys.path.append("protos")

from concurrent import futures

import grpc

import test_proxy_pb2_grpc

import grpc_handler

# Start gRPC server

server = grpc.server(futures.ThreadPoolExecutor(max_workers=10))

test_proxy_pb2_grpc.add_CloudBigtableV2TestProxyServicer_to_server(

grpc_handler.TestProxyGrpcServer(request_q, queue_pool), server

)

server.add_insecure_port("[::]:" + port)

server.start()

print("grpc_server_process started, listening on " + port)

server.wait_for_termination()

async def client_handler_process_async(request_q, queue_pool, client_type="async"):

"""

Defines a process that recives Bigtable requests from a grpc_server_process,

and runs the request using a client library instance

"""

import base64

import re

import asyncio

import warnings

import client_handler_data_async

warnings.filterwarnings("ignore", category=RuntimeWarning, message=".*Bigtable emulator.*")

def camel_to_snake(str):

return re.sub(r"(?<!^)(?=[A-Z])", "_", str).lower()

def format_dict(input_obj):

if isinstance(input_obj, list):

return [format_dict(x) for x in input_obj]

elif isinstance(input_obj, dict):

return {camel_to_snake(k): format_dict(v) for k, v in input_obj.items()}

elif isinstance(input_obj, str):

# check for time encodings

if re.match("^[0-9]+s$", input_obj):

return int(input_obj[:-1])

# check for encoded bytes

if re.match("^[A-Za-z0-9+/=]+$", input_obj):

try:

decoded_str = base64.b64decode(input_obj)

# if the string contains non-ascii bytes, raise exception

decoded_str.decode("ascii")

return decoded_str

except Exception:

pass

# check for int strings

try:

return int(input_obj)

except (ValueError, TypeError):

return input_obj

else:

return input_obj

# Listen to requests from grpc server process

print_msg = f"client_handler_process started with client_type={client_type}"

print(print_msg)

client_map = {}

background_tasks = set()

while True:

if not request_q.empty():

json_data = format_dict(request_q.get())

fn_name = json_data.pop("proxy_request")

print(f"--- running {fn_name} with {json_data}")

out_q = queue_pool[json_data.pop("response_queue_idx")]

client_id = json_data.pop("client_id")

client = client_map.get(client_id, None)

# handle special cases for client creation and deletion

if fn_name == "CreateClient":

if client_type == "legacy":

import client_handler_legacy

client = client_handler_legacy.LegacyTestProxyClientHandler(**json_data)

elif client_type == "sync":

import client_handler_data_sync_autogen

client = client_handler_data_sync_autogen.TestProxyClientHandler(**json_data)

else:

client = client_handler_data_async.TestProxyClientHandlerAsync(**json_data)

client_map[client_id] = client

out_q.put(True)

elif client is None:

out_q.put(RuntimeError("client not found"))

elif fn_name == "CloseClient":

client.close()

out_q.put(True)

elif fn_name == "RemoveClient":

client_map.pop(client_id, None)

out_q.put(True)

print("")

else:

# run actual rpc against client

async def _run_fn(out_q, fn, **kwargs):

result = await fn(**kwargs)

out_q.put(result)

fn = getattr(client, fn_name)

task = asyncio.create_task(_run_fn(out_q, fn, **json_data))

await asyncio.sleep(0)

background_tasks.add(task)

task.add_done_callback(background_tasks.remove)

await asyncio.sleep(0.01)

def client_handler_process(request_q, queue_pool, client_type="async"):

"""

Sync entrypoint for client_handler_process_async

"""

import asyncio

asyncio.run(client_handler_process_async(request_q, queue_pool, client_type))

p = argparse.ArgumentParser()

p.add_argument("--port", dest='port', default="50055")

p.add_argument("--client_type", dest='client_type', default="async", choices=["async", "sync", "legacy"])

if __name__ == "__main__":

port = p.parse_args().port

client_type = p.parse_args().client_type

# start and run both processes

# larger pools support more concurrent requests

response_queue_pool = [multiprocessing.Queue() for _ in range(100)]

request_q = multiprocessing.Queue()

CLIENT_IS_FOREGROUND=bool(os.environ.get("CLIENT_IS_FOREGROUND", True))

if CLIENT_IS_FOREGROUND:

# run client in forground and proxy in background

# breakpoints can be attached to handlers/client_handler_data.py

proxy = multiprocessing.Process(

target=grpc_server_process,

args=(

request_q,

response_queue_pool,

port

),

)

proxy.start()

client_handler_process(request_q, response_queue_pool, client_type)

proxy.join()

else:

# run proxy in forground and client in background

# breakpoints can be attached to handlers/grpc_handler.py

client = multiprocessing.Process(

target=client_handler_process,

args=(

request_q,

response_queue_pool,

),

)

client.start()

grpc_server_process(request_q, response_queue_pool, port)

client.join()