def test_encryption(self) -> None:
# Test creating encrypted clients.
px_client = pxapi.Client(
token=ACCESS_TOKEN,
server_url=self.url(),
use_encryption=True,
channel_fn=lambda url: grpc.insecure_channel(url),
conn_channel_fn=lambda url: grpc.aio.insecure_channel(url),
)
conn = px_client.connect_to_cluster(
px_client.list_healthy_clusters()[0])
# Create the script_executor.
script_executor = conn.prepare_script(pxl_script)
self.assertTrue(script_executor._use_encryption)
# Create table for cluster_uuid1.
http_table1 = self.http_table_factory.create_table(
test_utils.table_id1)
self.fake_vizier_service.add_fake_data(
conn.cluster_id,
[
# Initialize the table on the stream with the metadata.
http_table1.metadata_response(),
# Send over a single-row batch.
http_table1.row_batch_response([["foo"], [200]]),
# Send an end-of-stream for the table.
http_table1.end(),
])
# Use the results API to run and get the data from the http table.
for row in script_executor.results("http"):
self.assertEqual(row["resp_body"], "foo")
self.assertEqual(row["resp_status"], 200)
def test_shared_grpc_channel_for_cloud(self) -> None:
# 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)
# 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 main():
logging.debug("Authorizing Pixie client.")
px_client = pxapi.Client(token=pixie_api_key)
cluster_conn = px_client.connect_to_cluster(pixie_cluster_id)
logging.debug("Pixie client connected to %s cluster.", cluster_conn.name())
logging.debug("Authorizing Slack client.")
slack_client = WebClient(slack_bot_token)
# Slack channel for Slackbot to post in.
# Slack App must be a member of this channel.
SLACK_CHANNEL = "#pixie-alerts"
# Schedule sending a Slack channel message every 5 minutes.
schedule.every(5).minutes.do(lambda: send_slack_message(slack_client,
SLACK_CHANNEL,
cluster_conn))
logging.info("Message scheduled for %s Slack channel.", SLACK_CHANNEL)
while True:
schedule.run_pending()
time.sleep(5)
def test_shared_grpc_channel_on_conns(self) -> None:
# Make sure the shraed grpc channel does what we expect.
num_create_channel_calls = 0
def conn_channel_fn(url: str) -> grpc.aio.Channel:
# Nonlocal because we're incrementing the outer variable.
nonlocal num_create_channel_calls
num_create_channel_calls += 1
return grpc.aio.insecure_channel(url)
px_client = pxapi.Client(
token=ACCESS_TOKEN,
server_url=self.url(),
# Channel functions for testing.
channel_fn=lambda url: grpc.insecure_channel(url),
conn_channel_fn=conn_channel_fn,
)
# Connect to a cluster.
conn = px_client.connect_to_cluster(
px_client.list_healthy_clusters()[0])
# Create fake data for table for cluster_uuid1.
http_table1 = self.http_table_factory.create_table(
test_utils.table_id1)
self.fake_vizier_service.add_fake_data(
conn.cluster_id,
[
# Initialize the table on the stream with the metadata.
http_table1.metadata_response(),
# Send over a single-row batch.
http_table1.row_batch_response([["foo"], [200]]),
# Send an end-of-stream for the table.
http_table1.end(),
])
# Create the script executor.
script_executor = conn.prepare_script(pxl_script)
script_executor.run()
self.assertEqual(num_create_channel_calls, 1)
# Creating another script will not create another channel call.
script_executor = conn.prepare_script(pxl_script)
script_executor.run()
self.assertEqual(num_create_channel_calls, 1)
# Reset cache, so now must create a new channel.
conn._channel_cache = None
# Now we expect the channel to be created again.
script_executor = conn.prepare_script(pxl_script)
script_executor.run()
self.assertEqual(num_create_channel_calls, 2)
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())
#
# 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.
#
# SPDX-License-Identifier: Apache-2.0
import os
import pxapi
# You'll need to generate an API token.
# For more info, see: https://docs.px.dev/using-pixie/api-quick-start/
API_TOKEN = os.getenv("PX_API_KEY")
# create a Pixie client
px_client = pxapi.Client(token=API_TOKEN, use_encryption=True)
clusters = px_client.list_healthy_clusters()
# print results
for row in clusters:
print(row.id)
API_TOKEN = "<YOUR_API_TOKEN>"
# PxL script with 2 output tables
PXL_SCRIPT = """
import px
df = px.DataFrame('http_events')[['http_resp_status','http_req_path']]
df = df.head(10)
px.display(df, 'http_table')
df2 = px.DataFrame('process_stats')[['upid','cpu_ktime_ns']]
df2 = df2.head(10)
px.display(df2, 'stats')
"""
# create a Pixie client
px_client = pxapi.Client(token=API_TOKEN)
# find first healthy cluster
cluster = px_client.list_healthy_clusters()[0]
print("Selected cluster: ", cluster.name())
# connect to cluster
conn = px_client.connect_to_cluster(cluster)
# setup the PxL script
script = conn.prepare_script(PXL_SCRIPT)
# create a function to process the table rows
def http_fn(row: pxapi.Row) -> None:
print(row["http_resp_status"], row["http_req_path"])