query.py

from attrs import define
from typing import ClassVar


@define
class QueryContainer:
    """Container for SQL queries used in performance testing"""

    # Query 1: Complex JOIN with aggregation (tests query processing)
    COMPLEX_JOIN_AGGREGATION: ClassVar[str] = """
        SELECT
            p.ProductID,
            p.Name AS ProductName,
            pc.Name AS Category,
            psc.Name AS Subcategory,
            COUNT(sod.SalesOrderDetailID) AS TotalOrders,
            SUM(sod.OrderQty) AS TotalQuantity,
            SUM(sod.LineTotal) AS TotalRevenue,
            AVG(sod.UnitPrice) AS AvgPrice
        FROM Sales.SalesOrderDetail sod
        INNER JOIN Production.Product p ON sod.ProductID = p.ProductID
        INNER JOIN Production.ProductSubcategory psc ON p.ProductSubcategoryID = psc.ProductSubcategoryID
        INNER JOIN Production.ProductCategory pc ON psc.ProductCategoryID = pc.ProductCategoryID
        GROUP BY p.ProductID, p.Name, pc.Name, psc.Name
        HAVING SUM(sod.LineTotal) > 10000
        ORDER BY TotalRevenue DESC;
    """

    # Query 2: Large dataset retrieval (tests data transfer speed)
    LARGE_DATASET: ClassVar[str] = """
        SELECT
            soh.SalesOrderID,
            soh.OrderDate,
            soh.DueDate,
            soh.ShipDate,
            soh.Status,
            soh.SubTotal,
            soh.TaxAmt,
            soh.Freight,
            soh.TotalDue,
            c.CustomerID,
            p.FirstName,
            p.LastName,
            a.AddressLine1,
            a.City,
            sp.Name AS StateProvince,
            cr.Name AS Country
        FROM Sales.SalesOrderHeader soh
        INNER JOIN Sales.Customer c ON soh.CustomerID = c.CustomerID
        INNER JOIN Person.Person p ON c.PersonID = p.BusinessEntityID
        INNER JOIN Person.BusinessEntityAddress bea ON p.BusinessEntityID = bea.BusinessEntityID
        INNER JOIN Person.Address a ON bea.AddressID = a.AddressID
        INNER JOIN Person.StateProvince sp ON a.StateProvinceID = sp.StateProvinceID
        INNER JOIN Person.CountryRegion cr ON sp.CountryRegionCode = cr.CountryRegionCode
        WHERE soh.OrderDate >= '2013-01-01';
    """

    # Query 3: synthetic 1.2M rows result set
    VERY_LARGE_DATASET: ClassVar[str] = """
    SELECT
        sod.SalesOrderID,
        sod.SalesOrderDetailID,
        sod.ProductID,
        sod.OrderQty,
        sod.UnitPrice,
        sod.LineTotal,
        p.Name AS ProductName,
        p.ProductNumber,
        p.Color,
        p.ListPrice,
        n1.number AS RowMultiplier1
    FROM Sales.SalesOrderDetail sod
    CROSS JOIN (SELECT TOP 10 ROW_NUMBER() OVER (ORDER BY (SELECT NULL)) AS number
                FROM Sales.SalesOrderDetail) n1
    INNER JOIN Production.Product p ON sod.ProductID = p.ProductID;
    """
    # Query 4: CTE
    SUBQUERY_WITH_CTE: ClassVar[str] = """
        WITH SalesSummary AS (
            SELECT
                soh.SalesPersonID,
                YEAR(soh.OrderDate) AS OrderYear,
                SUM(soh.TotalDue) AS YearlyTotal
            FROM Sales.SalesOrderHeader soh
            WHERE soh.SalesPersonID IS NOT NULL
            GROUP BY soh.SalesPersonID, YEAR(soh.OrderDate)
        ),
        RankedSales AS (
            SELECT
                SalesPersonID,
                OrderYear,
                YearlyTotal,
                RANK() OVER (PARTITION BY OrderYear ORDER BY YearlyTotal DESC) AS SalesRank
            FROM SalesSummary
        )
        SELECT
            rs.SalesPersonID,
            p.FirstName,
            p.LastName,
            rs.OrderYear,
            rs.YearlyTotal,
            rs.SalesRank
        FROM RankedSales rs
        INNER JOIN Person.Person p ON rs.SalesPersonID = p.BusinessEntityID
        WHERE rs.SalesRank <= 10
        ORDER BY rs.OrderYear DESC, rs.SalesRank;
    """

    @classmethod
    def get_query(cls, key: str | None = None) -> dict[str, str]:
        """Retrieves a query or all queries.
        Args:
            key (str | None, optional): query name. Defaults to None.
        Raises:
            KeyError: error if the query does not exist.
        Returns:
            str | dict[str, str]: query definition or dictionary of all queries.
        """
        queries = {
            'Complex JOIN Aggregation': cls.COMPLEX_JOIN_AGGREGATION,
            'Large Dataset': cls.LARGE_DATASET,
            'Very Large Dataset': cls.VERY_LARGE_DATASET,
            'Subquery with CTE': cls.SUBQUERY_WITH_CTE
        }

        if key is None:
            return queries

        if key not in queries:
            raise KeyError(f"Query '{key}' not found. Available keys: {list(queries.keys())}")

        return {key: queries[key]}  # must return a dict either way