Wait for a pg-compatible database (includes materialized)
(
timeout_secs: int,
query: str,
dbname: str,
port: int,
host: str,
user: str,
password: str | None,
expected: Iterable[Any] | Literal["any"],
print_result: bool = False,
sslmode: str = "disable",
)
| 724 | |
| 725 | # TODO(benesch): replace with Docker health checks. |
| 726 | def _wait_for_pg( |
| 727 | timeout_secs: int, |
| 728 | query: str, |
| 729 | dbname: str, |
| 730 | port: int, |
| 731 | host: str, |
| 732 | user: str, |
| 733 | password: str | None, |
| 734 | expected: Iterable[Any] | Literal["any"], |
| 735 | print_result: bool = False, |
| 736 | sslmode: str = "disable", |
| 737 | ) -> None: |
| 738 | """Wait for a pg-compatible database (includes materialized)""" |
| 739 | obfuscated_password = password[0:1] if password is not None else "" |
| 740 | args = f"dbname={dbname} host={host} port={port} user={user} password='{obfuscated_password}...'" |
| 741 | ui.progress(f"waiting for {args} to handle {query!r}", "C") |
| 742 | error = None |
| 743 | for remaining in ui.timeout_loop(timeout_secs, tick=0.5): |
| 744 | try: |
| 745 | conn = psycopg.connect( |
| 746 | dbname=dbname, |
| 747 | host=host, |
| 748 | port=port, |
| 749 | user=user, |
| 750 | password=password, |
| 751 | connect_timeout=1, |
| 752 | sslmode=sslmode, |
| 753 | ) |
| 754 | # The default (autocommit = false) wraps everything in a transaction. |
| 755 | conn.autocommit = True |
| 756 | with conn.cursor() as cur: |
| 757 | cur.execute(query.encode()) |
| 758 | if expected == "any" and cur.rowcount == -1: |
| 759 | ui.progress(" success!", finish=True) |
| 760 | return |
| 761 | result = list(cur.fetchall()) |
| 762 | if expected == "any" or result == expected: |
| 763 | if print_result: |
| 764 | say(f"query result: {result}") |
| 765 | else: |
| 766 | ui.progress(" success!", finish=True) |
| 767 | return |
| 768 | else: |
| 769 | say( |
| 770 | f"host={host} port={port} did not return rows matching {expected} got: {result}" |
| 771 | ) |
| 772 | except Exception as e: |
| 773 | ui.progress(f"{e if print_result else ''} {int(remaining)}") |
| 774 | error = e |
| 775 | ui.progress(finish=True) |
| 776 | raise UIError(f"never got correct result for {args}: {error}") |
| 777 | |
| 778 | |
| 779 | def bootstrap_cluster_replica_size() -> str: |