97 lines
2.9 KiB
Python
97 lines
2.9 KiB
Python
"""Verify the countydata connection without exposing credentials or row data."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import os
|
|
import sys
|
|
|
|
import psycopg2
|
|
|
|
|
|
REQUIRED_ENV_VARS = (
|
|
"PGHOST",
|
|
"PGPORT",
|
|
"PGDATABASE",
|
|
"PGUSER",
|
|
"PGPASSWORD",
|
|
"PGSSLMODE",
|
|
)
|
|
|
|
|
|
def main() -> int:
|
|
missing = [name for name in REQUIRED_ENV_VARS if not os.environ.get(name)]
|
|
if missing:
|
|
print(
|
|
"Missing database environment variables: " + ", ".join(missing),
|
|
file=sys.stderr,
|
|
)
|
|
print("Run this through the VS Code task so .env is loaded safely.", file=sys.stderr)
|
|
return 1
|
|
|
|
if os.environ["PGSSLMODE"].lower() not in {"require", "verify-ca", "verify-full"}:
|
|
print("Refusing to connect because PGSSLMODE does not require TLS.", file=sys.stderr)
|
|
return 1
|
|
|
|
try:
|
|
connection = psycopg2.connect(
|
|
application_name="summer_project_vscode_check",
|
|
connect_timeout=10,
|
|
options=(
|
|
"-c default_transaction_read_only=on "
|
|
"-c statement_timeout=30000"
|
|
),
|
|
)
|
|
connection.set_session(readonly=True)
|
|
except psycopg2.Error as exc:
|
|
print(f"Connection failed: {exc.diag.message_primary or type(exc).__name__}", file=sys.stderr)
|
|
return 1
|
|
|
|
try:
|
|
with connection, connection.cursor() as cursor:
|
|
cursor.execute(
|
|
"""
|
|
SELECT
|
|
current_database(),
|
|
current_user,
|
|
current_setting('transaction_read_only'),
|
|
version()
|
|
"""
|
|
)
|
|
database, user, read_only, version = cursor.fetchone()
|
|
|
|
cursor.execute(
|
|
"""
|
|
SELECT ssl, version, cipher, bits
|
|
FROM pg_stat_ssl
|
|
WHERE pid = pg_backend_pid()
|
|
"""
|
|
)
|
|
ssl, tls_version, cipher, bits = cursor.fetchone()
|
|
|
|
cursor.execute(
|
|
"""
|
|
SELECT count(*)
|
|
FROM pg_class AS c
|
|
JOIN pg_namespace AS n ON n.oid = c.relnamespace
|
|
WHERE c.relkind IN ('r', 'p', 'v', 'm', 'f')
|
|
AND n.nspname NOT IN ('pg_catalog', 'information_schema')
|
|
AND n.nspname NOT LIKE 'pg_toast%'
|
|
"""
|
|
)
|
|
relation_count = cursor.fetchone()[0]
|
|
|
|
print("Countydata connection succeeded")
|
|
print(f" database: {database}")
|
|
print(f" user: {user}")
|
|
print(f" transaction read-only: {read_only}")
|
|
print(f" TLS: {ssl} ({tls_version}, {cipher}, {bits}-bit)")
|
|
print(f" visible non-system relations: {relation_count}")
|
|
print(f" server: {version}")
|
|
return 0
|
|
finally:
|
|
connection.close()
|
|
|
|
|
|
if __name__ == "__main__":
|
|
raise SystemExit(main())
|