def test_direct_conns(self) -> None:
# Test the direct connections.
# Create the direct conn server.
dc_server = grpc.server(futures.ThreadPoolExecutor(max_workers=1))
fake_dc_service = VizierServiceFake()
# Create one http table.
http_table1 = self.http_table_factory.create_table(
test_utils.table_id1)
cluster_id = "10000000-0000-0000-0000-000000000004"
fake_dc_service.add_fake_data(
cluster_id,
[
# Init "http".
http_table1.metadata_response(),
# Send data for "http".
http_table1.row_batch_response([["foo"], [200]]),
# End "http".
http_table1.end(),
])
vizierapi_pb2_grpc.add_VizierServiceServicer_to_server(
fake_dc_service, dc_server)
port = dc_server.add_insecure_port("[::]:0")
dc_server.start()
url = f"http://[::]:{port}"
token = cluster_id
self.fake_cloud_service.add_direct_conn_cluster(
cluster_id, "dc_cluster", url, token)
clusters = self.px_client.list_healthy_clusters()
self.assertSetEqual(set([c.name() for c in clusters]),
{"cluster1", "cluster2", "dc_cluster"})
conns = [
self.px_client.connect_to_cluster(c) for c in clusters
if c.name() == 'dc_cluster'
]
self.assertEqual(len(conns), 1)
script_executor = conns[0].prepare_script(pxl_script)
# Define callback function for "http" table.
def http_fn(row: pxapi.Row) -> None:
self.assertEqual(row["http_resp_body"], "foo")
self.assertEqual(row["http_resp_status"], 200)
script_executor.add_callback("http", http_fn)
# Run the script_executor synchronously.
script_executor.run()
def test_shared_grpc_channel_for_cloud(self) -> None:
# Setup a direct connect cluster.
# Create the direct conn server.
dc_server = grpc.server(futures.ThreadPoolExecutor(max_workers=1))
cluster_id = "10000000-0000-0000-0000-000000000004"
fake_dc_service = VizierServiceFake()
vizierapi_pb2_grpc.add_VizierServiceServicer_to_server(
fake_dc_service, dc_server)
port = dc_server.add_insecure_port("[::]:0")
dc_server.start()
url = f"http://[::]:{port}"
token = cluster_id
self.fake_cloud_service.add_direct_conn_cluster(
cluster_id, "dc_cluster", url, token)
# Make sure the shraed grpc channel are actually shared.
num_create_channel_calls = 0
def cloud_channel_fn(url: str) -> grpc.Channel:
# Nonlocal because we're incrementing the outer variable.
nonlocal num_create_channel_calls
num_create_channel_calls += 1
return grpc.insecure_channel(url)
px_client = pxapi.Client(
token=ACCESS_TOKEN,
server_url=self.url(),
# Channel functions for testing.
channel_fn=cloud_channel_fn,
conn_channel_fn=lambda url: grpc.aio.insecure_channel(url),
)
# Connect to a cluster.
healthy_clusters = px_client.list_healthy_clusters()
self.assertEqual(num_create_channel_calls, 1)
# No new channels made on later calls to cloud
px_client.connect_to_cluster(healthy_clusters[0])
self.assertEqual(num_create_channel_calls, 1)
# We need to make more cloud calls for direct connect clusters. Do not need new channels.
dc_cluster = [c for c in healthy_clusters if not c.passthrough()][0]
px_client.connect_to_cluster(dc_cluster)
self.assertEqual(num_create_channel_calls, 1)
# Reset cache, so now must create a new channel.
px_client._cloud_channel_cache = None
px_client.connect_to_cluster(healthy_clusters[0])
self.assertEqual(num_create_channel_calls, 2)
def setUp(self) -> None:
# Create a fake server for the VizierService
self.server = grpc.server(futures.ThreadPoolExecutor(max_workers=1))
self.fake_vizier_service = VizierServiceFake()
self.fake_cloud_service = CloudServiceFake()
vizierapi_pb2_grpc.add_VizierServiceServicer_to_server(
self.fake_vizier_service, self.server)
cloudapi_pb2_grpc.add_VizierClusterInfoServicer_to_server(
self.fake_cloud_service, self.server)
self.port = self.server.add_insecure_port("[::]:0")
self.server.start()
self.px_client = pxapi.Client(
token=ACCESS_TOKEN,
server_url=self.url(),
channel_fn=lambda url: grpc.insecure_channel(url),
conn_channel_fn=lambda url: grpc.aio.insecure_channel(url),
)
self.http_table_factory = test_utils.FakeTableFactory(
"http",
vpb.Relation(columns=[
test_utils.string_col("http_resp_body"),
test_utils.int64_col("http_resp_status"),
]))
self.stats_table_factory = test_utils.FakeTableFactory(
"stats",
vpb.Relation(columns=[
test_utils.uint128_col("upid"),
test_utils.int64_col("cpu_ktime_ns"),
test_utils.int64_col("rss_bytes"),
]))
asyncio.set_event_loop(asyncio.new_event_loop())