misc.python.materialize.util

Various utilities

  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
 10"""Various utilities"""
 11
 12from __future__ import annotations
 13
 14import hashlib
 15import json
 16import os
 17import pathlib
 18import random
 19import re
 20import subprocess
 21from collections.abc import Iterator
 22from dataclasses import dataclass
 23from enum import Enum
 24from pathlib import Path
 25from threading import Thread
 26from typing import Protocol, TypeVar
 27from urllib.parse import parse_qs, quote, unquote, urlparse
 28
 29import psycopg
 30import xxhash
 31import zstandard
 32
 33MZ_ROOT = Path(os.environ["MZ_ROOT"])
 34
 35
 36def nonce(digits: int) -> str:
 37    return "".join(random.choice("0123456789abcdef") for _ in range(digits))
 38
 39
 40T = TypeVar("T")
 41
 42
 43def all_subclasses(cls: type[T]) -> set[type[T]]:
 44    """Returns a recursive set of all subclasses of a class"""
 45    sc = cls.__subclasses__()
 46    return set(sc).union([subclass for c in sc for subclass in all_subclasses(c)])
 47
 48
 49NAUGHTY_STRINGS = None
 50
 51
 52def naughty_strings() -> list[str]:
 53    # Naughty strings taken from https://github.com/minimaxir/big-list-of-naughty-strings
 54    # Under MIT license, Copyright (c) 2015-2020 Max Woolf
 55    global NAUGHTY_STRINGS
 56    if not NAUGHTY_STRINGS:
 57        with open(MZ_ROOT / "misc" / "python" / "materialize" / "blns.json") as f:
 58            NAUGHTY_STRINGS = json.load(f)
 59    return NAUGHTY_STRINGS
 60
 61
 62class YesNoOnce(Enum):
 63    YES = 1
 64    NO = 2
 65    ONCE = 3
 66
 67
 68class PropagatingThread(Thread):
 69    def run(self):
 70        self.exc = None
 71        try:
 72            self.ret = self._target(*self._args, **self._kwargs)  # type: ignore
 73        except BaseException as e:
 74            self.exc = e
 75
 76    def join(self, timeout=None):
 77        super().join(timeout)
 78        if self.exc:
 79            raise self.exc
 80        if hasattr(self, "ret"):
 81            return self.ret
 82
 83
 84def decompress_zst_to_directory(
 85    zst_file_path: str, destination_dir_path: str
 86) -> list[str]:
 87    """
 88    :return: file paths in destination dir
 89    """
 90    input_file = pathlib.Path(zst_file_path)
 91    output_paths = []
 92
 93    with open(input_file, "rb") as compressed:
 94        decompressor = zstandard.ZstdDecompressor()
 95        output_path = pathlib.Path(destination_dir_path) / input_file.stem
 96        output_paths.append(str(output_path))
 97        with open(output_path, "wb") as destination:
 98            decompressor.copy_stream(compressed, destination)
 99
100    return output_paths
101
102
103def ensure_dir_exists(path_to_dir: str) -> None:
104    subprocess.run(
105        [
106            "mkdir",
107            "-p",
108            f"{path_to_dir}",
109        ],
110        check=True,
111    )
112
113
114def sha256_of_utf8_string(value: str) -> str:
115    return hashlib.sha256(bytes(value, encoding="utf-8")).hexdigest()
116
117
118def stable_int_hash(*values: str) -> int:
119    if len(values) == 1:
120        return xxhash.xxh64(values[0], seed=0).intdigest()
121
122    return stable_int_hash(",".join([str(stable_int_hash(entry)) for entry in values]))
123
124
125class HasName(Protocol):
126    name: str
127
128
129U = TypeVar("U", bound=HasName)
130
131
132def selected_by_name(selected: list[str], objs: list[U]) -> Iterator[U]:
133    for name in selected:
134        for obj in objs:
135            if obj.name == name:
136                yield obj
137                break
138        else:
139            raise ValueError(
140                f"Unknown object with name {name} in {[obj.name for obj in objs]}"
141            )
142
143
144@dataclass
145class PgConnInfo:
146    user: str
147    host: str
148    port: int
149    database: str
150    password: str | None = None
151    ssl: bool = False
152    cluster: str | None = None
153    autocommit: bool = False
154
155    def connect(self) -> psycopg.Connection:
156        conn = psycopg.connect(
157            host=self.host,
158            port=self.port,
159            user=self.user,
160            password=self.password,
161            dbname=self.database,
162            sslmode="require" if self.ssl else None,
163        )
164        # Set SO_LINGER(1, 0) so close() sends RST instead of FIN, bypassing
165        # TIME_WAIT. Prevents exhausting the ~28k ephemeral port range under
166        # high connection churn (e.g. benchmarks doing rapid connect/disconnect).
167        self._set_linger(conn)
168        if self.autocommit:
169            conn.autocommit = True
170        if self.cluster:
171            with conn.cursor() as cur:
172                cur.execute(f"SET cluster = {self.cluster}".encode())
173        return conn
174
175    @staticmethod
176    def _set_linger(conn: psycopg.Connection) -> None:
177        import socket
178        import struct
179
180        fd = conn.pgconn.socket
181        if fd < 0:
182            return
183        sock = socket.socket(fileno=fd)
184        try:
185            sock.setsockopt(
186                socket.SOL_SOCKET, socket.SO_LINGER, struct.pack("ii", 1, 0)
187            )
188        finally:
189            sock.detach()
190
191    def to_conn_string(self) -> str:
192        return (
193            f"postgres://{quote(self.user)}:{quote(self.password)}@{self.host}:{self.port}/{quote(self.database)}"
194            if self.password
195            else f"postgres://{quote(self.user)}@{self.host}:{self.port}/{quote(self.database)}"
196        )
197
198
199def parse_pg_conn_string(conn_string: str) -> PgConnInfo:
200    """Not supported natively by pg8000, so we have to parse ourselves"""
201    url = urlparse(conn_string)
202    query_params = parse_qs(url.query)
203    assert url.username
204    assert url.hostname
205    return PgConnInfo(
206        user=unquote(url.username),
207        password=unquote(url.password) if url.password else url.password,
208        host=url.hostname,
209        port=url.port or 5432,
210        database=url.path.lstrip("/"),
211        ssl=query_params.get("sslmode", ["disable"])[-1] != "disable",
212    )
213
214
215FILTERED_ARGS = [
216    # Secrets
217    "mzp_",
218    "-----BEGIN PRIVATE KEY-----",
219    "-----BEGIN CERTIFICATE-----",
220    "confluent-api-key=",
221    "confluent-api-secret=",
222    "aws-access-key-id=",
223    "aws-secret-access-key=",
224    "default-sql-server-password=",
225    "Authorization: Bearer ",
226    "client_secret=",
227    # Not a secret, but too spammy, filter too
228    "CLUSTER_REPLICA_SIZES",
229    "cluster-replica-sizes=",
230]
231
232
233def filter_cmd(args: list[str]) -> list[str]:
234    """Don't print out secrets in test logs"""
235    return [
236        (
237            "[REDACTED]"
238            if any(filtered_arg in arg for filtered_arg in FILTERED_ARGS)
239            else arg
240        )
241        for arg in args
242    ]
243
244
245def redact_secrets(text: str) -> str:
246    text = re.sub(
247        r"-----BEGIN [A-Z ]+-----.*?-----END [A-Z ]+-----",
248        "[REDACTED]",
249        text,
250        flags=re.DOTALL,
251    )
252    for secret in FILTERED_ARGS:
253        if secret in text:
254            text = re.sub(re.escape(secret) + r"\S*", "[REDACTED]", text)
255    return text
MZ_ROOT = PosixPath('/var/lib/buildkite-agent/builds/buildkite-15f2293-i-076fb2ae1f4c2fe07-1/materialize/deploy')
def nonce(digits: int) -> str: 
37def nonce(digits: int) -> str:
38    return "".join(random.choice("0123456789abcdef") for _ in range(digits))
def all_subclasses(cls: type[~T]) -> set[type[~T]]: 
44def all_subclasses(cls: type[T]) -> set[type[T]]:
45    """Returns a recursive set of all subclasses of a class"""
46    sc = cls.__subclasses__()
47    return set(sc).union([subclass for c in sc for subclass in all_subclasses(c)])

Returns a recursive set of all subclasses of a class

NAUGHTY_STRINGS = None
def naughty_strings() -> list[str]: 
53def naughty_strings() -> list[str]:
54    # Naughty strings taken from https://github.com/minimaxir/big-list-of-naughty-strings
55    # Under MIT license, Copyright (c) 2015-2020 Max Woolf
56    global NAUGHTY_STRINGS
57    if not NAUGHTY_STRINGS:
58        with open(MZ_ROOT / "misc" / "python" / "materialize" / "blns.json") as f:
59            NAUGHTY_STRINGS = json.load(f)
60    return NAUGHTY_STRINGS
class YesNoOnce(enum.Enum): 
63class YesNoOnce(Enum):
64    YES = 1
65    NO = 2
66    ONCE = 3
YES = <YesNoOnce.YES: 1>
NO = <YesNoOnce.NO: 2>
ONCE = <YesNoOnce.ONCE: 3>
class PropagatingThread(threading.Thread): 
69class PropagatingThread(Thread):
70    def run(self):
71        self.exc = None
72        try:
73            self.ret = self._target(*self._args, **self._kwargs)  # type: ignore
74        except BaseException as e:
75            self.exc = e
76
77    def join(self, timeout=None):
78        super().join(timeout)
79        if self.exc:
80            raise self.exc
81        if hasattr(self, "ret"):
82            return self.ret

A class that represents a thread of control.

This class can be safely subclassed in a limited fashion. There are two ways to specify the activity: by passing a callable object to the constructor, or by overriding the run() method in a subclass.

def run(self): 
70    def run(self):
71        self.exc = None
72        try:
73            self.ret = self._target(*self._args, **self._kwargs)  # type: ignore
74        except BaseException as e:
75            self.exc = e

Method representing the thread's activity.

You may override this method in a subclass. The standard run() method invokes the callable object passed to the object's constructor as the target argument, if any, with sequential and keyword arguments taken from the args and kwargs arguments, respectively.

def join(self, timeout=None): 
77    def join(self, timeout=None):
78        super().join(timeout)
79        if self.exc:
80            raise self.exc
81        if hasattr(self, "ret"):
82            return self.ret

Wait until the thread terminates.

This blocks the calling thread until the thread whose join() method is called terminates -- either normally or through an unhandled exception or until the optional timeout occurs.

When the timeout argument is present and not None, it should be a floating point number specifying a timeout for the operation in seconds (or fractions thereof). As join() always returns None, you must call is_alive() after join() to decide whether a timeout happened -- if the thread is still alive, the join() call timed out.

When the timeout argument is not present or None, the operation will block until the thread terminates.

A thread can be join()ed many times.

join() raises a RuntimeError if an attempt is made to join the current thread as that would cause a deadlock. It is also an error to join() a thread before it has been started and attempts to do so raises the same exception.

def decompress_zst_to_directory(zst_file_path: str, destination_dir_path: str) -> list[str]: 
 85def decompress_zst_to_directory(
 86    zst_file_path: str, destination_dir_path: str
 87) -> list[str]:
 88    """
 89    :return: file paths in destination dir
 90    """
 91    input_file = pathlib.Path(zst_file_path)
 92    output_paths = []
 93
 94    with open(input_file, "rb") as compressed:
 95        decompressor = zstandard.ZstdDecompressor()
 96        output_path = pathlib.Path(destination_dir_path) / input_file.stem
 97        output_paths.append(str(output_path))
 98        with open(output_path, "wb") as destination:
 99            decompressor.copy_stream(compressed, destination)
100
101    return output_paths
Returns

file paths in destination dir

def ensure_dir_exists(path_to_dir: str) -> None: 
104def ensure_dir_exists(path_to_dir: str) -> None:
105    subprocess.run(
106        [
107            "mkdir",
108            "-p",
109            f"{path_to_dir}",
110        ],
111        check=True,
112    )
def sha256_of_utf8_string(value: str) -> str: 
115def sha256_of_utf8_string(value: str) -> str:
116    return hashlib.sha256(bytes(value, encoding="utf-8")).hexdigest()
def stable_int_hash(*values: str) -> int: 
119def stable_int_hash(*values: str) -> int:
120    if len(values) == 1:
121        return xxhash.xxh64(values[0], seed=0).intdigest()
122
123    return stable_int_hash(",".join([str(stable_int_hash(entry)) for entry in values]))
class HasName(typing.Protocol): 
126class HasName(Protocol):
127    name: str

Base class for protocol classes.

Protocol classes are defined as::

class Proto(Protocol):
    def meth(self) -> int:
        ...

Such classes are primarily used with static type checkers that recognize structural subtyping (static duck-typing).

For example::

class C:
    def meth(self) -> int:
        return 0

def func(x: Proto) -> int:
    return x.meth()

func(C())  # Passes static type check

See PEP 544 for details. Protocol classes decorated with @typing.runtime_checkable act as simple-minded runtime protocols that check only the presence of given attributes, ignoring their type signatures. Protocol classes can be generic, they are defined as::

class GenProto[T](Protocol):
    def meth(self) -> T:
        ...
HasName(*args, **kwargs) 
1739def _no_init_or_replace_init(self, *args, **kwargs):
1740    cls = type(self)
1741
1742    if cls._is_protocol:
1743        raise TypeError('Protocols cannot be instantiated')
1744
1745    # Already using a custom `__init__`. No need to calculate correct
1746    # `__init__` to call. This can lead to RecursionError. See bpo-45121.
1747    if cls.__init__ is not _no_init_or_replace_init:
1748        return
1749
1750    # Initially, `__init__` of a protocol subclass is set to `_no_init_or_replace_init`.
1751    # The first instantiation of the subclass will call `_no_init_or_replace_init` which
1752    # searches for a proper new `__init__` in the MRO. The new `__init__`
1753    # replaces the subclass' old `__init__` (ie `_no_init_or_replace_init`). Subsequent
1754    # instantiation of the protocol subclass will thus use the new
1755    # `__init__` and no longer call `_no_init_or_replace_init`.
1756    for base in cls.__mro__:
1757        init = base.__dict__.get('__init__', _no_init_or_replace_init)
1758        if init is not _no_init_or_replace_init:
1759            cls.__init__ = init
1760            break
1761    else:
1762        # should not happen
1763        cls.__init__ = object.__init__
1764
1765    cls.__init__(self, *args, **kwargs)
name: str
def selected_by_name(selected: list[str], objs: list[~U]) -> Iterator[~U]: 
133def selected_by_name(selected: list[str], objs: list[U]) -> Iterator[U]:
134    for name in selected:
135        for obj in objs:
136            if obj.name == name:
137                yield obj
138                break
139        else:
140            raise ValueError(
141                f"Unknown object with name {name} in {[obj.name for obj in objs]}"
142            )
@dataclass
class PgConnInfo: 
145@dataclass
146class PgConnInfo:
147    user: str
148    host: str
149    port: int
150    database: str
151    password: str | None = None
152    ssl: bool = False
153    cluster: str | None = None
154    autocommit: bool = False
155
156    def connect(self) -> psycopg.Connection:
157        conn = psycopg.connect(
158            host=self.host,
159            port=self.port,
160            user=self.user,
161            password=self.password,
162            dbname=self.database,
163            sslmode="require" if self.ssl else None,
164        )
165        # Set SO_LINGER(1, 0) so close() sends RST instead of FIN, bypassing
166        # TIME_WAIT. Prevents exhausting the ~28k ephemeral port range under
167        # high connection churn (e.g. benchmarks doing rapid connect/disconnect).
168        self._set_linger(conn)
169        if self.autocommit:
170            conn.autocommit = True
171        if self.cluster:
172            with conn.cursor() as cur:
173                cur.execute(f"SET cluster = {self.cluster}".encode())
174        return conn
175
176    @staticmethod
177    def _set_linger(conn: psycopg.Connection) -> None:
178        import socket
179        import struct
180
181        fd = conn.pgconn.socket
182        if fd < 0:
183            return
184        sock = socket.socket(fileno=fd)
185        try:
186            sock.setsockopt(
187                socket.SOL_SOCKET, socket.SO_LINGER, struct.pack("ii", 1, 0)
188            )
189        finally:
190            sock.detach()
191
192    def to_conn_string(self) -> str:
193        return (
194            f"postgres://{quote(self.user)}:{quote(self.password)}@{self.host}:{self.port}/{quote(self.database)}"
195            if self.password
196            else f"postgres://{quote(self.user)}@{self.host}:{self.port}/{quote(self.database)}"
197        )
PgConnInfo( user: str, host: str, port: int, database: str, password: str | None = None, ssl: bool = False, cluster: str | None = None, autocommit: bool = False)
user: str
host: str
port: int
database: str
password: str | None = None
ssl: bool = False
cluster: str | None = None
autocommit: bool = False
def connect(self) -> psycopg.Connection: 
156    def connect(self) -> psycopg.Connection:
157        conn = psycopg.connect(
158            host=self.host,
159            port=self.port,
160            user=self.user,
161            password=self.password,
162            dbname=self.database,
163            sslmode="require" if self.ssl else None,
164        )
165        # Set SO_LINGER(1, 0) so close() sends RST instead of FIN, bypassing
166        # TIME_WAIT. Prevents exhausting the ~28k ephemeral port range under
167        # high connection churn (e.g. benchmarks doing rapid connect/disconnect).
168        self._set_linger(conn)
169        if self.autocommit:
170            conn.autocommit = True
171        if self.cluster:
172            with conn.cursor() as cur:
173                cur.execute(f"SET cluster = {self.cluster}".encode())
174        return conn
def to_conn_string(self) -> str: 
192    def to_conn_string(self) -> str:
193        return (
194            f"postgres://{quote(self.user)}:{quote(self.password)}@{self.host}:{self.port}/{quote(self.database)}"
195            if self.password
196            else f"postgres://{quote(self.user)}@{self.host}:{self.port}/{quote(self.database)}"
197        )
def parse_pg_conn_string(conn_string: str) -> PgConnInfo: 
200def parse_pg_conn_string(conn_string: str) -> PgConnInfo:
201    """Not supported natively by pg8000, so we have to parse ourselves"""
202    url = urlparse(conn_string)
203    query_params = parse_qs(url.query)
204    assert url.username
205    assert url.hostname
206    return PgConnInfo(
207        user=unquote(url.username),
208        password=unquote(url.password) if url.password else url.password,
209        host=url.hostname,
210        port=url.port or 5432,
211        database=url.path.lstrip("/"),
212        ssl=query_params.get("sslmode", ["disable"])[-1] != "disable",
213    )

Not supported natively by pg8000, so we have to parse ourselves

FILTERED_ARGS = ['mzp_', '-----BEGIN PRIVATE KEY-----', '-----BEGIN CERTIFICATE-----', 'confluent-api-key=', 'confluent-api-secret=', 'aws-access-key-id=', 'aws-secret-access-key=', 'default-sql-server-password=', 'Authorization: Bearer ', 'client_secret=', 'CLUSTER_REPLICA_SIZES', 'cluster-replica-sizes=']
def filter_cmd(args: list[str]) -> list[str]: 
234def filter_cmd(args: list[str]) -> list[str]:
235    """Don't print out secrets in test logs"""
236    return [
237        (
238            "[REDACTED]"
239            if any(filtered_arg in arg for filtered_arg in FILTERED_ARGS)
240            else arg
241        )
242        for arg in args
243    ]

Don't print out secrets in test logs

def redact_secrets(text: str) -> str: 
246def redact_secrets(text: str) -> str:
247    text = re.sub(
248        r"-----BEGIN [A-Z ]+-----.*?-----END [A-Z ]+-----",
249        "[REDACTED]",
250        text,
251        flags=re.DOTALL,
252    )
253    for secret in FILTERED_ARGS:
254        if secret in text:
255            text = re.sub(re.escape(secret) + r"\S*", "[REDACTED]", text)
256    return text