diff --git a/sqlite_utils/cli.py b/sqlite_utils/cli.py index a924439..9fe8d0a 100644 --- a/sqlite_utils/cli.py +++ b/sqlite_utils/cli.py @@ -44,6 +44,28 @@ from .utils import ( CONTEXT_SETTINGS = dict(help_option_names=["-h", "--help"]) + +def _register_db_for_cleanup(db): + """Register a database to be closed when the Click context is cleaned up.""" + ctx = click.get_current_context(silent=True) + if ctx is None: + return + if not hasattr(ctx, "_databases_to_close"): + ctx._databases_to_close = [] + ctx.call_on_close(lambda: _close_databases(ctx)) + ctx._databases_to_close.append(db) + + +def _close_databases(ctx): + """Close all databases registered for cleanup.""" + if hasattr(ctx, "_databases_to_close"): + for db in ctx._databases_to_close: + try: + db.close() + except Exception: + pass + + VALID_COLUMN_TYPES = ("INTEGER", "TEXT", "FLOAT", "REAL", "BLOB") UNICODE_ERROR = """ @@ -183,6 +205,7 @@ def tables( sqlite-utils tables trees.db """ db = sqlite_utils.Database(path) + _register_db_for_cleanup(db) _load_extensions(db, load_extension) headers = ["view" if views else "table"] if counts: @@ -309,6 +332,7 @@ def optimize(path, tables, no_vacuum, load_extension): sqlite-utils optimize chickens.db """ db = sqlite_utils.Database(path) + _register_db_for_cleanup(db) _load_extensions(db, load_extension) if not tables: tables = db.table_names(fts4=True) + db.table_names(fts5=True) @@ -336,6 +360,7 @@ def rebuild_fts(path, tables, load_extension): sqlite-utils rebuild-fts chickens.db chickens """ db = sqlite_utils.Database(path) + _register_db_for_cleanup(db) _load_extensions(db, load_extension) if not tables: tables = db.table_names(fts4=True) + db.table_names(fts5=True) @@ -360,6 +385,7 @@ def analyze(path, names): sqlite-utils analyze chickens.db """ db = sqlite_utils.Database(path) + _register_db_for_cleanup(db) try: if names: for name in names: @@ -384,7 +410,9 @@ def vacuum(path): \b sqlite-utils vacuum chickens.db """ - sqlite_utils.Database(path).vacuum() + db = sqlite_utils.Database(path) + _register_db_for_cleanup(db) + db.vacuum() @cli.command() @@ -403,6 +431,7 @@ def dump(path, load_extension): sqlite-utils dump chickens.db """ db = sqlite_utils.Database(path) + _register_db_for_cleanup(db) _load_extensions(db, load_extension) for line in db.iterdump(): click.echo(line) @@ -464,6 +493,7 @@ def add_column( sqlite-utils add-column chickens.db chickens weight float """ db = sqlite_utils.Database(path) + _register_db_for_cleanup(db) _load_extensions(db, load_extension) try: db[table].add_column( @@ -501,6 +531,7 @@ def add_foreign_key( sqlite-utils add-foreign-key my.db books author_id authors id """ db = sqlite_utils.Database(path) + _register_db_for_cleanup(db) _load_extensions(db, load_extension) try: db[table].add_foreign_key(column, other_table, other_column, ignore=ignore) @@ -528,6 +559,7 @@ def add_foreign_keys(path, foreign_key, load_extension): authors country_id countries id """ db = sqlite_utils.Database(path) + _register_db_for_cleanup(db) _load_extensions(db, load_extension) if len(foreign_key) % 4 != 0: raise click.ClickException( @@ -559,6 +591,7 @@ def index_foreign_keys(path, load_extension): sqlite-utils index-foreign-keys chickens.db """ db = sqlite_utils.Database(path) + _register_db_for_cleanup(db) _load_extensions(db, load_extension) db.index_foreign_keys() @@ -603,6 +636,7 @@ def create_index( sqlite-utils create-index chickens.db chickens -- -name """ db = sqlite_utils.Database(path) + _register_db_for_cleanup(db) _load_extensions(db, load_extension) # Treat -prefix as descending for columns columns = [] @@ -660,6 +694,7 @@ def enable_fts( fts_version = "FTS4" db = sqlite_utils.Database(path) + _register_db_for_cleanup(db) _load_extensions(db, load_extension) try: db[table].enable_fts( @@ -691,6 +726,7 @@ def populate_fts(path, table, column, load_extension): sqlite-utils populate-fts chickens.db chickens name """ db = sqlite_utils.Database(path) + _register_db_for_cleanup(db) _load_extensions(db, load_extension) db[table].populate_fts(column) @@ -712,6 +748,7 @@ def disable_fts(path, table, load_extension): sqlite-utils disable-fts chickens.db chickens """ db = sqlite_utils.Database(path) + _register_db_for_cleanup(db) _load_extensions(db, load_extension) db[table].disable_fts() @@ -734,6 +771,7 @@ def enable_wal(path, load_extension): """ for path_ in path: db = sqlite_utils.Database(path_) + _register_db_for_cleanup(db) _load_extensions(db, load_extension) db.enable_wal() @@ -756,6 +794,7 @@ def disable_wal(path, load_extension): """ for path_ in path: db = sqlite_utils.Database(path_) + _register_db_for_cleanup(db) _load_extensions(db, load_extension) db.disable_wal() @@ -777,6 +816,7 @@ def enable_counts(path, tables, load_extension): sqlite-utils enable-counts chickens.db """ db = sqlite_utils.Database(path) + _register_db_for_cleanup(db) _load_extensions(db, load_extension) if not tables: db.enable_counts() @@ -805,6 +845,7 @@ def reset_counts(path, load_extension): sqlite-utils reset-counts chickens.db """ db = sqlite_utils.Database(path) + _register_db_for_cleanup(db) _load_extensions(db, load_extension) db.reset_counts() @@ -964,6 +1005,7 @@ def insert_upsert_implementation( strict=False, ): db = sqlite_utils.Database(path) + _register_db_for_cleanup(db) _load_extensions(db, load_extension) _maybe_register_functions(db, functions) if (delimiter or quotechar or sniff or no_headers) and not tsv: @@ -1480,6 +1522,7 @@ def create_database(path, enable_wal, init_spatialite, load_extension): sqlite-utils create-database trees.db """ db = sqlite_utils.Database(path) + _register_db_for_cleanup(db) if enable_wal: db.enable_wal() @@ -1569,6 +1612,7 @@ def create_table( Valid column types are text, integer, float and blob. """ db = sqlite_utils.Database(path) + _register_db_for_cleanup(db) _load_extensions(db, load_extension) if len(columns) % 2 == 1: raise click.ClickException( @@ -1620,6 +1664,7 @@ def duplicate(path, table, new_table, ignore, load_extension): Create a duplicate of this table, copying across the schema and all row data. """ db = sqlite_utils.Database(path) + _register_db_for_cleanup(db) _load_extensions(db, load_extension) try: db[table].duplicate(new_table) @@ -1643,6 +1688,7 @@ def rename_table(path, table, new_name, ignore, load_extension): Rename this table. """ db = sqlite_utils.Database(path) + _register_db_for_cleanup(db) _load_extensions(db, load_extension) try: db.rename_table(table, new_name) @@ -1671,6 +1717,7 @@ def drop_table(path, table, ignore, load_extension): sqlite-utils drop-table chickens.db chickens """ db = sqlite_utils.Database(path) + _register_db_for_cleanup(db) _load_extensions(db, load_extension) try: db[table].drop(ignore=ignore) @@ -1707,6 +1754,7 @@ def create_view(path, view, select, ignore, replace, load_extension): 'select * from chickens where weight > 3' """ db = sqlite_utils.Database(path) + _register_db_for_cleanup(db) _load_extensions(db, load_extension) # Does view already exist? if view in db.view_names(): @@ -1741,6 +1789,7 @@ def drop_view(path, view, ignore, load_extension): sqlite-utils drop-view chickens.db heavy_chickens """ db = sqlite_utils.Database(path) + _register_db_for_cleanup(db) _load_extensions(db, load_extension) try: db[view].drop(ignore=ignore) @@ -1805,6 +1854,7 @@ def query( -p age 1 """ db = sqlite_utils.Database(path) + _register_db_for_cleanup(db) for alias, attach_path in attach: db.attach(alias, attach_path) _load_extensions(db, load_extension) @@ -1939,6 +1989,8 @@ def memory( sqlite-utils memory animals.csv --schema """ db = sqlite_utils.Database(memory=True) + if not return_db: + _register_db_for_cleanup(db) # If --dump or --save or --analyze used but no paths detected, assume SQL query is a path: if (dump or save or schema or analyze) and not paths: @@ -2004,6 +2056,7 @@ def memory( if save: db2 = sqlite_utils.Database(save) + _register_db_for_cleanup(db2) for line in db.iterdump(): db2.execute(line) return @@ -2140,6 +2193,7 @@ def search( sqlite-utils search data.db chickens lila """ db = sqlite_utils.Database(path) + _register_db_for_cleanup(db) _load_extensions(db, load_extension) # Check table exists table_obj = db[dbtable] @@ -2307,7 +2361,9 @@ def triggers( """ sql = "select name, tbl_name as \"table\", sql from sqlite_master where type = 'trigger'" if tables: - quote = sqlite_utils.Database(memory=True).quote + _quote_db = sqlite_utils.Database(memory=True) + _register_db_for_cleanup(_quote_db) + quote = _quote_db.quote sql += ' and "table" in ({})'.format( ", ".join(quote(table) for table in tables) ) @@ -2372,7 +2428,9 @@ def indexes( sqlite_master.type = 'table' """ if tables: - quote = sqlite_utils.Database(memory=True).quote + _quote_db = sqlite_utils.Database(memory=True) + _register_db_for_cleanup(_quote_db) + quote = _quote_db.quote sql += " and sqlite_master.name in ({})".format( ", ".join(quote(table) for table in tables) ) @@ -2415,6 +2473,7 @@ def schema( sqlite-utils schema trees.db """ db = sqlite_utils.Database(path) + _register_db_for_cleanup(db) _load_extensions(db, load_extension) if tables: for table in tables: @@ -2507,6 +2566,7 @@ def transform( --rename column2 column_renamed """ db = sqlite_utils.Database(path) + _register_db_for_cleanup(db) _load_extensions(db, load_extension) types = {} kwargs = {} @@ -2590,6 +2650,7 @@ def extract( sqlite-utils extract trees.db Street_Trees species """ db = sqlite_utils.Database(path) + _register_db_for_cleanup(db) _load_extensions(db, load_extension) kwargs = dict( columns=columns, @@ -2734,6 +2795,7 @@ def insert_files( yield row db = sqlite_utils.Database(path) + _register_db_for_cleanup(db) _load_extensions(db, load_extension) try: with db.conn: @@ -2792,6 +2854,7 @@ def analyze_tables( sqlite-utils analyze-tables data.db trees """ db = sqlite_utils.Database(path) + _register_db_for_cleanup(db) _load_extensions(db, load_extension) _analyze(db, tables, columns, save, common_limit, no_most, no_least) @@ -2991,6 +3054,7 @@ def convert( ): sqlite3.enable_callback_tracebacks(True) db = sqlite_utils.Database(db_path) + _register_db_for_cleanup(db) if output is not None and len(columns) > 1: raise click.ClickException("Cannot use --output with more than one column") if multi and len(columns) > 1: @@ -3133,6 +3197,7 @@ def add_geometry_column( By default, this command will try to load the SpatiaLite extension from usual paths. To load it from a specific path, use --load-extension.""" db = sqlite_utils.Database(db_path) + _register_db_for_cleanup(db) if not db[table].exists(): raise click.ClickException( "You must create a table before adding a geometry column" @@ -3165,6 +3230,7 @@ def create_spatial_index(db_path, table, column_name, load_extension): By default, this command will try to load the SpatiaLite extension from usual paths. To load it from a specific path, use --load-extension.""" db = sqlite_utils.Database(db_path) + _register_db_for_cleanup(db) if not db[table].exists(): raise click.ClickException( "You must create a table and add a geometry column before creating a spatial index" diff --git a/tests/conftest.py b/tests/conftest.py index 3932f05..a9a47bc 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -16,7 +16,9 @@ def pytest_configure(config): @pytest.fixture def fresh_db(): - return Database(memory=True) + db = Database(memory=True) + yield db + db.close() @pytest.fixture @@ -30,7 +32,8 @@ def existing_db(): INSERT INTO foo (text) values ("three"); """ ) - return database + yield database + database.close() @pytest.fixture @@ -38,4 +41,5 @@ def db_path(tmpdir): path = str(tmpdir / "test.db") db = sqlite3.connect(path) db.executescript(CREATE_TABLES) + db.close() return path diff --git a/tests/test_analyze_tables.py b/tests/test_analyze_tables.py index 9634cfc..2867bb5 100644 --- a/tests/test_analyze_tables.py +++ b/tests/test_analyze_tables.py @@ -44,6 +44,7 @@ def big_db_to_analyze_path(tmpdir): } ) db["stuff"].insert_all(to_insert) + db.close() return path @@ -137,6 +138,7 @@ def db_to_analyze_path(db_to_analyze, tmpdir): db = sqlite3.connect(path) sql = "\n".join(db_to_analyze.iterdump()) db.executescript(sql) + db.close() return path @@ -311,6 +313,7 @@ def test_analyze_table_validate_columns(tmpdir, args, expected_error): "age": 5, } ) + db.close() result = CliRunner().invoke( cli.cli, ["analyze-tables", path] + args, diff --git a/tests/test_cli.py b/tests/test_cli.py index 4198727..822748a 100644 --- a/tests/test_cli.py +++ b/tests/test_cli.py @@ -82,6 +82,7 @@ def test_tables_counts_and_columns(db_path): db = Database(db_path) with db.conn: db["lots"].insert_all([{"id": i, "age": i + 1} for i in range(30)]) + db.close() result = CliRunner().invoke(cli.cli, ["tables", "--counts", "--columns", db_path]) assert ( '[{"table": "Gosh", "count": 0, "columns": ["c1", "c2", "c3"]},\n' @@ -117,6 +118,7 @@ def test_tables_counts_and_columns_csv(db_path, format, expected): db = Database(db_path) with db.conn: db["lots"].insert_all([{"id": i, "age": i + 1} for i in range(30)]) + db.close() result = CliRunner().invoke( cli.cli, ["tables", "--counts", "--columns", format, db_path] ) @@ -127,6 +129,7 @@ def test_tables_schema(db_path): db = Database(db_path) with db.conn: db["lots"].insert_all([{"id": i, "age": i + 1} for i in range(30)]) + db.close() result = CliRunner().invoke(cli.cli, ["tables", "--schema", db_path]) assert ( '[{"table": "Gosh", "schema": "CREATE TABLE Gosh (c1 text, c2 text, c3 text)"},\n' @@ -188,6 +191,7 @@ def test_output_table(db_path, options, expected): for i in range(4) ] ) + db.close() result = CliRunner().invoke(cli.cli, ["rows", db_path, "rows"] + options) assert result.exit_code == 0 assert expected == result.output.strip() @@ -243,6 +247,7 @@ def test_create_index(db_path): CliRunner().invoke(cli.cli, create_index_unique_args + [option]).exit_code == 0 ) + db.close() def test_create_index_analyze(db_path): @@ -622,6 +627,7 @@ def test_optimize(db_path, tables): ) db["Gosh"].enable_fts(["c1", "c2", "c3"], fts_version="FTS4") db["Gosh2"].enable_fts(["c1", "c2", "c3"], fts_version="FTS5") + db.close() size_before_optimize = os.stat(db_path).st_size result = CliRunner().invoke(cli.cli, ["optimize", db_path] + tables) assert result.exit_code == 0 @@ -1450,6 +1456,7 @@ def test_drop_table_error(): with runner.isolated_filesystem(): db = Database("test.db") db["t"].create({"pk": int}, pk="pk") + db.close() result = runner.invoke( cli.cli, [ @@ -1474,6 +1481,7 @@ def test_drop_view(): db = Database("test.db") db.create_view("hello", "select 1") assert "hello" in db.view_names() + db.close() result = runner.invoke( cli.cli, [ @@ -1483,7 +1491,9 @@ def test_drop_view(): ], ) assert result.exit_code == 0 + db = Database("test.db") assert "hello" not in db.view_names() + db.close() def test_drop_view_error(): @@ -1491,6 +1501,7 @@ def test_drop_view_error(): with runner.isolated_filesystem(): db = Database("test.db") db["t"].create({"pk": int}, pk="pk") + db.close() result = runner.invoke( cli.cli, [ @@ -1728,10 +1739,13 @@ def test_transform(db_path, args, expected_schema): defaults={"age": 1}, pk="id", ) + db.close() result = CliRunner().invoke(cli.cli, ["transform", db_path, "dogs"] + args) print(result.output) assert result.exit_code == 0 + db = Database(db_path) schema = db["dogs"].schema + db.close() assert schema == expected_schema diff --git a/tests/test_cli_memory.py b/tests/test_cli_memory.py index 8483963..3d0e812 100644 --- a/tests/test_cli_memory.py +++ b/tests/test_cli_memory.py @@ -328,9 +328,11 @@ def test_memory_return_db(tmpdir): from sqlite_utils.cli import cli path = str(tmpdir / "dogs.csv") - open(path, "w").write("id,name\n1,Cleo") + with open(path, "w") as f: + f.write("id,name\n1,Cleo") with click.Context(cli) as ctx: db = ctx.invoke(cli.commands["memory"], paths=(path,), return_db=True) assert db.table_names() == ["dogs"] + db.close() diff --git a/tests/test_hypothesis.py b/tests/test_hypothesis.py index f12f865..54bdc9d 100644 --- a/tests/test_hypothesis.py +++ b/tests/test_hypothesis.py @@ -6,39 +6,39 @@ import sqlite_utils # SQLite integers are -(2^63) to 2^63 - 1 @given(st.integers(-9223372036854775808, 9223372036854775807)) def test_roundtrip_integers(integer): - db = sqlite_utils.Database(memory=True) - row = { - "integer": integer, - } - db["test"].insert(row) - assert list(db["test"].rows) == [row] + with sqlite_utils.Database(memory=True) as db: + row = { + "integer": integer, + } + db["test"].insert(row) + assert list(db["test"].rows) == [row] @given(st.text()) def test_roundtrip_text(text): - db = sqlite_utils.Database(memory=True) - row = { - "text": text, - } - db["test"].insert(row) - assert list(db["test"].rows) == [row] + with sqlite_utils.Database(memory=True) as db: + row = { + "text": text, + } + db["test"].insert(row) + assert list(db["test"].rows) == [row] @given(st.binary(max_size=1024 * 1024)) def test_roundtrip_binary(binary): - db = sqlite_utils.Database(memory=True) - row = { - "binary": binary, - } - db["test"].insert(row) - assert list(db["test"].rows) == [row] + with sqlite_utils.Database(memory=True) as db: + row = { + "binary": binary, + } + db["test"].insert(row) + assert list(db["test"].rows) == [row] @given(st.floats(allow_nan=False)) def test_roundtrip_floats(floats): - db = sqlite_utils.Database(memory=True) - row = { - "floats": floats, - } - db["test"].insert(row) - assert list(db["test"].rows) == [row] + with sqlite_utils.Database(memory=True) as db: + row = { + "floats": floats, + } + db["test"].insert(row) + assert list(db["test"].rows) == [row]