def test_subscribe_all(self) -> None:
# Tests `subscribe_all_tables()`.
# Connect to a single fake cluster.
conn = self.px_client.connect_to_cluster(
self.px_client.list_healthy_clusters()[0])
# Create two tables and simulate them sent over as part of the ExecuteScript call.
http_table1 = self.http_table_factory.create_table(
test_utils.table_id1)
stats_table1 = self.stats_table_factory.create_table(
test_utils.table_id3)
self.fake_vizier_service.add_fake_data(conn.cluster_id, [
http_table1.metadata_response(),
http_table1.row_batch_response([["foo"], [200]]),
stats_table1.metadata_response(),
stats_table1.row_batch_response([
[vpb.UInt128(high=123, low=456)],
[1000],
[999],
]),
http_table1.end(),
stats_table1.end(),
])
# Create script_executor.
script_executor = conn.prepare_script(pxl_script)
# Get a subscription to all of the tables that arrive over the stream.
tables = script_executor.subscribe_all_tables()
# Async function to run on the "http" table.
async def process_http_tb(table_sub: pxapi.TableSub) -> None:
num_rows = 0
async for row in table_sub:
self.assertEqual(row["http_resp_body"], "foo")
self.assertEqual(row["http_resp_status"], 200)
num_rows += 1
self.assertEqual(num_rows, 1)
# Async function that processes the tables subscription and runs the
# async function above to process the "http" table when that table shows
# we see thtparticular table.
async def process_all_tables(
tables_gen: pxapi.TableSubGenerator) -> None:
table_names = set()
async for table in tables_gen:
table_names.add(table.table_name)
if table.table_name == "http":
# Once we find the http_tb, process it.
await process_http_tb(table)
# Make sure we see both tables on the generator.
self.assertEqual(table_names, {"http", "stats"})
# Run the script_executor and process_all_tables function concurrently.
# We expect no errors.
loop = asyncio.get_event_loop()
loop.run_until_complete(
run_script_and_tasks(script_executor,
[process_all_tables(tables())]))
def test_fail_on_multi_run(self) -> None:
# Tests to show that queries may only be run once. After a script_executor has been
# run, calling data grabbing methods like subscribe, add_callback, etc. will
# raise an error.
# Connect to a single fake cluster.
conn = self.px_client.connect_to_cluster(
self.px_client.list_healthy_clusters()[0])
stats_table1 = self.stats_table_factory.create_table(
test_utils.table_id3)
self.fake_vizier_service.add_fake_data(conn.cluster_id, [
stats_table1.metadata_response(),
stats_table1.row_batch_response([
[vpb.UInt128(high=123, low=456)],
[1000],
[999],
]),
stats_table1.end(),
])
script_executor = conn.prepare_script(pxl_script)
# Create a dummy callback.
def stats_cb(row: pxapi.Row) -> None:
pass
script_executor.add_callback("stats", stats_cb)
# Run the script_executor for the first time. Should not return an error.
script_executor.run()
# Each of the following methods should fail if called after script_executor.run()
script_ran_message = "Script already executed"
# Adding a callback should fail.
with self.assertRaisesRegex(ValueError, script_ran_message):
script_executor.add_callback("stats", stats_cb)
# Subscribing to a table should fail.
with self.assertRaisesRegex(ValueError, script_ran_message):
script_executor.subscribe("stats")
# Subscribing to all tables should fail.
with self.assertRaisesRegex(ValueError, script_ran_message):
script_executor.subscribe_all_tables()
# Synchronous run should error out.
with self.assertRaisesRegex(ValueError, script_ran_message):
script_executor.run()
# Async run should error out.
loop = asyncio.new_event_loop()
asyncio.set_event_loop(loop)
with self.assertRaisesRegex(ValueError, script_ran_message):
loop.run_until_complete(script_executor.run_async())
def make_upid() -> vpb.UInt128:
return vpb.UInt128(high=123, low=456)
def test_run_script_callback(self) -> None:
# Test the callback API. Callback API is a simpler alternative to the TableSub
# API that allows you to designate a function that runs on individual rows. Users
# can process data without worrying about async processing by using this API.
# Connect to a single fake cluster.
conn = self.px_client.connect_to_cluster(
self.px_client.list_healthy_clusters()[0])
# Create two tables: "http" and "stats"
http_table1 = self.http_table_factory.create_table(
test_utils.table_id1)
stats_table1 = self.stats_table_factory.create_table(
test_utils.table_id3)
self.fake_vizier_service.add_fake_data(
conn.cluster_id,
[
# Init "http".
http_table1.metadata_response(),
# Send data for "http".
http_table1.row_batch_response([["foo"], [200]]),
# Init "stats".
stats_table1.metadata_response(),
# Send data for "stats".
stats_table1.row_batch_response([
[vpb.UInt128(high=123, low=456)],
[1000],
[999],
]),
# End "http".
http_table1.end(),
# End "stats".
stats_table1.end(),
])
script_executor = conn.prepare_script(pxl_script)
http_counter = 0
stats_counter = 0
# Define callback function for "http" table.
def http_fn(row: pxapi.Row) -> None:
nonlocal http_counter
http_counter += 1
self.assertEqual(row["http_resp_body"], "foo")
self.assertEqual(row["http_resp_status"], 200)
script_executor.add_callback("http", http_fn)
# Define a callback function for the stats_fn.
def stats_fn(row: pxapi.Row) -> None:
nonlocal stats_counter
stats_counter += 1
self.assertEqual(row["upid"],
uuid.UUID('00000000-0000-007b-0000-0000000001c8'))
self.assertEqual(row["cpu_ktime_ns"], 1000)
self.assertEqual(row["rss_bytes"], 999)
script_executor.add_callback("stats", stats_fn)
# Run the script_executor synchronously.
script_executor.run()
# We expect each callback function to only be called once.
self.assertEqual(stats_counter, 1)
self.assertEqual(http_counter, 1)
def test_one_conn_two_tables(self) -> None:
# Connect to a single fake cluster.
conn = self.px_client.connect_to_cluster(
self.px_client.list_healthy_clusters()[0])
# We will send two tables for this test "http" and "stats".
http_table1 = self.http_table_factory.create_table(
test_utils.table_id1)
stats_table1 = self.stats_table_factory.create_table(
test_utils.table_id3)
self.fake_vizier_service.add_fake_data(
conn.cluster_id,
[
# Initialize "http" on the stream.
http_table1.metadata_response(),
# Send over a row-batch from "http".
http_table1.row_batch_response([["foo"], [200]]),
# Initialize "stats" on the stream.
stats_table1.metadata_response(),
# Send over a row-batch from "stats".
stats_table1.row_batch_response([
[vpb.UInt128(high=123, low=456)],
[1000],
[999],
]),
# Send an end-of-stream for "http".
http_table1.end(),
# Send an end-of-stream for "stats".
stats_table1.end(),
])
script_executor = conn.prepare_script(pxl_script)
# Subscribe to both tables.
http_tb = script_executor.subscribe("http")
stats_tb = script_executor.subscribe("stats")
# Async function that makes sure "http" table returns the expected rows.
async def process_http_tb(table_sub: pxapi.TableSub) -> None:
num_rows = 0
async for row in table_sub:
self.assertEqual(row["http_resp_body"], "foo")
self.assertEqual(row["http_resp_status"], 200)
num_rows += 1
self.assertEqual(num_rows, 1)
# Async function that makes sure "stats" table returns the expected rows.
async def process_stats_tb(table_sub: pxapi.TableSub) -> None:
num_rows = 0
async for row in table_sub:
self.assertEqual(
row["upid"],
uuid.UUID('00000000-0000-007b-0000-0000000001c8'))
self.assertEqual(row["cpu_ktime_ns"], 1000)
self.assertEqual(row["rss_bytes"], 999)
num_rows += 1
self.assertEqual(num_rows, 1)
# Run the script_executor and the processing tasks concurrently.
loop = asyncio.get_event_loop()
loop.run_until_complete(
run_script_and_tasks(
script_executor,
[process_http_tb(http_tb),
process_stats_tb(stats_tb)]))
