1
0
mirror of https://github.com/arsenetar/dupeguru.git synced 2024-10-30 05:15:57 +00:00
dupeguru/core/pe/cache_sqlite.py
Andrew Senetar a81069be61
fix: Photo matching fixes
- Correct bad query introduced in rotation matching
- Promote get_orientation from "private" on photo class
- Fix prepare_pictures to only generate the needed blocks, add check for missing blocks when rotation matchin is true
- Fix cache test inputs to match schema
2024-05-11 00:11:27 -07:00

200 lines
7.2 KiB
Python

# Copyright 2016 Virgil Dupras
#
# This software is licensed under the "GPLv3" License as described in the "LICENSE" file,
# which should be included with this package. The terms are also available at
# http://www.gnu.org/licenses/gpl-3.0.html
import os
import os.path as op
import logging
import sqlite3 as sqlite
from core.pe.cache import bytes_to_colors, colors_to_bytes
class SqliteCache:
"""A class to cache picture blocks in a sqlite backend."""
schema_version = 2
schema_version_description = "Added blocks for all 8 orientations."
create_table_query = (
"CREATE TABLE IF NOT EXISTS "
"pictures(path TEXT, mtime_ns INTEGER, blocks BLOB, blocks2 BLOB, blocks3 BLOB, "
"blocks4 BLOB, blocks5 BLOB, blocks6 BLOB, blocks7 BLOB, blocks8 BLOB)"
)
create_index_query = "CREATE INDEX IF NOT EXISTS idx_path on pictures (path)"
drop_table_query = "DROP TABLE IF EXISTS pictures"
drop_index_query = "DROP INDEX IF EXISTS idx_path"
def __init__(self, db=":memory:", readonly=False):
# readonly is not used in the sqlite version of the cache
self.dbname = db
self.con = None
self._create_con()
def __contains__(self, key):
sql = "select count(*) from pictures where path = ?"
result = self.con.execute(sql, [key]).fetchall()
return result[0][0] > 0
def __delitem__(self, key):
if key not in self:
raise KeyError(key)
sql = "delete from pictures where path = ?"
self.con.execute(sql, [key])
# Optimized
def __getitem__(self, key):
if isinstance(key, int):
sql = (
"select blocks, blocks2, blocks3, blocks4, blocks5, blocks6, blocks7, blocks8 "
"from pictures "
"where rowid = ?"
)
else:
sql = (
"select blocks, blocks2, blocks3, blocks4, blocks5, blocks6, blocks7, blocks8 "
"from pictures "
"where path = ?"
)
blocks = self.con.execute(sql, [key]).fetchone()
if blocks:
result = [bytes_to_colors(block) for block in blocks]
return result
else:
raise KeyError(key)
def __iter__(self):
sql = "select path from pictures"
result = self.con.execute(sql)
return (row[0] for row in result)
def __len__(self):
sql = "select count(*) from pictures"
result = self.con.execute(sql).fetchall()
return result[0][0]
def __setitem__(self, path_str, blocks):
blocks = [colors_to_bytes(block) for block in blocks]
if op.exists(path_str):
mtime = int(os.stat(path_str).st_mtime)
else:
mtime = 0
if path_str in self:
sql = (
"update pictures set blocks = ?, blocks2 = ?, blocks3 = ?, blocks4 = ?, blocks5 = ?, blocks6 = ?, "
"blocks7 = ?, blocks8 = ?, mtime_ns = ?"
"where path = ?"
)
else:
sql = (
"insert into pictures(blocks,blocks2,blocks3,blocks4,blocks5,blocks6,blocks7,blocks8,mtime_ns,path) "
"values(?,?,?,?,?,?,?,?,?,?)"
)
try:
self.con.execute(sql, blocks + [mtime, path_str])
except sqlite.OperationalError:
logging.warning("Picture cache could not set value for key %r", path_str)
except sqlite.DatabaseError as e:
logging.warning("DatabaseError while setting value for key %r: %s", path_str, str(e))
def _create_con(self, second_try=False):
try:
self.con = sqlite.connect(self.dbname, isolation_level=None)
self._check_upgrade()
except sqlite.DatabaseError as e: # corrupted db
if second_try:
raise # Something really strange is happening
logging.warning("Could not create picture cache because of an error: %s", str(e))
self.con.close()
os.remove(self.dbname)
self._create_con(second_try=True)
def _check_upgrade(self) -> None:
with self.con as conn:
has_schema = conn.execute(
"SELECT NAME FROM sqlite_master WHERE type='table' AND name='schema_version'"
).fetchall()
version = None
if has_schema:
version = conn.execute("SELECT version FROM schema_version ORDER BY version DESC").fetchone()[0]
else:
conn.execute("CREATE TABLE schema_version (version int PRIMARY KEY, description TEXT)")
if version != self.schema_version:
conn.execute(self.drop_table_query)
conn.execute(
"INSERT OR REPLACE INTO schema_version VALUES (:version, :description)",
{"version": self.schema_version, "description": self.schema_version_description},
)
conn.execute(self.create_table_query)
conn.execute(self.create_index_query)
def clear(self):
self.close()
if self.dbname != ":memory:":
os.remove(self.dbname)
self._create_con()
def close(self):
if self.con is not None:
self.con.close()
self.con = None
def filter(self, func):
to_delete = [key for key in self if not func(key)]
for key in to_delete:
del self[key]
def get_id(self, path):
sql = "select rowid from pictures where path = ?"
result = self.con.execute(sql, [path]).fetchone()
if result:
return result[0]
else:
raise ValueError(path)
def get_multiple(self, rowids):
ids = ",".join(map(str, rowids))
sql = (
"select rowid, blocks, blocks2, blocks3, blocks4, blocks5, blocks6, blocks7, blocks8 "
f"from pictures where rowid in ({ids})"
)
cur = self.con.execute(sql)
return (
(
rowid,
[
bytes_to_colors(blocks),
bytes_to_colors(blocks2),
bytes_to_colors(blocks3),
bytes_to_colors(blocks4),
bytes_to_colors(blocks5),
bytes_to_colors(blocks6),
bytes_to_colors(blocks7),
bytes_to_colors(blocks8),
],
)
for rowid, blocks, blocks2, blocks3, blocks4, blocks5, blocks6, blocks7, blocks8 in cur
)
def purge_outdated(self):
"""Go through the cache and purge outdated records.
A record is outdated if the picture doesn't exist or if its mtime is greater than the one in
the db.
"""
todelete = []
sql = "select rowid, path, mtime_ns from pictures"
cur = self.con.execute(sql)
for rowid, path_str, mtime_ns in cur:
if mtime_ns and op.exists(path_str):
picture_mtime = os.stat(path_str).st_mtime
if int(picture_mtime) <= mtime_ns:
# not outdated
continue
todelete.append(rowid)
if todelete:
sql = "delete from pictures where rowid in (%s)" % ",".join(map(str, todelete))
self.con.execute(sql)