Spaces:
Running
Running
File size: 1,439 Bytes
1c61a98 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 | import os
from urllib.parse import parse_qsl, urlencode, urlsplit, urlunsplit
import psycopg
from dotenv import load_dotenv
load_dotenv()
ALLOWED_POSTGRES_PARAMS = {
"application_name",
"connect_timeout",
"dbname",
"fallback_application_name",
"gssencmode",
"host",
"hostaddr",
"keepalives",
"keepalives_count",
"keepalives_idle",
"keepalives_interval",
"load_balance_hosts",
"options",
"passfile",
"password",
"port",
"replication",
"require_auth",
"requiressl",
"service",
"sslcert",
"sslcompression",
"sslcrl",
"sslcrldir",
"sslkey",
"sslmode",
"sslnegotiation",
"sslpassword",
"sslrootcert",
"sslsni",
"target_session_attrs",
"tcp_user_timeout",
"user",
}
def get_postgres_url():
postgres_url = os.getenv("POSTGRES_URL")
if not postgres_url:
raise ValueError("POSTGRES_URL is not set")
return clean_postgres_url(postgres_url)
def clean_postgres_url(postgres_url):
parts = urlsplit(postgres_url.strip().strip("\"'"))
query = urlencode(
[
(key, value)
for key, value in parse_qsl(parts.query, keep_blank_values=True)
if key in ALLOWED_POSTGRES_PARAMS
]
)
return urlunsplit((parts.scheme, parts.netloc, parts.path, query, parts.fragment))
def get_connection():
return psycopg.connect(get_postgres_url())
|