datastore/tests/integration/test_prometheus_protocols/generate_pb2.py

#!/usr/bin/env python3

# This is a helper utility.
# It generates files in the "pb2" folder using the protocol buffer compiler.
# This script must be called manually after updating contrib/prometheus-protobufs/prompb/remote.proto

import grpc_tools  # pip3 install grpcio-tools

import os, subprocess


SCRIPT_DIR = os.path.dirname(os.path.realpath(__file__))
ROOT_DIR = os.path.abspath(os.path.join(SCRIPT_DIR, "../../.."))

PROTO_DIR = os.path.abspath(os.path.join(ROOT_DIR, "contrib/prometheus-protobufs"))
REMOTE_PROTO = "prompb/remote.proto"
TYPES_PROTO = "prompb/types.proto"

GOGO_INCLUDE_DIR = os.path.abspath(
    os.path.join(ROOT_DIR, "contrib/prometheus-protobufs-gogo")
)
GOGO_PROTO = "gogoproto/gogo.proto"

# Files in the "pb2" folder which will be generated by this script.
PB2_DIR = os.path.join(SCRIPT_DIR, "pb2")
PB2_FILENAMES = ["prompb/remote_pb2.py", "prompb/types_pb2.py", "gogoproto/gogo_pb2.py"]


# Processes the protobuf schema with the protocol buffer compiler and generates the "pb2" folder.
def generate_pb2():
    print(f"Generating files:")
    for filename in PB2_FILENAMES:
        print(os.path.join(PB2_DIR, filename))

    os.makedirs(PB2_DIR, exist_ok=True)

    cmd = [
        "python3",
        "-m",
        "grpc_tools.protoc",
        "-I" + PROTO_DIR,
        "-I" + GOGO_INCLUDE_DIR,
        "--python_out=" + PB2_DIR,
        os.path.join(PROTO_DIR, REMOTE_PROTO),
        os.path.join(PROTO_DIR, TYPES_PROTO),
        os.path.join(GOGO_INCLUDE_DIR, GOGO_PROTO),
    ]
    subprocess.run(cmd)

    for filename in PB2_FILENAMES:
        assert os.path.exists(os.path.join(PB2_DIR, filename))
    print("Done! (generate_pb2)")


# MAIN
if __name__ == "__main__":
    generate_pb2()