1from __future__ import annotations
  2
  3# Copyright (c) Kenneth Reitz & individual contributors
  4# All rights reserved.
  5# Redistribution and use in source and binary forms, with or without modification,
  6# are permitted provided that the following conditions are met:
  7#     1. Redistributions of source code must retain the above copyright notice,
  8#        this list of conditions and the following disclaimer.
  9#     2. Redistributions in binary form must reproduce the above copyright
 10#        notice, this list of conditions and the following disclaimer in the
 11#        documentation and/or other materials provided with the distribution.
 12#     3. Neither the name of Plain nor the names of its contributors may be used
 13#        to endorse or promote products derived from this software without
 14#        specific prior written permission.
 15# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND
 16# ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED
 17# WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
 18# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR CONTRIBUTORS BE LIABLE FOR
 19# ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES
 20# (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;
 21# LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON
 22# ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
 23# (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS
 24# SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
 25import urllib.parse as urlparse
 26from typing import Any, TypedDict
 27
 28SCHEMES = {"postgres", "postgresql", "pgsql"}
 29
 30
 31class DatabaseConfig(TypedDict, total=False):
 32    HOST: str
 33    DATABASE: str  # Required (validated in _configure_settings)
 34    OPTIONS: dict[str, Any]
 35    PASSWORD: str
 36    PORT: int | None
 37    USER: str
 38
 39
 40# Register database schemes in URLs.
 41for scheme in SCHEMES:
 42    urlparse.uses_netloc.append(scheme)
 43
 44
 45def parse_database_url(url: str) -> DatabaseConfig:
 46    """Parse a database URL into a fully-populated `DatabaseConfig`.
 47
 48    All keys are present; empty values are `""` (strings), `None` (PORT),
 49    or `{}` (OPTIONS) — callers can index without `.get(...)` defaults.
 50    """
 51    spliturl = urlparse.urlsplit(url)
 52
 53    if spliturl.scheme not in SCHEMES:
 54        raise ValueError(
 55            f"No support for '{spliturl.scheme}'. We support: {', '.join(sorted(SCHEMES))}"
 56        )
 57
 58    path = spliturl.path[1:]
 59    query = urlparse.parse_qs(spliturl.query)
 60
 61    # Handle percent-encoded hostnames (e.g. socket paths).
 62    hostname = spliturl.hostname or ""
 63    if "%" in hostname:
 64        # Use netloc to avoid lowercased paths, strip credentials if present.
 65        hostname = spliturl.netloc
 66        if "@" in hostname:
 67            hostname = hostname.rsplit("@", 1)[1]
 68        hostname = urlparse.unquote(hostname)
 69
 70    return {
 71        "DATABASE": urlparse.unquote(path or ""),
 72        "USER": urlparse.unquote(spliturl.username or ""),
 73        "PASSWORD": urlparse.unquote(spliturl.password or ""),
 74        "HOST": hostname,
 75        "PORT": spliturl.port,
 76        "OPTIONS": {key: values[-1] for key, values in query.items()},
 77    }
 78
 79
 80_CLI_FLAGS: list[tuple[str, str]] = [("USER", "-U"), ("HOST", "-h"), ("PORT", "-p")]
 81_CLI_OPTION_ENV_VARS: dict[str, str] = {
 82    "passfile": "PGPASSFILE",
 83    "sslmode": "PGSSLMODE",
 84    "sslrootcert": "PGSSLROOTCERT",
 85    "sslcert": "PGSSLCERT",
 86    "sslkey": "PGSSLKEY",
 87}
 88
 89
 90def postgres_cli_args(config: DatabaseConfig) -> list[str]:
 91    """Build connection flags for libpq-based tools (psql, pg_dump, pg_restore)."""
 92    args: list[str] = []
 93    for key, flag in _CLI_FLAGS:
 94        if value := config.get(key):
 95            args += [flag, str(value)]
 96    return args
 97
 98
 99def postgres_cli_env(config: DatabaseConfig) -> dict[str, str]:
100    """Build env vars for libpq-based tools (psql, pg_dump, pg_restore)."""
101    env: dict[str, str] = {}
102    if password := config.get("PASSWORD"):
103        env["PGPASSWORD"] = str(password)
104    options = config.get("OPTIONS", {})
105    for option_key, env_var in _CLI_OPTION_ENV_VARS.items():
106        if value := options.get(option_key):
107            env[env_var] = str(value)
108    return env
109
110
111def build_database_url(config: DatabaseConfig) -> str:
112    """Build a database URL from a configuration dictionary."""
113    options = config.get("OPTIONS", {})
114    query = urlparse.urlencode(list(options.items()))
115
116    user = urlparse.quote(str(config.get("USER", "")))
117    password = urlparse.quote(str(config.get("PASSWORD", "")))
118    host = config.get("HOST", "")
119    port = config.get("PORT")
120    name = urlparse.quote(str(config.get("DATABASE", "")))
121
122    netloc = ""
123    if user or password:
124        netloc += user
125        if password:
126            netloc += f":{password}"
127        netloc += "@"
128    netloc += host
129    if port:
130        netloc += f":{port}"
131
132    return urlparse.urlunsplit(("postgresql", netloc, f"/{name}", query, ""))
133
134
135def replace_database_name(url: str, name: str) -> str:
136    """Return the URL with the database name (path segment) replaced.
137
138    Preserves scheme, netloc, query string, and fragment exactly — only the
139    path changes. Avoids the round-trip through parse/build, which normalizes
140    the scheme and collapses duplicate query keys.
141    """
142    spliturl = urlparse.urlsplit(url)
143    return urlparse.urlunsplit(spliturl._replace(path=f"/{urlparse.quote(name)}"))