vacp2p/pod-api-requester
1
0
Fork 0
mirror of https://github.com/vacp2p/pod-api-requester.git synced 2026-08-27 09:41:16 +00:00
Code Issues Packages Projects Releases Wiki Activity

Refactor to use configs and server

Browse Source 
- Use configs to define actions

- Run a server to allow running commands interactively
This commit is contained in:
Pearson White
2026-01-11 16:37:14 -05:00
parent 5d23c3e2c8
commit 953da9b0c0
13 changed files with 1230 additions and 670 deletions
Download Patch File Download Diff File Expand all files Collapse all files
Dockerfile
+28 -3
View File
@@ -1,5 +1,30 @@
FROM python:3.11.9-alpine
FROM python:3.11.9-alpine AS base
WORKDIR /app
ADD store_msg_retriever.py /app/store_msg_retriever.py
COPY requirements.txt .
RUN pip install --no-cache-dir --break-system-packages -r requirements.txt
COPY api_requester.py utils.py configs.py /app/
RUN pip install requests pydantic
FROM base AS debug
WORKDIR /app
RUN apk add --no-cache \
bash \
bind-tools \
curl \
ethtool \
iputils \
jq \
net-tools \
tcpdump \
vim \
wget \
ws \
nodejs \
npm \
&& npm install -g wscat
ENTRYPOINT ["sleep", "infinity"]
FROM base AS production
WORKDIR /app
ENTRYPOINT ["python", "./api_requester.py", "--mode", "server", "--config", "/mount/config.yaml"]
README.md
+56 -26
View File
@@ -1,40 +1,70 @@
## Waku Storage Retriever
This Python script retrieves messages
from a Waku storage service using an HTTP API.
It supports pagination and resolves DNS for the service host.
This Python script facilitates arbitrary GET and POST requests
to pods in a Kubernetes cluster.
The script is designed to run inside a Docker container.
### Usage
Run the script with:
The script can be run in one of two modes: `server` or `batch`.
#### Batch Mode
Simply runs all actions sequentially
#### Server Mode
`python ./api_requester.py --mode server --config /mount/config.yaml`
Runs a server, allowing scripts to call API endpoints, causing this pod to make API requests to other pods.
See endpoints under `def create_app` in `api_requester.py` for usage details.
### Config Format
The ConfigMap in `config.yaml` defines the config objects.
Class definitions are in `configs.py`.
The idea is that a user can define various pieces of the config,
then combine them as needed. Each object stands independently, but
they work together when running an action or making a request.
Each config object has a name by which it can be referenced.
Some fields are optional.
Endpoints - Defines an API endpoint for a request.
Targets - Defines a set of filters to use to determine if pods on a cluster are part of the target.
Requests - Contains an Endpoint and some additional information for retries and delays.
Actions - Combines Targets and Requests into a defined action, representing a series of requests.
#### How an Action is performed
1. For each ConfigTarget, add all pods to the list "all_pods". Note: No deduplication is done.
2. Sort the list of pods according to `order`.
3. Starting at `pod_start_index`, take `pod_count` pods. Note: Loop through the list if needed to get `pod_count` items.
4. According to `loop_order`, make every request in `requests` to every pod in the remaining list.
See docstrings in the `ConfigAction` class for more details.
### Files
```
python script.py [-c CONTENT_TOPIC] [-p PUBSUB_TOPIC] [-ps PAGE_SIZE] [-cs CURSOR]
api_requester.py Main code that will run on the pod-api-requester pod
bind.yaml Necessary Kubernetes Role + RoleBinding to give permission to list pods
config.yaml Kubernetes config containing the definitions for Targets, Endpoints, Requests, and Actions
build.sh Sample commands to build the Docker container
Dockerfile File to build the pod-api-requester image
deployment.yaml Sample pod for development/testing
client.py Sample code to make API requests directed at a pod running this code
```
### Arguments
- `-c`, `--contentTopics` (default: `/my-app/1/dst/proto`): Content topic to query.
- `-p`, `--pubsubTopic` (default: `/waku/2/rs/2/0`): Pubsub topic.
- `-ps`, `--pageSize` (default: 60): Number of messages per request.
- `-cs`, `--cursor` (optional): Cursor for pagination.
### Example in Kubernetes yaml
```
containers:
- name: container
image: <your-registry>/get_store_messages:v1.0.0
imagePullPolicy: IfNotPresent
command:
- sh
- -c
- python /app/store_msg_retriever.py --contentTopics=/my-app/1/dst/proto
```
### How It Works
Queries a random waku node by selecting a random ip from `"zerotesting-service:8645"`.
It keeps querying that node until all messages are retrieved.
### Changelog
- `v2.0.0`:
- Changed to using a ConfigMap to define Targets, Endpoints, Requests, and Actions
- Added server capability
- Removed --debug mode logic
- `v1.0.1`:
- Added `--debug` mode. Makes multiple API requests to each IP
- Added `--select-types` mode
api_requester.py
+327
View File
@@ -0,0 +1,327 @@
import argparse
import random
import traceback
from argparse import Namespace
from collections import defaultdict
from pathlib import Path
from typing import Dict, List
import requests
import uvicorn
import yaml
from fastapi import Depends, FastAPI, HTTPException
from kubernetes import client, config
from kubernetes.client.models.v1_pod import V1Pod
from pydantic import BaseModel, ConfigDict
from configs import ConfigAction, ConfigEndpoint, ConfigRequest, ConfigTarget
from utils import paged_request, setup_logger
logger = setup_logger(__file__)
app = FastAPI()
class TargetPodInfo(BaseModel):
model_config = ConfigDict(arbitrary_types_allowed=True)
config_target: ConfigTarget
pod: V1Pod
@property
def pod_name(self) -> str:
return self.pod.metadata.name
def do_request(request: ConfigRequest, pod_info: TargetPodInfo):
raise NotImplementedError()
def call_endpoint(endpoint: ConfigEndpoint, pod_info: TargetPodInfo) -> dict:
result_data = {"request": {"configEndpoint": endpoint}}
request_data = {"params": endpoint.params, "headers": endpoint.headers}
try:
request_data["url"] = endpoint.url.format(
node=pod_info.pod.status.pod_ip, port=pod_info.config_target.port
)
request_data["pod"] = f"{pod_info.pod.metadata.name}"
logger.info(f"request_data: {request_data}")
if endpoint.paged:
if endpoint.type != "GET":
raise NotImplementedError("Paged requests only implemented for GET requests.")
result = paged_request(request=request_data, max_attempts=1, page_request_delay=0)
else:
if endpoint.type == "POST":
result = requests.post(
request_data["url"],
json=request_data["params"],
headers=request_data["headers"],
)
elif endpoint.type == "GET":
result = requests.post(
request_data["url"],
json=request_data["params"],
headers=request_data["headers"],
)
else:
raise AttributeError(f"Unknown request type. request: `{endpoint}`")
result_data["request"].update(request_data)
result_data["response"] = {"status_code": result.status_code, "text": result.text}
except Exception as e:
error = traceback.format_exc()
logger.error(
f"Exception attempting API request. endpoint: `{endpoint}`, exception: `{e}`, error: `{error}`"
)
result_data["exception"] = error
logger.info(result_data)
return result_data
def assert_unique_attr(objects: List[object], attribute: str):
names = [getattr(endpoint, attribute) for endpoint in objects]
duplicates = set()
seen = []
for name in names:
if any(name == item for item in seen):
duplicates.add(name)
else:
seen.append(name)
assert not duplicates, (
f"At least one object has the same attribute as another. "
f"Attribute name: `{attribute}`. "
f"Duplicate attributes: `{duplicates}`. "
f"Objects: `{objects}`"
)
def get_pods_for_target(target: ConfigTarget) -> List[str]:
config.load_incluster_config()
v1 = client.CoreV1Api()
namespace = open("/var/run/secrets/kubernetes.io/serviceaccount/namespace").read() or "default"
if target.service is not None:
pods = v1.list_namespaced_pod(namespace)
else:
service = v1.read_namespaced_service(target.service, namespace)
selector = service.spec.selector
selector_str = ",".join([f"{k}={v}" for k, v in selector.items()])
pods = v1.list_namespaced_pod(namespace, label_selector=selector_str)
return list(filter(lambda pod: target.matches(pod), pods.items))
def parse_config(config: Dict[str, List[object]]) -> Dict[str, Dict[str, object]]:
targets = [ConfigTarget.model_validate(targ) for targ in config.get("targets", [])]
targets_dict = {target.name: target for target in targets}
assert_unique_attr(targets, "name")
endpoints = [ConfigEndpoint.model_validate(endpoint) for endpoint in config["endpoints"]]
endpoints_dict = {endpoint.name: endpoint for endpoint in endpoints}
assert_unique_attr(endpoints, "name")
requests = []
for request_dict in config["requests"]:
request_dict["endpoint"] = endpoints_dict[request_dict["endpoint"]]
requests.append(ConfigRequest.model_validate(request_dict))
requests_dict = {request.name: request for request in requests}
assert_unique_attr(requests, "name")
actions: List[ConfigAction] = []
for action in config["actions"]:
try:
action["requests"] = [requests_dict[req] for req in action["requests"]]
except KeyError as e:
raise ValueError(
f"Action contains unknown request. action: `{action}` requests: `{requests_dict}`"
) from e
try:
action["targets"] = [targets_dict[targ] for targ in action["targets"]]
except KeyError as e:
raise ValueError(
f"Action contains unknown target. action: `{action}` targets: `{targets_dict}`"
) from e
actions.append(ConfigAction.model_validate(action))
actions_dict = {action.name: action for action in actions}
assert_unique_attr(actions, "name")
return {
"targets": targets_dict,
"endpoints": endpoints_dict,
"requests": requests_dict,
"actions": actions_dict,
}
def load_configs(config_files: List[str]) -> Dict[str, Dict[str, object]]:
logger.info(f"Loading configs: {config_files}")
full_config = defaultdict(list)
for config_file in config_files:
with open(config_file, "r") as file:
config = yaml.safe_load(file)
for key, value in config.items():
full_config[key].extend(value)
return parse_config(full_config)
def do_action(
action: ConfigAction,
pods: List[TargetPodInfo],
) -> List[TargetPodInfo]:
target_names = [target.name for target in action.targets]
possible_pods = [pod for pod in pods if pod.config_target.name in target_names]
if action.order == "random":
random.shuffle(possible_pods)
elif action.order == "ascending":
possible_pods.sort(key=lambda pod: pod.pod_name)
elif action.order == "descending":
possible_pods.sort(key=lambda pod: pod.pod_name, reverse=True)
else:
raise ValueError(f"Unknown order for action: {action.order}")
pods = []
count = len(possible_pods) if action.pod_count == "all" else action.pod_count
index = action.pod_start_index
for _ in range(count):
pods.append(possible_pods[index])
index = (index + 1) % len(possible_pods)
if action.loop_order == "foreach_pod_make_all_requests":
for pod in pods:
for request in action.requests:
# time.sleep(delay_between_requests) TODO
call_endpoint(request, pod)
elif action.loop_order == "foreach_request_target_each_pod":
for request in action.requests:
# TODO: ensure time between requests has elapsed
for pod in pods:
call_endpoint(request, pod)
else:
raise ValueError(f"Unknown loop_order for action: {action}")
def get_pod_infos(targets: List[ConfigTarget]) -> List[TargetPodInfo]:
pods_info: List[TargetPodInfo] = []
for target in targets:
pods = get_pods_for_target(target)
for pod in pods:
pods_info.append(TargetPodInfo(config_target=target, pod=pod))
return pods_info
def create_app(config) -> FastAPI:
app = FastAPI()
async def get_config():
logger.info(f"todo get config: {config}")
return config
class InvokeRequestData(BaseModel):
target: ConfigTarget | str
endpoint: ConfigEndpoint | str
@app.post("/process")
# TODO: Implement try/catch return error in decorator. It will be the same for all endpoints.
def process_data(data: InvokeRequestData, config=Depends(get_config)):
"""
Performs an API request to the given endpoint on the given target.
:param data: Contains target and endpoint.
For each, the argument may either the name from the config,
or a custom object passed in as a dict.
Sample usage (from outside the cluster):
data = {
"target": {
"name": "dummy",
"service": "zerotesting-lightpush-client",
"name_template": "lpclient-0-0",
},
"endpoint": "lightpush-publish-static-sharding",
}
url = f"http://{external_ip}:{node_port}/process"
response = requests.post(url, json=data)
"""
try:
try:
# Treat target as the name of a preset target from config.
target = config["targets"][data.target]
except TypeError:
# If no target with that name exists, treat as custom target.
target = data.target
try:
# Treat endpoint as the name of preset endpoint from config.
endpoint = config["endpoints"][data.endpoint]
except TypeError:
# If no endpoint with that name exists, treat as custom endpoint.
endpoint = data.endpoint
request = ConfigRequest(
name="dummy_request", endpoint=endpoint, retries=0, retry_delay=0
)
pod_info = next(iter(get_pod_infos([target])))
result = call_endpoint(request.endpoint, pod_info)
return result
except Exception as e:
# TODO: Add hints to errors (eg. Action doesn't exist, etc)
logger.error(HTTPException(status_code=500, detail=f"{e!r}\n{traceback.format_exc()}"))
raise HTTPException(status_code=500, detail=f"{e!r}\n{traceback.format_exc()}")
return app
def main(args: Namespace):
config = load_configs(args.config_files)
available_endpoints = [endpoint.name for endpoint in config["endpoints"].values()]
logger.debug(f"Loaded config. Available endpoints: {available_endpoints}")
if args.mode == "server":
app = create_app(config)
uvicorn.run(app, host="0.0.0.0", port=args.port, log_config=None)
else:
pods_info = get_pod_infos(config["targets"])
for action in config["actions"]:
do_action(action, pods_info)
def mode_type(value):
if value not in ["batch", "server"]:
raise argparse.ArgumentTypeError(f"Invalid mode: {value}. Must be 'batch' or 'server'.")
return value
def parse_args() -> argparse.Namespace:
parser = argparse.ArgumentParser(description="Waku storage retriever")
parser.add_argument(
"--config",
type=Path,
action="append",
dest="config_files",
required=True,
help="Paths to config files. Can be passed multiple times.",
)
parser.add_argument(
"--mode",
type=mode_type,
default="server",
help="Batch: Run actions immediately. Server: Wait for API calls to /action/<myaction> to run.",
)
parser.add_argument(
"--port",
type=int,
default=8645,
help="Port for the action HTTP server (default 8000)",
)
args = parser.parse_args()
return args
if __name__ == "__main__":
args = parse_args()
main(args)
bind.yaml
+24
View File
@@ -0,0 +1,24 @@
apiVersion: rbac.authorization.k8s.io/v1
kind: Role
metadata:
namespace: zerotesting
name: pod-service-reader
rules:
- apiGroups: [""]
resources: ["pods", "services"]
verbs: ["get", "list", "watch"]
---
apiVersion: rbac.authorization.k8s.io/v1
kind: RoleBinding
metadata:
name: pod-service-reader-binding
namespace: zerotesting
subjects:
- kind: ServiceAccount
name: default # service account your pod uses
namespace: zerotesting
roleRef:
kind: Role
name: pod-service-reader
apiGroup: rbac.authorization.k8s.io
build.sh
+8 -2
View File
@@ -1,2 +1,8 @@
docker build -t <your-registry>/get_store_messages:v1.0.0 .
docker push <your-registry>/get_store_messages:v1.0.0
docker build -t <your-registry>/pod-api-requester:<tag> --target debug .
# or
docker build -t <your-registry>/pod-api-requester:<tag> --target production .
# The default build is production
docker build -t <your-registry>/pod-api-requester:<tag>
docker push <your-registry>/pod-api-requester:<tag>
client.py
+136
View File
@@ -0,0 +1,136 @@
import asyncio
import json
from typing import Literal, Optional
import requests
from kubernetes import client, config
from pydantic import NonNegativeInt
from utils import setup_logger
logger = setup_logger(__file__)
async def main():
raise NotImplementedError("Choose your Kubernetes config path and remove this.")
config.load_kube_config("/path_to_kube_config.yaml")
config.load_kube_config() # WARNING! LOCAL
publish_message(
namespace="zerotesting",
message_type="lightpush",
pod_name_template="lpclient-0-0",
service="zerotesting-lightpush-client",
)
PublishType = Literal["lightpush", "relay"]
async def publish_message(
namespace: str,
message_type: PublishType,
*,
pod_name_template: Optional[str] = None,
service: Optional[str] = None,
stateful_set_name: Optional[str] = None,
port: NonNegativeInt = 80,
):
if message_type == "lightpush":
endpoint = "lightpush-publish-static-sharding"
elif message_type == "relay":
raise NotImplementedError()
else:
raise ValueError("Unknown message type")
data = {
"target": {
"name": "dummy",
"service": service,
"name_template": pod_name_template,
"stateful_set": stateful_set_name,
"port": port,
},
"endpoint": endpoint,
}
return await pod_api_request(
namespace=namespace,
service_name="zerotesting-publisher",
app="zerotenkay-publisher",
data=data,
)
class PodApiRequestError(Exception):
pass
async def pod_api_request(
namespace: str,
service_name: str,
app: str,
data: dict,
*,
publisher_pod: str | NonNegativeInt = 0,
) -> dict:
v1 = client.CoreV1Api()
try:
pods = v1.list_namespaced_pod(namespace=namespace, label_selector=f"app={app}")
if isinstance(publisher_pod, str):
pod = next(pod for pod in pods.items if pod.metadata.name == publisher_pod)
else:
pod = pods.items[publisher_pod]
except IndexError as e:
logger.error(f"No pod found. app: `{app}` pod_index: `{publisher_pod}`")
raise ValueError() from e
except StopIteration as e:
logger.error(f"No pod found. app: `{app}` pod_name: `{publisher_pod}`")
raise ValueError() from e
# Get publisher IP.
node = v1.read_node(name=pod.spec.node_name)
target_ip = kube_utils.get_node_ip(node)
# Get publisher port.
service = v1.read_namespaced_service(service_name, namespace)
node_port = service.spec.ports[0].node_port
if node_port is None:
raise ValueError(f"Failed to find port for service. Service: `{service.metadata.name}`")
url = f"http://{target_ip}:{node_port}/process"
logger.info(f"publishing message. url: `{url}` data: `{data}`")
response = requests.post(url, json=data)
response_obj = json.loads(response.text)
if response.status_code != 200:
err = response_obj["detail"].replace("\n", "\n")
logger.error(err)
raise PodApiRequestError(response_obj)
try:
# Assuming that the pod we made the API request to returns a response with a JSON object.
inner_response_obj = json.loads(response_obj["response"]["text"])
response_obj["inner_response"] = inner_response_obj
if response_obj["response"]["status_code"] != 200:
# JsWaku puts the error under the key "error".
try:
err = inner_response_obj["error"].replace("\n", "\n")
except KeyError as e:
err = "<Failed to extract inner error>"
logger.error(f"Publisher request returned failure. inner_error: `{err}`")
raise PodApiRequestError(response_obj)
except json.JSONDecodeError as e:
# Response was not a Json object.
pass
except KeyError as e:
err = response_obj["exception"].replace("\n", "\n")
logger.error(f"The publisher's API request attempt failed. Exception: `{err}`")
raise PodApiRequestError(response_obj) from e
logger.info(f"Response: `{response_obj}`")
return response_obj
if __name__ == "__main__":
asyncio.run(main())
config.yaml
+62
View File
@@ -0,0 +1,62 @@
apiVersion: v1
kind: ConfigMap
metadata:
name: api-requester-config
namespace: zerotesting
data:
config.yaml: |
targets:
- name: firstClients
service: zerotesting-lightpush-client
name_template: "^client-0-(1[0-2]|[0-9])$"
stateful_set: "client-0"
- name: jswakuClients
service: zerotesting-lightpush-client
stateful_set: "client-0"
port: 8080
- name: nwakuClients
service: zerotesting-lightpush-client
stateful_set: "client-0"
port: 8645
endpoints:
- name: set-debug
headers: {"accept": "text/plain"}
params: {"logLevel": "DEBUG"}
url: "http://{node}/admin/v1/log-level/DEBUG"
type: "POST"
paged: False
- name: lightpush-publish-static-sharding
url: "http://{node}:{port}/lightpush/v3/message"
headers: {"Content-Type": "application/json"}
params: {"pubsubTopic": "/waku/2/rs/2/0", "message": {"contentTopic": "/test/1/cross-network/proto", "payload": "W1EsIDIsIDNd"}}
type: "POST"
paged: False
- name: lightpush-publish-auto-sharding
url: "http://{node}:{port}/lightpush/v3/message"
headers: {"Content-Type": "application/json"}
params: { "pubsubTopic":"", "message": {"contentTopic": "/test/1/cross-network/proto", "payload" : "W1EsIDIsIDNd"}}
type: "POST"
paged: False
requests:
- name: publish-request
endpoint: lightpush-publish-static-sharding
retries: 0
retry_delay: 0.3
actions:
- name: publish-to-clients-random
requests: ["publish-request"]
targets: ["nwakuClients"]
pod_start_index: 0
pod_count: 91
order: ascending
loop_order: foreach_request_target_each_pod
configs.py
+222
View File
@@ -0,0 +1,222 @@
import datetime
import logging
import re
from typing import List, Literal, Optional
from kubernetes import client
from kubernetes.client.models.v1_pod import V1Pod
from pydantic import BaseModel, NonNegativeFloat, NonNegativeInt, PositiveInt
class UTCFormatter(logging.Formatter):
def formatTime(self, record, datefmt=None):
# Get UTC time and format with milliseconds
utc_dt = datetime.datetime.utcfromtimestamp(record.created)
if datefmt:
s = utc_dt.strftime(datefmt)
# Add milliseconds
s = s + f".{int(record.msecs):03d}"
return s
else:
t = utc_dt.strftime("%Y-%m-%d %H:%M:%S")
s = f"{t}.{int(record.msecs):03d}"
return s
logfmt = "%(asctime)s %(levelname)s [%(name)s] %(message)s"
datefmt = "%Y-%m-%d %H:%M:%S"
handler = logging.StreamHandler()
handler.setFormatter(UTCFormatter(logfmt, datefmt=datefmt))
logging.basicConfig(level=logging.INFO, handlers=[handler])
logger = logging.getLogger(__file__)
class ConfigEndpoint(BaseModel):
"""Describes an endpoint on a pod in the cluster.
This endpoint may exist on multiple pods, or just a single pod, or no pod at all.
It is the responsibility of the caller to combine a defined endpoint with a proper pod."""
name: str
"""The name of this config object."""
headers: dict
"""
The header to send with the request.
Typically either
headers: {"accept": "text/plain"}
or
{"Content-Type": "application/json"}
"""
params: dict
"""HTTP POST data to include with request."""
url: str
"""Url for the endpoint.
Instances of `{node}` and `{port}` will be replaced with
the pod IP and the target port respectively.
For example, when calling the following endpoint:
`http://{node}:{port}/lightpush/v3/message`
on the node at index `2` of the a target with:
`stateful_set: "client", port: 8645`, then the following url will be used:
`http://client-2:8645/lightpush/v3/message`
"""
type: Literal["POST", "GET"]
"""Specifies the method of the request. Either `POST` or `GET`."""
paged: bool
"""Use `True` if the request returns paged data. Otherwise, use `False`."""
class ConfigRequest(BaseModel):
"""A request to be made to a pod.
Contains the `endpoint` and some additional data for retries/delays."""
name: str
"""The name of this config object."""
endpoint: ConfigEndpoint
"""The Endpoint to use for this request."""
retries: NonNegativeInt
"""Number of times to retry the request if it fails."""
retry_delay: NonNegativeFloat
"""The delay between each retry attempt for this request."""
class ConfigTarget(BaseModel):
"""A config describing pods.
This is a list of filters to apply to any pod
to see if that pod is part of the target group.
"""
name: str
"""The name of this config object. Not the pod name."""
service: Optional[str] = None
"""The name of the service that any target pod must belong.
Example: zerotesting-bootstrap"""
name_template: Optional[str] = None
"""Regex describing the pod names. Example: ^client-([0-9])$"""
stateful_set: Optional[str] = None
"""Name of the StatefulSet that any target pod must belong to.
Example: "bootstrap"
"""
port: NonNegativeInt = 80
"""Port to use for requests to endpoints with this target.
Default is 80."""
def matches(self, pod: V1Pod) -> bool:
"""Check if pod is a valid target of self"""
if self.stateful_set is not None:
if pod.metadata.owner_references is None:
return False
if not all(
[
owner.kind == "StatefulSet" and owner.name == self.stateful_set
for owner in pod.metadata.owner_references
]
):
return False
if self.name_template is not None:
if not re.search(self.name_template, pod.metadata.name):
return False
if self.service is not None:
v1 = client.CoreV1Api()
namespace = (
open("/var/run/secrets/kubernetes.io/serviceaccount/namespace").read() or "default"
)
service = v1.read_namespaced_service(self.service, namespace)
selector = service.spec.selector
if not all([pod.metadata.labels.get(key) == value for key, value in selector.items()]):
return False
return True
class ConfigAction(BaseModel):
"""Description of an action to take. Here is how an action is performed:
1. For each ConfigTarget, add all pods to the list.
2. Sort the list of pods according to `order`.
3. Starting at `pod_start_index`, take `pod_count` pods.
4. According to `loop_order`, make every request in `requests` to every pod in the remaining list.
"""
name: str
"""The name of this config object."""
loop_order: Literal["foreach_pod_make_all_requests", "foreach_request_target_each_pod"]
"""Which algorithm to use to determine how requests should be made to pods.
`foreach_pod_make_all_requests`: Loop through the list of pods.
At each pod, make all the requests in `requests`
`foreach_request_target_each_pod`: Loop through the `requests` list.
For each `ConfigRequest`, execute that request on all pods in the
list of pods derived from the algorithm described above.
"""
pod_start_index: NonNegativeInt = 0
"""Allows a user to "skip" a certain amount of pods.
This is applied to a list created by combining the lists of pods from `targets`,
and sorting the list according to `order`.
Assumes that `pod_start_index < len(all_pods)`.
"""
pod_count: PositiveInt | Literal["all"] = "all"
"""The number of pods for this action.
This can be used to limit the total number of pods considered for requests.
Like, `pod_start_index`, this applies to the list of pods created via combining
pods from `targets` and sorting them.
If `pod_count` is `"all"`, then all pods will be used. This will not deduplicate any pods in the list.
If `pod_count > len(all_pods)`, then the cursor will loop back to the beginning of the list and continue
adding pods until the list of pods to use has exactly `pod_count` elements in it.
"""
order: Literal["ascending", "descending", "random"] | None
"""Once the list of possible pods is gathered by combining the lists of pods for each `ConfigTarget`,
they will be sorted by this ordering before applying `pod_start_index` and `pod_count`.
"""
targets: List[ConfigTarget]
"""A list of all `ConfigTarget`s used to gather the list of pods.
For each target, all pods will be added to as potential targets.
Then the list will be sorted according to `order`, and spliced
according to `pod_start_index` and `pod_count`.
Note: A pod may match multiple `ConfigTarget`s. In this case,
the pod will be added to the list as many times as it matches.
For example, with `ConfigTarget`s {"stateful_set": "some_pod"} and {"name_template": "^some_pod-[1-2]$"},
where the StatefulSet of some_pod has `replicas: 4`, the list of pods to use would be:
["some_pod-0", "some_pod-1", "some_pod-2", "some_pod-3", "some_pod-1", "some_pod-2"]
which then may be sorted by `ascending` to look like:
["some_pod-0", "some_pod-1", "some_pod-1", "some_pod-2", "some_pod-2", "some_pod-3"]
then, the list would be spliced using `all_pods[pod_start_index:pod_start_index+pod_count]`,
assuming that `pod_start_index+pod_count < len(all_pods)`.
"""
requests: List[ConfigRequest]
"""The list of requests to do for each pod that ends up in the list of pods to request to.
Every request will be executed, but it may not be on all pods matching every `ConfigTarget` in `targets`. See `targets`."""
deployment.yaml
+43
View File
@@ -0,0 +1,43 @@
apiVersion: v1
kind: Pod
metadata:
name: publisher
namespace: zerotesting
labels:
app: zerotenkay-publisher
spec:
restartPolicy: Never
dnsConfig:
searches:
- zerotesting-publisher.zerotesting.svc.cluster.local
volumes:
- name: api-requester-config-volume
configMap:
name: api-requester-config
containers:
- name: publisher-container
image: pearsonwhite/pod-api-requester:1e161cdf41478000bea17f9332f3624e9aad0829
imagePullPolicy: Always
command:
[
"python",
"/app/api_requester.py",
"--mode",
"server",
"--config",
"/mount/config.yaml",
]
ports:
- containerPort: 8645
- containerPort: 8008
- containerPort: 8080
volumeMounts:
- name: api-requester-config-volume
mountPath: /mount
resources:
requests:
memory: 64Mi
cpu: 150m
limits:
memory: 600Mi
cpu: 400m
publisher-service.yaml
+14
View File
@@ -0,0 +1,14 @@
apiVersion: v1
kind: Service
metadata:
name: zerotesting-publisher
namespace: zerotesting-pwhite
spec:
type: NodePort
selector:
app: zerotenkay-publisher
ports:
- protocol: TCP
port: 8000
targetPort: 8645
nodePort: 30080
requirements.txt
+7
View File
@@ -0,0 +1,7 @@
aiohttp==3.9.3
fastapi==0.124.0
kubernetes==27.2.0
pydantic==2.12.5
PyYAML==6.0.3
Requests==2.32.5
uvicorn==0.38.0
store_msg_retriever.py
-639
View File
@@ -1,639 +0,0 @@
# Python Imports
import argparse
import datetime
import json
import logging
import socket
import time
import traceback
from argparse import Namespace
from copy import deepcopy
from pathlib import Path
from typing import Any, Dict, List, Tuple
import requests
from pydantic import BaseModel, Field, PositiveInt
class UTCFormatter(logging.Formatter):
def formatTime(self, record, datefmt=None):
# Get UTC time and format with milliseconds
utc_dt = datetime.datetime.utcfromtimestamp(record.created)
if datefmt:
s = utc_dt.strftime(datefmt)
# Add milliseconds
s = s + f".{int(record.msecs):03d}"
return s
else:
t = utc_dt.strftime("%Y-%m-%d %H:%M:%S")
s = f"{t}.{int(record.msecs):03d}"
return s
# Usage
logfmt = "%(asctime)s %(levelname)s [%(name)s] %(message)s"
datefmt = "%Y-%m-%d %H:%M:%S"
handler = logging.StreamHandler()
handler.setFormatter(UTCFormatter(logfmt, datefmt=datefmt))
logging.basicConfig(level=logging.INFO, handlers=[handler])
logger = logging.getLogger(__file__)
def next_cursor(data: Dict) -> str | None:
cursor = data.get("paginationCursor")
if not cursor:
logger.info("No more messages")
return None
return cursor
def fetch_all_messages(base_url: str, initial_params: Dict, headers: Dict) -> List[str]:
all_messages = []
params = initial_params.copy()
while True:
logger.info(
f"requests.get: url: `{base_url}` init_params: `{initial_params}` params: `{params}`"
)
response = requests.get(base_url, headers=headers, params=params)
logger.info(f"response: `{response.text}`")
if response.status_code != 200:
logger.error(f"Error fetching data: {response.status_code}")
logger.error(response.text)
break
data = response.json()
logger.info(data)
if data["statusCode"] != 200:
logger.info(f"failed. statusCode: `{data['statusCode']}`")
paged_messages = [message["messageHash"] for message in data["messages"]]
logger.info(f"Retrieved {len(paged_messages)} messages")
all_messages.extend([message["messageHash"] for message in data["messages"]])
cursor = next_cursor(data)
if not cursor:
break
params["cursor"] = cursor
return all_messages
def dict_extract(obj: dict, path: Path):
def extract(obj: Any, parts: list, is_list=False):
if isinstance(obj, list):
results = []
for item in obj:
results.extend(extract(item, parts, is_list=True))
return results
if not parts:
return [obj] if is_list else obj
next_obj = obj[parts[0]]
return extract(next_obj, parts[1:], is_list)
return extract(obj, path.parts)
def paged_request(request: dict, max_attempts: PositiveInt, page_request_delay: float) -> dict:
"""
GET request with a "paged" param.
:param request: Must contain "params":dict.
"""
attempt_num = 1
url = request["url"]
all_messages = []
pages_data = []
params = request["params"]
status_codes = []
inner_status_codes = []
while True:
time.sleep(page_request_delay)
logger.info(f"Making paged request. request: `{request}`, params=`{params}`")
response = requests.get(url, headers=request["headers"], params=params)
try:
data = response.json()
except requests.exceptions.JSONDecodeError:
data = response.text
status_codes.append(response.status_code)
pages_data.append(data)
logger.info(f"response to paged request: `{response}`")
if response.status_code != 200:
logger.error(
f"Error fetching paged data. status_code: `{response.status_code}` data: `{data}`"
)
break
inner_status_codes.append(data["statusCode"])
logger.info(f"Response data: `{data}`")
if data["statusCode"] != 200:
logger.info(
f"inner_status_code != 200: status_code: `{data['statusCode']}`, attempt: `{attempt_num}`"
)
if attempt_num >= max_attempts:
logger.info(f"Exhausted all attempts: `{attempt_num}`")
break
attempt_num += 1
continue
logger.info(f"inner_status_code == 200: attempt: `{attempt_num}`")
if attempt_num > 1:
logger.info("A previous attempt failed, but now it worked.")
paged_data = dict_extract(data, request.get("extract_keys", Path()))
logger.info(f"Retrieved {len(paged_data)} messages on attempt `{attempt_num}`")
all_messages.extend(paged_data)
cursor = next_cursor(data)
if not cursor:
logger.info(f"page request finished with !cursor on attempt `{attempt_num}`")
break
params["cursor"] = cursor
attempt_num = 1
logger.info("finished page request")
return {
"request": request,
"response": {
"statusCodes": status_codes,
"inner_statusCodes": inner_status_codes,
"messages": all_messages,
"pages": pages_data,
"attempt_num": attempt_num,
},
}
def api_request(action, request) -> dict:
url = request["url"]
response = action(url, request["headers"], request.get("params"))
try:
data = response.json()
except requests.exceptions.JSONDecodeError:
data = response.text
if response.status_code != 200:
logger.error(f"Error fetching data: {response.status_code}")
logger.error(data)
return {
"request": request,
"response": {
"statusCode": response.status_code,
"contents": data,
},
}
def get_node_info(
name: str, node: str, api_args: dict, delay_between_requests=0.3
) -> Dict[str, dict]:
all_requests = {
"debug": {
"headers": {"accept": "text/plain"},
"params": {"logLevel": "DEBUG"},
"url": "http://{node}/admin/v1/log-level/DEBUG",
"type": "POST",
},
"info": {
"url": "http://{node}/debug/v1/info",
"headers": {"accept": "application/json"},
"type": "GET",
},
"peers": {
"url": "http://{node}/admin/v1/peers",
"headers": {"accept": "application/json"},
"type": "GET",
},
"mesh": {
"url": "http://{node}/admin/v1/peers/mesh",
"headers": {"accept": "application/json"},
"type": "GET",
},
"stats": {
"url": "http://{node}/admin/v1/peers/stats",
"headers": {"accept": "application/json"},
"type": "GET",
},
"connected": {
"url": "http://{node}/admin/v1/peers/connected",
"headers": {"accept": "application/json"},
"type": "GET",
},
"service": {
"url": "http://{node}/admin/v1/peers/service",
"headers": {"accept": "application/json"},
"type": "GET",
},
"store_messages": {
"url": f"http://{node}/store/v3/messages",
"headers": {"accept": "application/json"},
"paged": True,
"params": api_args,
"extract_keys": Path("messages", "messageHash"),
},
}
request_data = {}
for key, node_request in all_requests.items():
request = deepcopy(node_request)
request["url"] = request["url"].format(node=node)
request["node"] = name
if request.get("type") == "POST":
action = lambda url, headers, params: requests.post(url, data=params, headers=headers)
elif request.get("type") == "GET":
action = lambda url, headers, params: requests.get(url, headers=headers, params=params)
try:
if request.get("paged"):
result = paged_request(request=request, max_attempts=1, page_request_delay=0)
else:
result = api_request(action, request)
request_data[key] = result
except Exception as e:
error = traceback.format_exc()
logger.error(
f"Exception attempting API request. request: `{request}`, exception: `{e}`, error: `{error}`"
)
request_data[key] = {
"request": request,
"exception": error,
}
time.sleep(delay_between_requests)
return request_data
def serializer(obj):
if isinstance(obj, Path):
return obj.as_posix()
if isinstance(obj, set):
return list(obj)
raise TypeError(f"Object of type {obj.__class__.__name__} is not JSON serializable")
def resolve_dns(node: str) -> Tuple[str, str]:
start_time = time.time()
name, port = node.split(":")
ip_address = socket.gethostbyname(name)
entire_hostname = socket.gethostbyaddr(ip_address)
hostname = entire_hostname[0].split(".")[0]
elapsed = (time.time() - start_time) * 1000
logger.info(f"{node} DNS Response took {elapsed} ms")
logger.info(f"Talking with {hostname}, ip address: {ip_address}")
return (entire_hostname, f"{ip_address}:{port}")
class NodeType(BaseModel):
name_template: str
"""Format string for node name. Eg. fserver-0-{index}"""
service: str
count_key: str
namespace: str = Field(default="zerotesting")
def dns_name(self, index: PositiveInt) -> str:
"""Return name for DNS lookup.
<pod-name>.<headless-service-name>
"""
return f"{self.get_node_name(index)}.{self.service}"
def get_node_name(self, index: PositiveInt) -> str:
return self.name_template.format(index=index)
node_types = [
NodeType(
name_template="store-0-{index}",
service="zerotesting-store",
count_key="store",
),
NodeType(
# Note the plural "nodes" with an 's'!
# This is to match the name used in regression tests.
name_template="nodes-0-{index}",
service="zerotesting-service",
count_key="relay",
),
NodeType(
name_template="fserver-0-{index}",
service="zerotesting-filter",
count_key="filter_server",
),
NodeType(
name_template="fclient-0-{index}",
service="zerotesting-filter",
count_key="filter_client",
),
NodeType(
name_template="lpserver-0-{index}",
service="zerotesting-lightpush-server",
count_key="lightpush_server",
),
NodeType(
name_template="lpclient-0-{index}",
service="zerotesting-lightpush-client",
count_key="lightpush_client",
),
NodeType(
name_template="bootstrap-{index}",
service="zerotesting-bootstrap",
count_key="bootstrap",
),
]
def get_ips_by_type(args: dict, *, namespace=None) -> List[Tuple[str, str]]:
"""
Get node ips based on type flags (--store, --relay, etc) starting at start_index for each node type.
:return: (name, ip) tuples for node specified.
:rtype: List[str, str]
"""
# TODO: Handle multiple shards.
results = []
for node_type in node_types:
start_index = args.get("start_index", 0)
if args[node_type.count_key] == "all":
try:
_, _, ip_list = socket.gethostbyname_ex(node_type.service)
count = len(ip_list) - start_index
except socket.gaierror:
# This happens when either:
# 1. The service doesn't exist.
# 2. No pods with the matching app selector exist, thus though the service exists, it isn't running on any pod.
count = 0
# TODO: Check at the end if all `count` ips have been found.
# TODO: Add "unknown-{index}" for ips not in {nodetype}-0-{index}
# Note that if node types share the same service, count will be set to the total.
# for example fserver/fclient both use zerotesting-filter.
else:
try:
count = int(args[node_type.count_key])
except (KeyError, TypeError):
logger.info(f"No count for nodetype specified. `{node_type}`")
continue
logger.info(
f"Getting {count} IPs from nodes of type `{node_type.name_template}` starting at index {start_index}"
)
for index in range(start_index, start_index + count):
dns = node_type.dns_name(index)
try:
_, _, ips = socket.gethostbyname_ex(dns)
results.append((node_type.get_node_name(index), ips[0]))
except Exception as e:
error = traceback.format_exc()
logger.error(
f"Failed to resolve dns. dns: `{dns}`, node_type: `{node_type}`, exception: `{e}`, error: {error}"
)
return results
def get_api_args(args_dict: dict) -> dict:
"""These are the arguments that should be passed on to the GET request for store messages."""
return {
key: value
for key, value in args_dict.items()
if key
in [
"contentTopics",
"pubsubTopic",
"pageSize",
"cursor",
]
}
def positive_int_or_all(value):
if value == "all":
return value
try:
int_value = int(value)
assert int_value >= 0
return int_value
except (ValueError, AssertionError):
raise argparse.ArgumentTypeError(f"{value} is not an integer or 'all'.")
def parse_args() -> argparse.Namespace:
parser = argparse.ArgumentParser(description="Waku storage retriever")
parser.add_argument(
"-c", "--contentTopics", type=str, help="Content topic", default="/my-app/1/dst/proto"
)
parser.add_argument(
"-p", "--pubsubTopic", type=str, help="Pubsub topic", default="/waku/2/rs/2/0"
)
parser.add_argument(
"-ps", "--pageSize", type=int, help="Number of messages to retrieve per page", default=60
)
parser.add_argument(
"-cs",
"--cursor",
type=str,
help="Cursor field intended for pagination purposes. ",
default="",
)
parser.add_argument(
"-d",
"--debug",
action="store_true",
help="",
dest="debug",
)
parser.add_argument(
"-rd",
"--request-delay",
type=float,
default=0.3,
help="Delay between each REST API call on a node. Only applicable in --debug mode.",
dest="delay_between_requests",
)
parser.add_argument(
"-t",
"--select-types",
action="store_true",
help="If specified, gathers ips for nodes of the types indicated by additional flags (e.g., --store, --relay). If not specified, selects a random node of any type.",
dest="select_types",
)
parser.add_argument(
"-s",
"--store",
type=positive_int_or_all,
help="Number of store nodes",
dest="store",
)
parser.add_argument(
"-r",
"--relay",
type=positive_int_or_all,
help="Number of plain relay nodes",
dest="relay",
)
parser.add_argument(
"-fs",
"--filter-server",
type=positive_int_or_all,
help="Number of fserver nodes",
dest="filter_server",
)
parser.add_argument(
"-fc",
"--filter-client",
type=positive_int_or_all,
help="Number of fclient nodes",
dest="filter_client",
)
parser.add_argument(
"-lps",
"--lightpush-server",
type=positive_int_or_all,
help="Number of lpserver nodes",
dest="lightpush_server",
)
parser.add_argument(
"-lpc",
"--lightpush-client",
type=positive_int_or_all,
help="Number of lpclient nodes",
dest="lightpush_client",
)
parser.add_argument(
"-bn",
"--bootstrap",
type=positive_int_or_all,
help="Number of bootstrap nodes",
dest="bootstrap",
)
parser.add_argument(
"-si",
"--start-index",
type=int,
default=0,
help="Start looking for at index: {nodetype}-0-{index}",
dest="start_index",
)
args = parser.parse_args()
assert args.select_types == any(
[
args.relay,
args.store,
args.filter_server,
args.filter_client,
args.lightpush_server,
args.lightpush_client,
args.bootstrap,
]
), "--select-types should be True if any node types have been specified and False otherwise."
return args
def main(args: Namespace):
args_dict = vars(args)
api_args = get_api_args(args_dict)
logger.info(f"Arguments: {args_dict}")
nodes = get_ips(args)
messages = []
for index, (name, node) in enumerate(nodes):
try:
logger.info(
f"fetching messages. name: `{name}` url: `{node}` index: {index+1}/{len(nodes)} ({100* (index+1) / len(nodes):.2f}%)"
)
url = f"http://{node}/store/v3/messages"
logger.info(f"Query to {url}")
headers = {"accept": "application/json"}
new_messages = fetch_all_messages(url, api_args, headers)
messages.extend(new_messages)
except Exception as e:
error = traceback.format_exc()
print(f"exception while fetching messages. exception: `{e}`, error: `{error}`")
logger.info("List of messages")
# # We do a print here, so it is easier to parse when reading from victoria logs
print(messages)
def get_ips(args) -> Tuple[str, str]:
port = 8645
if args.select_types:
ips = get_ips_by_type(vars(args))
logger.info(f"ips: ({len(ips)}): ```{ips}```")
return [(name, f"{ip}:{port}") for name, ip in ips]
else:
service = f"zerotesting-service:{port}"
return [resolve_dns(service)]
def main_debug(args: Namespace):
args_dict = vars(args)
api_args = get_api_args(args_dict)
logger.info(f"Arguments: {args_dict}")
nodes = get_ips(args)
for name, node in nodes:
attempt = 1
max_attempts = 10
delay = 0.5
while True:
time.sleep(delay)
try:
logger.info(f"fetching messages. name: `{name}` url: `{node}` attempt: `{attempt}`")
logger.info(f"fetching messages. name: `{name}` url: `{node}`")
node_info = get_node_info(name, node, api_args, args.delay_between_requests)
node_info["attempt"] = attempt
logger.info(
f"store_msg_retriever::node_info: ```{json.dumps(node_info, default=serializer)}```"
)
if all(
code == 200
for code in node_info["store_messages"]["response"]["inner_statusCodes"]
):
logging.info("No inner status failures")
if attempt > 1:
logger.info("main::A previous attempt failed, but now it worked.")
break
logging.info("Inner status failures detected")
if attempt >= max_attempts:
break
attempt += 1
except Exception as e:
error = traceback.format_exc()
logging.error(
f"exception while fetching messages. exception: `{e}`, error: `{error}`"
)
if __name__ == "__main__":
args = parse_args()
if args.debug:
main_debug(args)
else:
main(args)
utils.py
+303
View File
@@ -0,0 +1,303 @@
import datetime
import logging
import socket
import time
import traceback
from pathlib import Path
from typing import Any, Dict, List, Tuple
import requests
from pydantic import BaseModel, Field, PositiveInt
LOGFMT = "%(asctime)s %(levelname)s [%(name)s] %(message)s"
DATEFMT = "%Y-%m-%d %H:%M:%S"
class UTCFormatter(logging.Formatter):
"""Formatter that outputs UTC timestamps with milliseconds."""
def formatTime(self, record, datefmt=None):
dt = datetime.datetime.fromtimestamp(record.created, datetime.timezone.utc)
if datefmt is None:
datefmt = DATEFMT
base = dt.strftime(datefmt)
return f"{base}.{int(record.msecs):03d}"
def setup_logger(name: str) -> logging.Logger:
handler = logging.StreamHandler()
handler.setFormatter(UTCFormatter(LOGFMT, datefmt=DATEFMT))
logging.basicConfig(level=logging.INFO, handlers=[handler], force=True)
return logging.getLogger(name)
logger = setup_logger(__file__)
def get_ips_by_service(service: str) -> List[str]:
try:
_, _, ips = socket.gethostbyname_ex(service)
return ips[0]
except Exception as e:
error = traceback.format_exc()
logger.error(
f"Failed to resolve dns. service: `{service}`, exception: `{e}`, error: {error}"
)
raise
class Target(BaseModel):
pod_name: str
ip: str
service: str
dns_name: str
class NodeType(BaseModel):
name_template: str
"""Format string for node name. Eg. fserver-0-{index}"""
service: str
count_key: str
namespace: str = Field(default="zerotesting")
def dns_name(self, index: PositiveInt) -> str:
"""Return name for DNS lookup.
<pod-name>.<headless-service-name>
"""
return f"{self.get_node_name(index)}.{self.service}"
def get_node_name(self, index: PositiveInt) -> str:
return self.name_template.format(index=index)
node_types = [
NodeType(
name_template="store-0-{index}",
service="zerotesting-store",
count_key="store",
),
NodeType(
# Note the plural "nodes" with an 's'!
# This is to match the name used in regression tests.
name_template="nodes-0-{index}",
service="zerotesting-service",
count_key="relay",
),
NodeType(
name_template="fserver-0-{index}",
service="zerotesting-filter",
count_key="filter_server",
),
NodeType(
name_template="fclient-0-{index}",
service="zerotesting-filter",
count_key="filter_client",
),
NodeType(
name_template="lpserver-0-{index}",
service="zerotesting-lightpush-server",
count_key="lightpush_server",
),
NodeType(
name_template="lpclient-0-{index}",
service="zerotesting-lightpush-client",
count_key="lightpush_client",
),
NodeType(
name_template="bootstrap-{index}",
service="zerotesting-bootstrap",
count_key="bootstrap",
),
]
def get_ips_by_type(args: dict, *, namespace=None) -> List[Tuple[str, str]]:
"""
Get node ips based on type flags (--store, --relay, etc) starting at start_index for each node type.
:return: (name, ip) tuples for node specified.
:rtype: List[str, str]
"""
# TODO: Handle multiple shards.
results = []
for node_type in node_types:
start_index = args.get("start_index", 0)
if args[node_type.count_key] == "all":
try:
_, _, ip_list = socket.gethostbyname_ex(node_type.service)
count = len(ip_list) - start_index
except socket.gaierror:
# This happens when either:
# 1. The service doesn't exist.
# 2. No pods with the matching app selector exist, thus though the service exists, it isn't running on any pod.
count = 0
# TODO: Check at the end if all `count` ips have been found.
# TODO: Add "unknown-{index}" for ips not in {nodetype}-0-{index}
# Note that if node types share the same service, count will be set to the total.
# for example fserver/fclient both use zerotesting-filter.
else:
try:
count = int(args[node_type.count_key])
except (KeyError, TypeError):
logger.info(f"No count for nodetype specified. `{node_type}`")
continue
logger.info(
f"Getting {count} IPs from nodes of type `{node_type.name_template}` starting at index {start_index}"
)
for index in range(start_index, start_index + count):
dns = node_type.dns_name(index)
try:
_, _, ips = socket.gethostbyname_ex(dns)
results.append((node_type.get_node_name(index), ips[0]))
except Exception as e:
error = traceback.format_exc()
logger.error(
f"Failed to resolve dns. dns: `{dns}`, node_type: `{node_type}`, exception: `{e}`, error: {error}"
)
return results
def resolve_dns(node: str) -> Tuple[str, str]:
start_time = time.time()
name, port = node.split(":")
ip_address = socket.gethostbyname(name)
entire_hostname = socket.gethostbyaddr(ip_address)
hostname = entire_hostname[0].split(".")[0]
elapsed = (time.time() - start_time) * 1000
logger.info(f"{node} DNS Response took {elapsed} ms")
logger.info(f"Talking with {hostname}, ip address: {ip_address}")
return (entire_hostname, f"{ip_address}:{port}")
def get_ips(args) -> Tuple[str, str]:
port = 8645
if args.select_types:
ips = get_ips_by_type(vars(args))
logger.info(f"ips: ({len(ips)}): ```{ips}```")
return [(name, f"{ip}:{port}") for name, ip in ips]
else:
service = f"zerotesting-service:{port}"
return [resolve_dns(service)]
# TODO: Extraneous code? (unused)
def get_api_args(args_dict: dict) -> dict:
"""These are the arguments that should be passed on to the GET request for store messages."""
return {
key: value
for key, value in args_dict.items()
if key
in [
"contentTopics",
"pubsubTopic",
"pageSize",
"cursor",
]
}
def dict_extract(obj: dict, path: Path):
def extract(obj: Any, parts: list, is_list=False):
if isinstance(obj, list):
results = []
for item in obj:
results.extend(extract(item, parts, is_list=True))
return results
if not parts:
return [obj] if is_list else obj
next_obj = obj[parts[0]]
return extract(next_obj, parts[1:], is_list)
return extract(obj, path.parts)
def next_cursor(data: Dict) -> str | None:
cursor = data.get("paginationCursor")
if not cursor:
logger.info("No more messages")
return None
return cursor
def paged_request(request: dict, max_attempts: PositiveInt, page_request_delay: float) -> dict:
"""
GET request with a "paged" param.
:param request: Must contain "params":dict.
"""
attempt_num = 1
url = request["url"]
all_messages = []
pages_data = []
params = request["params"]
status_codes = []
inner_status_codes = []
while True:
time.sleep(page_request_delay)
logger.info(f"Making paged request. request: `{request}`, params=`{params}`")
response = requests.get(url, headers=request["headers"], params=params)
try:
data = response.json()
except requests.exceptions.JSONDecodeError:
data = response.text
status_codes.append(response.status_code)
pages_data.append(data)
logger.info(f"response to paged request: `{response}`")
if response.status_code != 200:
logger.error(
f"Error fetching paged data. status_code: `{response.status_code}` data: `{data}`"
)
break
inner_status_codes.append(data["statusCode"])
logger.info(f"Response data: `{data}`")
if data["statusCode"] != 200:
logger.info(
f"inner_status_code != 200: status_code: `{data['statusCode']}`, attempt: `{attempt_num}`"
)
if attempt_num >= max_attempts:
logger.info(f"Exhausted all attempts: `{attempt_num}`")
break
attempt_num += 1
continue
logger.info(f"inner_status_code == 200: attempt: `{attempt_num}`")
if attempt_num > 1:
logger.info("A previous attempt failed, but now it worked.")
paged_data = dict_extract(data, request.get("extract_keys", Path()))
logger.info(f"Retrieved {len(paged_data)} messages on attempt `{attempt_num}`")
all_messages.extend(paged_data)
cursor = next_cursor(data)
if not cursor:
logger.info(f"page request finished with !cursor on attempt `{attempt_num}`")
break
params["cursor"] = cursor
attempt_num = 1
logger.info("finished page request")
return {
"request": request,
"response": {
"statusCodes": status_codes,
"inner_statusCodes": inner_status_codes,
"messages": all_messages,
"pages": pages_data,
"attempt_num": attempt_num,
},
}