mirror of
https://github.com/simonw/sqlite-utils.git
synced 2026-09-12 03:24:23 +02:00
Fix all remaining resource warnings, refs #693
https://gistpreview.github.io/?0bb8e869b82f6ff0db647de755182502
This commit is contained in:
parent
77359bea30
commit
dc9947a5e1
7 changed files with 110 additions and 61 deletions
|
|
@ -905,7 +905,7 @@ def insert_upsert_options(*, require_pk=False):
|
|||
required=True,
|
||||
),
|
||||
click.argument("table"),
|
||||
click.argument("file", type=click.File("rb"), required=True),
|
||||
click.argument("file", type=click.File("rb", lazy=True), required=True),
|
||||
click.option(
|
||||
"--pk",
|
||||
help="Columns to use as the primary key, e.g. id",
|
||||
|
|
@ -2000,6 +2000,7 @@ def memory(
|
|||
for i, path in enumerate(paths):
|
||||
# Path may have a :format suffix
|
||||
fp = None
|
||||
should_close_fp = False
|
||||
if ":" in path and path.rsplit(":", 1)[-1].upper() in Format.__members__:
|
||||
path, suffix = path.rsplit(":", 1)
|
||||
format = Format[suffix.upper()]
|
||||
|
|
@ -2017,29 +2018,32 @@ def memory(
|
|||
file_table = stem
|
||||
stem_counts[stem] = stem_counts.get(stem, 1) + 1
|
||||
fp = file_path.open("rb")
|
||||
rows, format_used = rows_from_file(fp, format=format, encoding=encoding)
|
||||
tracker = None
|
||||
if format_used in (Format.CSV, Format.TSV) and not no_detect_types:
|
||||
tracker = TypeTracker()
|
||||
rows = tracker.wrap(rows)
|
||||
if flatten:
|
||||
rows = (_flatten(row) for row in rows)
|
||||
should_close_fp = True
|
||||
try:
|
||||
rows, format_used = rows_from_file(fp, format=format, encoding=encoding)
|
||||
tracker = None
|
||||
if format_used in (Format.CSV, Format.TSV) and not no_detect_types:
|
||||
tracker = TypeTracker()
|
||||
rows = tracker.wrap(rows)
|
||||
if flatten:
|
||||
rows = (_flatten(row) for row in rows)
|
||||
|
||||
db[file_table].insert_all(rows, alter=True)
|
||||
if tracker is not None:
|
||||
db[file_table].transform(types=tracker.types)
|
||||
# Add convenient t / t1 / t2 views
|
||||
view_names = ["t{}".format(i + 1)]
|
||||
if i == 0:
|
||||
view_names.append("t")
|
||||
for view_name in view_names:
|
||||
if not db[view_name].exists():
|
||||
db.create_view(
|
||||
view_name, "select * from {}".format(quote_identifier(file_table))
|
||||
)
|
||||
|
||||
if fp:
|
||||
fp.close()
|
||||
db[file_table].insert_all(rows, alter=True)
|
||||
if tracker is not None:
|
||||
db[file_table].transform(types=tracker.types)
|
||||
# Add convenient t / t1 / t2 views
|
||||
view_names = ["t{}".format(i + 1)]
|
||||
if i == 0:
|
||||
view_names.append("t")
|
||||
for view_name in view_names:
|
||||
if not db[view_name].exists():
|
||||
db.create_view(
|
||||
view_name,
|
||||
"select * from {}".format(quote_identifier(file_table)),
|
||||
)
|
||||
finally:
|
||||
if should_close_fp and fp:
|
||||
fp.close()
|
||||
|
||||
if analyze:
|
||||
_analyze(db, tables=None, columns=None, save=False)
|
||||
|
|
|
|||
|
|
@ -9,7 +9,29 @@ import json
|
|||
import os
|
||||
import sys
|
||||
from . import recipes
|
||||
from typing import Dict, cast, BinaryIO, Iterable, Optional, Tuple, Type
|
||||
from typing import Dict, cast, BinaryIO, Iterable, Iterator, Optional, Tuple, Type
|
||||
|
||||
|
||||
class _CloseableIterator(Iterator[dict]):
|
||||
"""Iterator wrapper that closes a file when iteration is complete."""
|
||||
|
||||
def __init__(self, iterator: Iterator[dict], closeable: io.IOBase):
|
||||
self._iterator = iterator
|
||||
self._closeable = closeable
|
||||
|
||||
def __iter__(self) -> "_CloseableIterator":
|
||||
return self
|
||||
|
||||
def __next__(self) -> dict:
|
||||
try:
|
||||
return next(self._iterator)
|
||||
except StopIteration:
|
||||
self._closeable.close()
|
||||
raise
|
||||
|
||||
def close(self) -> None:
|
||||
self._closeable.close()
|
||||
|
||||
|
||||
import click
|
||||
|
||||
|
|
@ -299,7 +321,8 @@ def rows_from_file(
|
|||
reader = csv.DictReader(decoded_fp, dialect=dialect)
|
||||
else:
|
||||
reader = csv.DictReader(decoded_fp)
|
||||
return _extra_key_strategy(reader, ignore_extras, extras_key), Format.CSV
|
||||
rows = _extra_key_strategy(reader, ignore_extras, extras_key)
|
||||
return _CloseableIterator(iter(rows), decoded_fp), Format.CSV
|
||||
elif format == Format.TSV:
|
||||
rows = rows_from_file(
|
||||
fp, format=Format.CSV, dialect=csv.excel_tab, encoding=encoding
|
||||
|
|
|
|||
|
|
@ -19,9 +19,11 @@ def _supports_pragma_function_list():
|
|||
db = Database(memory=True)
|
||||
try:
|
||||
db.execute("select * from pragma_function_list()")
|
||||
return True
|
||||
except Exception:
|
||||
return False
|
||||
return True
|
||||
finally:
|
||||
db.close()
|
||||
|
||||
|
||||
def _has_compiled_ext():
|
||||
|
|
|
|||
|
|
@ -2,6 +2,14 @@ from sqlite_utils.db import Index, View, Database, XIndex, XIndexColumn
|
|||
import pytest
|
||||
|
||||
|
||||
def _check_supports_strict():
|
||||
"""Check if SQLite supports strict tables without leaking the database."""
|
||||
db = Database(memory=True)
|
||||
result = db.supports_strict
|
||||
db.close()
|
||||
return result
|
||||
|
||||
|
||||
def test_table_names(existing_db):
|
||||
assert ["foo"] == existing_db.table_names()
|
||||
|
||||
|
|
@ -282,7 +290,7 @@ def test_use_rowid(fresh_db):
|
|||
|
||||
|
||||
@pytest.mark.skipif(
|
||||
not Database(memory=True).supports_strict,
|
||||
not _check_supports_strict(),
|
||||
reason="Needs SQLite version that supports strict",
|
||||
)
|
||||
@pytest.mark.parametrize(
|
||||
|
|
|
|||
|
|
@ -9,9 +9,11 @@ def _supports_pragma_function_list():
|
|||
db = Database(memory=True)
|
||||
try:
|
||||
db.execute("select * from pragma_function_list()")
|
||||
return True
|
||||
except Exception:
|
||||
return False
|
||||
return True
|
||||
finally:
|
||||
db.close()
|
||||
|
||||
|
||||
def test_register_commands():
|
||||
|
|
|
|||
|
|
@ -14,8 +14,11 @@ def test_recreate_ignored_for_in_memory():
|
|||
|
||||
def test_recreate_not_allowed_for_connection():
|
||||
conn = sqlite3.connect(":memory:")
|
||||
with pytest.raises(AssertionError):
|
||||
Database(conn, recreate=True)
|
||||
try:
|
||||
with pytest.raises(AssertionError):
|
||||
Database(conn, recreate=True)
|
||||
finally:
|
||||
conn.close()
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
|
|
|
|||
|
|
@ -42,39 +42,46 @@ def test_register_function_deterministic(fresh_db):
|
|||
|
||||
|
||||
def test_register_function_deterministic_tries_again_if_exception_raised(fresh_db):
|
||||
# Save the original connection so we can close it later
|
||||
original_conn = fresh_db.conn
|
||||
fresh_db.conn = MagicMock()
|
||||
fresh_db.conn.create_function = MagicMock()
|
||||
|
||||
@fresh_db.register_function(deterministic=True)
|
||||
def to_lower_2(s):
|
||||
return s.lower()
|
||||
try:
|
||||
|
||||
fresh_db.conn.create_function.assert_called_with(
|
||||
"to_lower_2", 1, to_lower_2, deterministic=True
|
||||
)
|
||||
@fresh_db.register_function(deterministic=True)
|
||||
def to_lower_2(s):
|
||||
return s.lower()
|
||||
|
||||
first = True
|
||||
fresh_db.conn.create_function.assert_called_with(
|
||||
"to_lower_2", 1, to_lower_2, deterministic=True
|
||||
)
|
||||
|
||||
def side_effect(*args, **kwargs):
|
||||
# Raise exception only first time this is called
|
||||
nonlocal first
|
||||
if first:
|
||||
first = False
|
||||
raise sqlite3.NotSupportedError()
|
||||
first = True
|
||||
|
||||
# But if sqlite3.NotSupportedError is raised, it tries again
|
||||
fresh_db.conn.create_function.reset_mock()
|
||||
fresh_db.conn.create_function.side_effect = side_effect
|
||||
def side_effect(*args, **kwargs):
|
||||
# Raise exception only first time this is called
|
||||
nonlocal first
|
||||
if first:
|
||||
first = False
|
||||
raise sqlite3.NotSupportedError()
|
||||
|
||||
@fresh_db.register_function(deterministic=True)
|
||||
def to_lower_3(s):
|
||||
return s.lower()
|
||||
# But if sqlite3.NotSupportedError is raised, it tries again
|
||||
fresh_db.conn.create_function.reset_mock()
|
||||
fresh_db.conn.create_function.side_effect = side_effect
|
||||
|
||||
# Should have been called once with deterministic=True and once without
|
||||
assert fresh_db.conn.create_function.call_args_list == [
|
||||
call("to_lower_3", 1, to_lower_3, deterministic=True),
|
||||
call("to_lower_3", 1, to_lower_3),
|
||||
]
|
||||
@fresh_db.register_function(deterministic=True)
|
||||
def to_lower_3(s):
|
||||
return s.lower()
|
||||
|
||||
# Should have been called once with deterministic=True and once without
|
||||
assert fresh_db.conn.create_function.call_args_list == [
|
||||
call("to_lower_3", 1, to_lower_3, deterministic=True),
|
||||
call("to_lower_3", 1, to_lower_3),
|
||||
]
|
||||
finally:
|
||||
# Close the original connection that was replaced with the mock
|
||||
original_conn.close()
|
||||
|
||||
|
||||
def test_register_function_replace(fresh_db):
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue