Skip to content

Commit

Permalink
fix(dev): add management command to manually send API usage to influx (
Browse files Browse the repository at this point in the history
  • Loading branch information
matthewelwell authored Jun 13, 2024
1 parent 9a6e335 commit 77eeaa7
Showing 1 changed file with 103 additions and 0 deletions.
103 changes: 103 additions & 0 deletions api/app_analytics/management/commands/sendapiusagetoinflux.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,103 @@
import argparse
from datetime import datetime

from django.conf import settings
from django.core.management import BaseCommand
from influxdb_client import InfluxDBClient, Point
from influxdb_client.client.write_api import SYNCHRONOUS


class Command(BaseCommand):
help = "Send API usage to InfluxDB"

def add_arguments(self, parser: argparse.ArgumentParser) -> None:
parser.add_argument(
"--tag",
type=str,
action="append",
default=None,
dest="tags",
required=True,
help="Tag to send with the record in the format 'key=value'. "
"Send multiple tags by providing multiple --tag arguments.",
)
parser.add_argument(
"--bucket",
type=str,
dest="bucket_name",
required=True,
help="Influx bucket name to write to",
)

parser.add_argument(
"--influx-url",
type=str,
default=None,
dest="influx_url",
required=False,
help="Influx API url to use. Defaults to settings.INFLUXDB_URL.",
)
parser.add_argument(
"--influx-token",
type=str,
default=None,
dest="influx_token",
required=False,
help="Influx API token to use. Defaults to settings.INFLUXDB_TOKEN.",
)
parser.add_argument(
"--influx-org",
type=str,
default=None,
dest="influx_org",
required=False,
help="Influx organisation to use. Defaults to settings.INFLUXDB_ORG.",
)
parser.add_argument(
"--request-count",
type=int,
dest="request_count",
required=False,
default=1,
help="Count of requests to send with the data point. Defaults to 1.",
)
parser.add_argument(
"--time",
type=datetime.fromisoformat,
required=False,
default=None,
dest="time",
help="Time to send the data point with. Defaults to current time.",
)

def handle(
self,
*args,
tags: list[str],
bucket_name: str,
influx_url: str,
influx_token: str,
influx_org: str,
request_count: int,
time: datetime,
**kwargs,
) -> None:
influxdb_client = InfluxDBClient(
url=influx_url or settings.INFLUXDB_URL,
token=influx_token or settings.INFLUXDB_TOKEN,
org=influx_org or settings.INFLUXDB_ORG,
timeout=3000,
)

record = (
Point("api_call")
.field("request_count", request_count)
.time(time or datetime.now())
)

for tag in tags:
k, v = tag.split("=")
record.tag(k, v)

write_api = influxdb_client.write_api(write_options=SYNCHRONOUS)
write_api.write(bucket=bucket_name, record=record)

0 comments on commit 77eeaa7

Please sign in to comment.