mirror of
https://github.com/neondatabase/neon.git
synced 2025-12-22 21:59:59 +00:00
e60b70b475
## Problem When I switched `psycopg2.connect` from context manager to a regular function call in https://github.com/neondatabase/neon/pull/4382 I embarrassingly forgot about commit, so it doesn't really put data into DB 😞 ## Summary of changes - Enable autocommit for data ingestion scripts
119 lines
3.1 KiB
Python
119 lines
3.1 KiB
Python
#!/usr/bin/env python3
|
|
import argparse
|
|
import logging
|
|
import os
|
|
import re
|
|
import sys
|
|
from contextlib import contextmanager
|
|
from pathlib import Path
|
|
|
|
import backoff
|
|
import psycopg2
|
|
|
|
CREATE_TABLE = """
|
|
CREATE TABLE IF NOT EXISTS regress_test_results (
|
|
id SERIAL PRIMARY KEY,
|
|
reference CHAR(255),
|
|
revision CHAR(40),
|
|
build_type CHAR(16),
|
|
data JSONB
|
|
)
|
|
"""
|
|
|
|
|
|
def err(msg):
|
|
print(f"error: {msg}")
|
|
sys.exit(1)
|
|
|
|
|
|
@contextmanager
|
|
def get_connection_cursor():
|
|
connstr = os.getenv("DATABASE_URL")
|
|
if not connstr:
|
|
err("DATABASE_URL environment variable is not set")
|
|
|
|
@backoff.on_exception(backoff.expo, psycopg2.OperationalError, max_time=150)
|
|
def connect(connstr):
|
|
conn = psycopg2.connect(connstr, connect_timeout=30)
|
|
conn.autocommit = True
|
|
return conn
|
|
|
|
conn = connect(connstr)
|
|
try:
|
|
with conn.cursor() as cur:
|
|
yield cur
|
|
finally:
|
|
if conn is not None:
|
|
conn.close()
|
|
|
|
|
|
def create_table(cur):
|
|
cur.execute(CREATE_TABLE)
|
|
|
|
|
|
def ingest_regress_test_result(
|
|
cursor, reference: str, revision: str, build_type: str, data_file: Path
|
|
):
|
|
data = data_file.read_text()
|
|
# In the JSON report we can have lines related to LazyFixture with escaped double-quote
|
|
# It's hard to insert them into jsonb field as is, so replace \" with ' to make it easier for us
|
|
#
|
|
# "<LazyFixture \"vanilla_compare\">" -> "<LazyFixture 'vanilla_compare'>"
|
|
data = re.sub(r'("<LazyFixture )\\"([^\\]+)\\"(>")', r"\g<1>'\g<2>'\g<3>", data)
|
|
values = (
|
|
reference,
|
|
revision,
|
|
build_type,
|
|
data,
|
|
)
|
|
cursor.execute(
|
|
"""
|
|
INSERT INTO regress_test_results (
|
|
reference,
|
|
revision,
|
|
build_type,
|
|
data
|
|
) VALUES (%s, %s, %s, %s)
|
|
""",
|
|
values,
|
|
)
|
|
|
|
|
|
def main():
|
|
parser = argparse.ArgumentParser(
|
|
description="Regress test result uploader. \
|
|
Database connection string should be provided via DATABASE_URL environment variable",
|
|
)
|
|
parser.add_argument("--initdb", action="store_true", help="Initialuze database")
|
|
parser.add_argument(
|
|
"--reference", type=str, required=True, help="git reference, for example refs/heads/main"
|
|
)
|
|
parser.add_argument("--revision", type=str, required=True, help="git revision")
|
|
parser.add_argument(
|
|
"--build-type", type=str, required=True, help="build type: release, debug or remote"
|
|
)
|
|
parser.add_argument(
|
|
"--ingest", type=Path, required=True, help="Path to regress test result file"
|
|
)
|
|
|
|
args = parser.parse_args()
|
|
with get_connection_cursor() as cur:
|
|
if args.initdb:
|
|
create_table(cur)
|
|
|
|
if not args.ingest.exists():
|
|
err(f"ingest path {args.ingest} does not exist")
|
|
|
|
ingest_regress_test_result(
|
|
cur,
|
|
reference=args.reference,
|
|
revision=args.revision,
|
|
build_type=args.build_type,
|
|
data_file=args.ingest,
|
|
)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
logging.getLogger("backoff").addHandler(logging.StreamHandler())
|
|
main()
|