misc.python.materialize.optbench.sql

  1# Copyright Materialize, Inc. and contributors. All rights reserved.
  2#
  3# Use of this software is governed by the Business Source License
  4# included in the LICENSE file at the root of this repository.
  5#
  6# As of the Change Date specified in that file, in accordance with
  7# the Business Source License, use of this software will be governed
  8# by the Apache License, Version 2.0.
  9
 10import logging
 11import re
 12from enum import Enum
 13from pathlib import Path
 14from typing import Any, cast
 15
 16import numpy as np
 17import psycopg
 18import sqlparse
 19
 20from . import Scenario, util
 21
 22
 23class Dialect(Enum):
 24    PG = 0
 25    MZ = 1
 26
 27
 28class Query:
 29    """An API for manipulating workload queries."""
 30
 31    def __init__(self, query: str) -> None:
 32        self.query = query
 33
 34    def __str__(self) -> str:
 35        return self.query
 36
 37    def name(self) -> str:
 38        """Extracts and returns the name of this query from a '-- name: {name}' comment.
 39        Returns 'anonymous' if the name is not set."""
 40        p = r"-- name\: (?P<name>.+)"
 41        m = re.search(p, self.query, re.MULTILINE)
 42        return m.group("name") if m else "anonoymous"
 43
 44    def explain(self, timing: bool, dialect: Dialect = Dialect.MZ) -> str:
 45        """Prepends 'EXPLAIN ...' to the query respecting the given dialect."""
 46
 47        if dialect == Dialect.PG:
 48            if timing:
 49                return "\n".join(["EXPLAIN (ANALYZE, TIMING TRUE)", self.query])
 50            else:
 51                return "\n".join(["EXPLAIN", self.query])
 52        else:
 53            if timing:
 54                return "\n".join(["EXPLAIN WITH(timing)", self.query])
 55            else:
 56                return "\n".join(["EXPLAIN", self.query])
 57
 58
 59class ExplainOutput:
 60    """An API for manipulating 'EXPLAIN ... PLAN FOR' results."""
 61
 62    def __init__(self, output: str) -> None:
 63        self.output = output
 64
 65    def __str__(self) -> str:
 66        return self.output
 67
 68    def optimization_time(self) -> np.timedelta64 | None:
 69        """Optionally, returns the optimization_time time for an 'EXPLAIN' output."""
 70        p = r"(Optimization time|Planning Time)\: (?P<time>[0-9]+(\.[0-9]+)?\s?\S+)"
 71        m = re.search(p, self.output, re.MULTILINE)
 72        return util.duration_to_timedelta(m["time"]) if m else None
 73
 74
 75class Database:
 76    """An API to the database under test."""
 77
 78    def __init__(
 79        self,
 80        port: int,
 81        host: str,
 82        user: str,
 83        password: str | None,
 84        database: str | None,
 85        require_ssl: bool,
 86    ) -> None:
 87        logging.debug(f"Initialize Database with host={host} port={port}, user={user}")
 88
 89        self.conn = psycopg.connect(
 90            host=host,
 91            port=port,
 92            user=user,
 93            password=password,
 94            dbname=database,
 95            sslmode="require" if require_ssl else "disable",
 96        )
 97        self.conn.autocommit = True
 98        self.dialect = Dialect.MZ if "Materialize" in self.version() else Dialect.PG
 99
100    def close(self) -> None:
101        self.conn.close()
102
103    def version(self) -> str:
104        result = self.query_one("SELECT version()")
105        return cast(str, result[0])
106
107    def mz_version(self) -> str | None:
108        if self.dialect == Dialect.MZ:
109            result = self.query_one("SELECT mz_version()")
110            return cast(str, result[0])
111        else:
112            return None
113
114    def drop_database(self, scenario: Scenario) -> None:
115        logging.debug(f'Drop database "{scenario}"')
116        self.execute(f"DROP DATABASE IF EXISTS {scenario}")
117
118    def create_database(self, scenario: Scenario) -> None:
119        logging.debug(f'Create database "{scenario}"')
120        self.execute(f"CREATE DATABASE {scenario}")
121
122    def explain(self, query: Query, timing: bool) -> "ExplainOutput":
123        result = self.query_all(query.explain(timing, self.dialect))
124        return ExplainOutput("\n".join([col for line in result for col in line]))
125
126    def execute(self, statement: str) -> None:
127        with self.conn.cursor() as cursor:
128            cursor.execute(statement.encode())
129
130    def execute_all(self, statements: list[str]) -> None:
131        with self.conn.cursor() as cursor:
132            for statement in statements:
133                cursor.execute(statement.encode())
134
135    def query_one(self, query: str) -> dict[Any, Any]:
136        with self.conn.cursor() as cursor:
137            cursor.execute(query.encode())
138            return cast(dict[Any, Any], cursor.fetchone())
139
140    def query_all(self, query: str) -> dict[Any, Any]:
141        with self.conn.cursor() as cursor:
142            cursor.execute(query.encode())
143            return cast(dict[Any, Any], cursor.fetchall())
144
145
146# Utility functions
147# -----------------
148
149
150def parse_from_file(path: Path) -> list[str]:
151    """Parses a *.sql file to a list of queries."""
152    return sqlparse.split(path.read_text())