Claude Code helped fix a ton of .close() warnings, refs #692

https://gistpreview.github.io/?730f0c5dc38528a1dd0615f330bd5481

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"

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

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,

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
 
 

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()

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]