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)]))