nautobot graphql performance test command

graphql_performance_tests.py

import uuid
import time

from graphql import get_default_backend
from rich.console import Console
from rich.table import Table

from django.core.management.base import BaseCommand

from graphene_django.settings import graphene_settings
from django.db import connection

from django.test.client import RequestFactory
from nautobot.users.models import User
from cacheops.invalidation import invalidate_all
 
NUMBER_OF_TRIES = 5

QUERIES = {
    "devices_interfaces_ips_edge": """
    {
        devices(role: "edge") {
            name
            interfaces {
                name
                ip_addresses {
                    address
                }
            }
        }
    }
""", 
    "devices_interfaces_all": """
    {
        devices {
            name
            interfaces {
                name
            }
        }
    }
""", 
    "devices_sites_regions_all": """
    {
        devices {
            name
            site {
                name
                region {
                    name
                }
            }
        }
    }
""",
    "golden_config_device_extensive": """
{
  devices(name: "ams-edge-01") {
    config_context
    name
    position
    serial
    primary_ip4 {
      id
      primary_ip4_for {
        id
        name
      }
    }
    tenant {
      name
    }
    tags {
      name
      slug
    }
    device_role {
      name
    }
    platform {
      name
      slug
      manufacturer {
        name
      }
      napalm_driver
    }
    site {
      name
      slug
      vlans {
        id
        name
        vid
      }
      vlan_groups {
        id
      }
    }
    interfaces {
      description
      mac_address
      enabled
      name
      ip_addresses {
        address
        tags {
          id
        }
      }
      connected_circuit_termination {
        circuit {
          cid
          commit_rate
          provider {
            name
          }
        }
      }
      tagged_vlans {
        id
      }
      untagged_vlan {
        id
      }
      cable {
        termination_a_type
        status {
          name
        }
        color
      }
      tagged_vlans {
        site {
          name
        }
        id
      }
      tags {
        id
      }
    }
  }
}
""" 
}

class Command(BaseCommand):
    help = "Execute some GraphQL query and measure the response time."

    def handle(self, *model_names, **options):

        # Generate a fake Web request to attach a user
        # This should be simpler once PR 5632 has been merged Upstream
        # https://github.com/netbox-community/netbox/pull/5632
        request = RequestFactory().request(SERVER_NAME="WebRequestContext")
        request.id = uuid.uuid4()
        user = User.objects.get(username="demo")
        request.user = user

        backend = get_default_backend()
        schema = graphene_settings.SCHEMA

        table = Table(title=f"GraphQL Performance Tests | {NUMBER_OF_TRIES} iterations")

        table.add_column("Query", justify="right", style="cyan", no_wrap=True)
        table.add_column(f"No Cache (min / avg / max)", style="magenta")
        table.add_column(f"Cache (min / avg / max)", style="magenta")
        # table.add_column(f"Nbr SQL Queries", style="magenta")

        for query_name, query in QUERIES.items():
            results = [query_name]
            nbr_db_query = None
            
            for use_cache in [False, True]:
                query_results = []

                for i in range(0, NUMBER_OF_TRIES):
                    
                    if not use_cache:
                        invalidate_all()
                    
                    # nbr_query_start = len(connection.queries)
                    print(f"Executing '{query_name}', {i+1}/{NUMBER_OF_TRIES} : Cache : {use_cache}")

                    start_time = time.time()
                    document = backend.document_from_string(schema, query)
                    result = document.execute(context_value=request)

                    # Calculate the number of SQL queries on the first run
                    # if not nbr_db_query:
                    #     nbr_db_query = len(connection.queries) - nbr_query_start

                    execution_time = time.time() - start_time
                    query_results.append(execution_time)
            
                average = sum(query_results) / NUMBER_OF_TRIES
                min_time = min(query_results)
                max_time = max(query_results)
                results.append(f"{min_time:.4f} / {average:.4f} / {max_time:.4f}")
            
            # results.append(f"{nbr_db_query}")
            table.add_row(*results)

        console = Console()
        console.print(table)