-
Notifications
You must be signed in to change notification settings - Fork 21
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
Merge pull request #176 from WildMeOrg/sql/migration-script
Connecting logic for Postgres URI and a migration script from SQLite to Postgres
- Loading branch information
Showing
21 changed files
with
1,090 additions
and
473 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,95 @@ | ||
| # -*- coding: utf-8 -*- | ||
| import logging | ||
| import sys | ||
| import click | ||
|
|
||
| from wbia.dtool.copy_sqlite_to_postgres import ( | ||
| SqliteDatabaseInfo, | ||
| PostgresDatabaseInfo, | ||
| compare_databases, | ||
| DEFAULT_CHECK_PC, | ||
| DEFAULT_CHECK_MIN, | ||
| DEFAULT_CHECK_MAX, | ||
| ) | ||
|
|
||
|
|
||
| logger = logging.getLogger('wbia') | ||
|
|
||
|
|
||
| @click.command() | ||
| @click.option( | ||
| '--db-dir', | ||
| multiple=True, | ||
| help='SQLite databases location', | ||
| ) | ||
| @click.option( | ||
| '--sqlite-uri', | ||
| multiple=True, | ||
| help='SQLite database URI (e.g. sqlite:////path.sqlite3)', | ||
| ) | ||
| @click.option( | ||
| '--pg-uri', | ||
| multiple=True, | ||
| help='Postgres connection URI (e.g. postgresql://user:pass@host)', | ||
| ) | ||
| @click.option( | ||
| '--check-pc', | ||
| type=float, | ||
| default=DEFAULT_CHECK_PC, | ||
| help=f'Percentage of table to check, default {DEFAULT_CHECK_PC} ({int(DEFAULT_CHECK_PC * 100)}% of the table)', | ||
| ) | ||
| @click.option( | ||
| '--check-max', | ||
| type=int, | ||
| default=DEFAULT_CHECK_MAX, | ||
| help=f'Maximum number of rows to check, default {DEFAULT_CHECK_MAX} (0 for no limit)', | ||
| ) | ||
| @click.option( | ||
| '--check-min', | ||
| type=int, | ||
| default=DEFAULT_CHECK_MIN, | ||
| help=f'Minimum number of rows to check, default {DEFAULT_CHECK_MIN}', | ||
| ) | ||
| @click.option( | ||
| '-v', | ||
| '--verbose', | ||
| is_flag=True, | ||
| default=False, | ||
| help='Show debug messages', | ||
| ) | ||
| def main(db_dir, sqlite_uri, pg_uri, check_pc, check_max, check_min, verbose): | ||
| if verbose: | ||
| logger.setLevel(logging.DEBUG) | ||
| else: | ||
| logger.setLevel(logging.INFO) | ||
|
|
||
| logger.addHandler(logging.StreamHandler()) | ||
|
|
||
| if len(db_dir) + len(sqlite_uri) + len(pg_uri) != 2: | ||
| raise click.BadParameter('exactly 2 db_dir or sqlite_uri or pg_uri must be given') | ||
| db_infos = [] | ||
| for db_dir_ in db_dir: | ||
| db_infos.append(SqliteDatabaseInfo(db_dir_)) | ||
| for sqlite_uri_ in sqlite_uri: | ||
| db_infos.append(SqliteDatabaseInfo(sqlite_uri_)) | ||
| for pg_uri_ in pg_uri: | ||
| db_infos.append(PostgresDatabaseInfo(pg_uri_)) | ||
| exact = not (sqlite_uri and pg_uri) | ||
| differences = compare_databases( | ||
| *db_infos, | ||
| exact=exact, | ||
| check_pc=check_pc, | ||
| check_max=check_max, | ||
| check_min=check_min, | ||
| ) | ||
| if differences: | ||
| click.echo(f'Databases {db_infos[0]} and {db_infos[1]} are different:') | ||
| for line in differences: | ||
| click.echo(line) | ||
| sys.exit(1) | ||
| else: | ||
| click.echo(f'Databases {db_infos[0]} and {db_infos[1]} are the same') | ||
|
|
||
|
|
||
| if __name__ == '__main__': | ||
| main() |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,143 @@ | ||
| # -*- coding: utf-8 -*- | ||
| import logging | ||
| import re | ||
| import subprocess | ||
| import sys | ||
| from pathlib import Path | ||
|
|
||
| import click | ||
| import sqlalchemy | ||
|
|
||
| from wbia.dtool.copy_sqlite_to_postgres import ( | ||
| copy_sqlite_to_postgres, | ||
| SqliteDatabaseInfo, | ||
| PostgresDatabaseInfo, | ||
| compare_databases, | ||
| ) | ||
|
|
||
|
|
||
| logger = logging.getLogger('wbia') | ||
|
|
||
|
|
||
| @click.command() | ||
| @click.option( | ||
| '--db-dir', required=True, type=click.Path(exists=True), help='database location' | ||
| ) | ||
| @click.option( | ||
| '--db-uri', | ||
| required=True, | ||
| help='Postgres connection URI (e.g. postgres://user:pass@host)', | ||
| ) | ||
| @click.option( | ||
| '--force', | ||
| is_flag=True, | ||
| default=False, | ||
| help='Delete all tables in the public schema in postgres', | ||
| ) | ||
| @click.option( | ||
| '-v', | ||
| '--verbose', | ||
| is_flag=True, | ||
| default=False, | ||
| help='Show debug messages', | ||
| ) | ||
| def main(db_dir, db_uri, force, verbose): | ||
| """""" | ||
| # Set up logging | ||
| if verbose: | ||
| logger.setLevel(logging.DEBUG) | ||
| else: | ||
| logger.setLevel(logging.INFO) | ||
| logger.addHandler(logging.StreamHandler()) | ||
|
|
||
| logger.info(f'using {db_dir} ...') | ||
|
|
||
| # Create the database if it doesn't exist | ||
| engine = sqlalchemy.create_engine(db_uri) | ||
| try: | ||
| engine.connect() | ||
| except sqlalchemy.exc.OperationalError as e: | ||
| m = re.search(r'database "([^"]*)" does not exist', str(e)) | ||
| if m: | ||
| dbname = m.group(1) | ||
| engine = sqlalchemy.create_engine(db_uri.rsplit('/', 1)[0]) | ||
| logger.info(f'Creating "{dbname}"...') | ||
| engine.execution_options(isolation_level='AUTOCOMMIT').execute( | ||
| f'CREATE DATABASE {dbname}' | ||
| ) | ||
| else: | ||
| raise | ||
| finally: | ||
| engine.dispose() | ||
|
|
||
| # Check that the database hasn't already been migrated. | ||
| db_infos = [ | ||
| SqliteDatabaseInfo(Path(db_dir)), | ||
| PostgresDatabaseInfo(db_uri), | ||
| ] | ||
| differences = compare_databases(*db_infos) | ||
|
|
||
| if not differences: | ||
| logger.info('Database already migrated') | ||
| sys.exit(0) | ||
|
|
||
| # Make sure there are no tables in the public schema in postgresql | ||
| # because we're using it as the workspace for the migration | ||
| if 'public' in db_infos[1].get_schema(): | ||
| table_names = [ | ||
| t for schema, t in db_infos[1].get_table_names() if schema == 'public' | ||
| ] | ||
| if not force: | ||
| click.echo( | ||
| f'Tables in public schema in postgres database: {", ".join(table_names)}' | ||
| ) | ||
| click.echo('Use --force to remove the tables in public schema') | ||
| sys.exit(1) | ||
| else: | ||
| click.echo(f'Dropping all tables in public schema: {", ".join(table_names)}') | ||
| for table_name in table_names: | ||
| db_infos[1].engine.execute(f'DROP TABLE {table_name} CASCADE') | ||
|
|
||
| # Migrate | ||
| problems = {} | ||
| with click.progressbar(length=100000, show_eta=True) as bar: | ||
| for path, completed_future, db_size, total_size in copy_sqlite_to_postgres( | ||
| Path(db_dir), db_uri | ||
| ): | ||
| try: | ||
| completed_future.result() | ||
| except Exception as exc: | ||
| logger.info( | ||
| f'\nfailed while processing {str(path)}\n{completed_future.exception()}' | ||
| ) | ||
| problems[path] = exc | ||
| else: | ||
| logger.info(f'\nfinished processing {str(path)}') | ||
| finally: | ||
| bar.update(int(db_size / total_size * bar.length)) | ||
|
|
||
| # Report problems | ||
| for path, exc in problems.items(): | ||
| logger.info('*' * 60) | ||
| logger.info(f'There was a problem migrating {str(path)}') | ||
| logger.exception(exc) | ||
| if isinstance(exc, subprocess.CalledProcessError): | ||
| logger.info('-' * 30) | ||
| logger.info(exc.stdout.decode()) | ||
|
|
||
| # Verify the migration | ||
| differences = compare_databases(*db_infos) | ||
|
|
||
| if differences: | ||
| logger.info(f'Databases {db_infos[0]} and {db_infos[1]} are different:') | ||
| for line in differences: | ||
| logger.info(line) | ||
| sys.exit(1) | ||
| else: | ||
| logger.info(f'Database {db_infos[0]} successfully migrated to {db_infos[1]}') | ||
|
|
||
| sys.exit(0) | ||
|
|
||
|
|
||
| if __name__ == '__main__': | ||
| main() |
Oops, something went wrong.