Changeset View
Changeset View
Standalone View
Standalone View
swh/core/cli/db.py
#!/usr/bin/env python3 | #!/usr/bin/env python3 | ||||
# Copyright (C) 2018-2020 The Software Heritage developers | # Copyright (C) 2018-2020 The Software Heritage developers | ||||
# See the AUTHORS file at the top-level directory of this distribution | # See the AUTHORS file at the top-level directory of this distribution | ||||
# License: GNU General Public License version 3, or any later version | # License: GNU General Public License version 3, or any later version | ||||
# See top-level LICENSE file for more information | # See top-level LICENSE file for more information | ||||
import logging | import logging | ||||
from os import environ, path | from os import environ, path | ||||
from typing import Collection, Tuple | from typing import Collection, Optional, Tuple | ||||
import warnings | import warnings | ||||
import click | import click | ||||
from swh.core.cli import CONTEXT_SETTINGS | from swh.core.cli import CONTEXT_SETTINGS | ||||
from swh.core.cli import swh as swh_cli_group | from swh.core.cli import swh as swh_cli_group | ||||
warnings.filterwarnings("ignore") # noqa prevent psycopg from telling us sh*t | warnings.filterwarnings("ignore") # noqa prevent psycopg from telling us sh*t | ||||
▲ Show 20 Lines • Show All 69 Lines • ▼ Show 20 Lines | |||||
@click.argument("module", required=True) | @click.argument("module", required=True) | ||||
@click.option( | @click.option( | ||||
"--db-name", | "--db-name", | ||||
"-d", | "-d", | ||||
help="Database name.", | help="Database name.", | ||||
default="softwareheritage-dev", | default="softwareheritage-dev", | ||||
show_default=True, | show_default=True, | ||||
) | ) | ||||
def db_init(module, db_name): | @click.option( | ||||
"--flavor", help="Database flavor.", default=None, | |||||
douardda: not a blocker, but would it be possible to have this a click.Choice option? | |||||
Done Inline ActionsI wish it was possible, but in general it depends on which database you're trying to initialize, and you'd have to parse their SQL files to know which flavors exist. The function does:
olasd: I wish it was possible, but in general it depends on which database you're trying to initialize… | |||||
) | |||||
def db_init(module, db_name, flavor): | |||||
"""Initialize a database for the Software Heritage <module>. | """Initialize a database for the Software Heritage <module>. | ||||
Example: | Example: | ||||
swh db init -d swh-test storage | swh db init -d swh-test storage | ||||
If you want to specify non-default postgresql connection parameters, | If you want to specify non-default postgresql connection parameters, | ||||
please provide them using standard environment variables. | please provide them using standard environment variables. | ||||
See psql(1) man page (section ENVIRONMENTS) for details. | See psql(1) man page (section ENVIRONMENTS) for details. | ||||
Example: | Examples: | ||||
PGPORT=5434 swh db-init indexer | PGPORT=5434 swh db init indexer | ||||
swh db init -d postgresql://user:passwd@pghost:5433/swh-storage storage | swh db init -d postgresql://user:passwd@pghost:5433/swh-storage storage | ||||
swh db init --flavor read_replica -d swh-storage storage | |||||
""" | """ | ||||
logger.debug("db_init %s dn_name=%s", module, db_name) | logger.debug("db_init %s flavor=%s dn_name=%s", module, flavor, db_name) | ||||
initialized, dbversion = populate_database_for_package(module, db_name) | initialized, dbversion, dbflavor = populate_database_for_package( | ||||
module, db_name, flavor | |||||
) | |||||
# TODO: Ideally migrate the version from db_version to the latest | # TODO: Ideally migrate the version from db_version to the latest | ||||
# db version | # db version | ||||
click.secho( | click.secho( | ||||
"DONE database for {} {} at version {}".format( | "DONE database for {} {}{} at version {}".format( | ||||
module, "initialized" if initialized else "exists", dbversion | module, | ||||
"initialized" if initialized else "exists", | |||||
f" (flavor {dbflavor})" if dbflavor is not None else "", | |||||
dbversion, | |||||
), | ), | ||||
Done Inline Actionsshouldn't the "already exists" message also print the flavor? douardda: shouldn't the "already exists" message also print the flavor? | |||||
Done Inline ActionsYeah, it should. This means adding a new function to "get" the database flavor, and having an even larger populate_database_... return type, which is annoying and why I didn't do it initially. olasd: Yeah, it should. This means adding a new function to "get" the database flavor, and having an… | |||||
fg="green", | fg="green", | ||||
bold=True, | bold=True, | ||||
) | ) | ||||
if flavor is not None and dbflavor != flavor: | |||||
click.secho( | |||||
f"WARNING requested flavor '{flavor}' != recorded flavor '{dbflavor}'", | |||||
fg="red", | |||||
bold=True, | |||||
) | |||||
def get_sql_for_package(modname): | def get_sql_for_package(modname): | ||||
import glob | import glob | ||||
from importlib import import_module | from importlib import import_module | ||||
from swh.core.utils import numfile_sortkey as sortkey | from swh.core.utils import numfile_sortkey as sortkey | ||||
if not modname.startswith("swh."): | if not modname.startswith("swh."): | ||||
modname = "swh.{}".format(modname) | modname = "swh.{}".format(modname) | ||||
try: | try: | ||||
m = import_module(modname) | m = import_module(modname) | ||||
except ImportError: | except ImportError: | ||||
raise click.BadParameter("Unable to load module {}".format(modname)) | raise click.BadParameter("Unable to load module {}".format(modname)) | ||||
sqldir = path.join(path.dirname(m.__file__), "sql") | sqldir = path.join(path.dirname(m.__file__), "sql") | ||||
if not path.isdir(sqldir): | if not path.isdir(sqldir): | ||||
raise click.BadParameter( | raise click.BadParameter( | ||||
"Module {} does not provide a db schema " "(no sql/ dir)".format(modname) | "Module {} does not provide a db schema " "(no sql/ dir)".format(modname) | ||||
) | ) | ||||
return sorted(glob.glob(path.join(sqldir, "*.sql")), key=sortkey) | return sorted(glob.glob(path.join(sqldir, "*.sql")), key=sortkey) | ||||
def populate_database_for_package(modname: str, conninfo: str) -> Tuple[bool, int]: | def populate_database_for_package( | ||||
modname: str, conninfo: str, flavor: Optional[str] = None | |||||
) -> Tuple[bool, int, Optional[str]]: | |||||
"""Populate the database, pointed at with `conninfo`, using the SQL files found in | """Populate the database, pointed at with `conninfo`, using the SQL files found in | ||||
the package `modname`. | the package `modname`. | ||||
Args: | Args: | ||||
modname: Name of the module of which we're loading the files | modname: Name of the module of which we're loading the files | ||||
conninfo: connection info string for the SQL database | conninfo: connection info string for the SQL database | ||||
flavor: the module-specific flavor which we want to initialize the database under | |||||
Returns: | Returns: | ||||
Tuple with two elements: whether the database has been initialized; the current | Tuple with three elements: whether the database has been initialized; the current | ||||
version of the database. | version of the database; if it exists, the flavor of the database. | ||||
""" | """ | ||||
from swh.core.db.db_utils import swh_db_version | from swh.core.db.db_utils import swh_db_flavor, swh_db_version | ||||
current_version = swh_db_version(conninfo) | current_version = swh_db_version(conninfo) | ||||
if current_version is not None: | if current_version is not None: | ||||
return False, current_version | dbflavor = swh_db_flavor(conninfo) | ||||
return False, current_version, dbflavor | |||||
sqlfiles = get_sql_for_package(modname) | sqlfiles = get_sql_for_package(modname) | ||||
sqlfiles = [fname for fname in sqlfiles if "-superuser-" not in fname] | sqlfiles = [fname for fname in sqlfiles if "-superuser-" not in fname] | ||||
execute_sqlfiles(sqlfiles, conninfo) | execute_sqlfiles(sqlfiles, conninfo, flavor) | ||||
current_version = swh_db_version(conninfo) | current_version = swh_db_version(conninfo) | ||||
assert current_version is not None | assert current_version is not None | ||||
return True, current_version | dbflavor = swh_db_flavor(conninfo) | ||||
return True, current_version, dbflavor | |||||
def create_database_for_package( | def create_database_for_package( | ||||
modname: str, conninfo: str, template: str = "template1" | modname: str, conninfo: str, template: str = "template1" | ||||
): | ): | ||||
"""Create the database pointed at with `conninfo`, and initialize it using | """Create the database pointed at with `conninfo`, and initialize it using | ||||
-superuser- SQL files found in the package `modname`. | -superuser- SQL files found in the package `modname`. | ||||
Show All 31 Lines | ): | ||||
# the remaining initialization process -- running -superuser- SQL files -- | # the remaining initialization process -- running -superuser- SQL files -- | ||||
# is done using the given conninfo, thus connecting to the newly created | # is done using the given conninfo, thus connecting to the newly created | ||||
# database | # database | ||||
sqlfiles = get_sql_for_package(modname) | sqlfiles = get_sql_for_package(modname) | ||||
sqlfiles = [fname for fname in sqlfiles if "-superuser-" in fname] | sqlfiles = [fname for fname in sqlfiles if "-superuser-" in fname] | ||||
execute_sqlfiles(sqlfiles, conninfo) | execute_sqlfiles(sqlfiles, conninfo) | ||||
def execute_sqlfiles(sqlfiles: Collection[str], conninfo: str): | def execute_sqlfiles( | ||||
sqlfiles: Collection[str], conninfo: str, flavor: Optional[str] = None | |||||
): | |||||
"""Execute a list of SQL files on the database pointed at with `conninfo`. | """Execute a list of SQL files on the database pointed at with `conninfo`. | ||||
Args: | Args: | ||||
sqlfiles: List of SQL files to execute | sqlfiles: List of SQL files to execute | ||||
conninfo: connection info string for the SQL database | conninfo: connection info string for the SQL database | ||||
flavor: the database flavor to initialize | |||||
""" | """ | ||||
import subprocess | import subprocess | ||||
for sqlfile in sqlfiles: | psql_command = [ | ||||
logger.debug(f"execute SQL file {sqlfile} db_name={conninfo}") | |||||
subprocess.check_call( | |||||
[ | |||||
"psql", | "psql", | ||||
"--quiet", | "--quiet", | ||||
"--no-psqlrc", | "--no-psqlrc", | ||||
"-v", | "-v", | ||||
"ON_ERROR_STOP=1", | "ON_ERROR_STOP=1", | ||||
"-d", | "-d", | ||||
conninfo, | conninfo, | ||||
"-f", | |||||
sqlfile, | |||||
] | ] | ||||
flavor_set = False | |||||
for sqlfile in sqlfiles: | |||||
logger.debug(f"execute SQL file {sqlfile} db_name={conninfo}") | |||||
subprocess.check_call(psql_command + ["-f", sqlfile]) | |||||
if flavor is not None and not flavor_set and sqlfile.endswith("-flavor.sql"): | |||||
logger.debug("Setting database flavor %s", flavor) | |||||
query = f"insert into dbflavor (flavor) values ('{flavor}')" | |||||
subprocess.check_call(psql_command + ["-c", query]) | |||||
flavor_set = True | |||||
if flavor is not None and not flavor_set: | |||||
logger.warn( | |||||
"Asked for flavor %s, but module does not support database flavors", flavor, | |||||
) | ) |
not a blocker, but would it be possible to have this a click.Choice option?