1 files changed, 98 insertions, 0 deletions
diff --git a/gallery_dl/archive.py b/gallery_dl/archive.py
new file mode 100644
index 0000000..5f05bbf
--- /dev/null
+++ b/gallery_dl/archive.py
@@ -0,0 +1,98 @@
+# -*- coding: utf-8 -*-
+
+# Copyright 2024 Mike Fährmann
+#
+# This program is free software; you can redistribute it and/or modify
+# it under the terms of the GNU General Public License version 2 as
+# published by the Free Software Foundation.
+
+"""Download Archives"""
+
+import os
+import sqlite3
+from . import formatter
+
+
+class DownloadArchive():
+
+    def __init__(self, path, format_string, pragma=None,
+                 cache_key="_archive_key"):
+        try:
+            con = sqlite3.connect(path, timeout=60, check_same_thread=False)
+        except sqlite3.OperationalError:
+            os.makedirs(os.path.dirname(path))
+            con = sqlite3.connect(path, timeout=60, check_same_thread=False)
+        con.isolation_level = None
+
+        self.keygen = formatter.parse(format_string).format_map
+        self.connection = con
+        self.close = con.close
+        self.cursor = cursor = con.cursor()
+        self._cache_key = cache_key
+
+        if pragma:
+            for stmt in pragma:
+                cursor.execute("PRAGMA " + stmt)
+
+        try:
+            cursor.execute("CREATE TABLE IF NOT EXISTS archive "
+                           "(entry TEXT PRIMARY KEY) WITHOUT ROWID")
+        except sqlite3.OperationalError:
+            # fallback for missing WITHOUT ROWID support (#553)
+            cursor.execute("CREATE TABLE IF NOT EXISTS archive "
+                           "(entry TEXT PRIMARY KEY)")
+
+    def add(self, kwdict):
+        """Add item described by 'kwdict' to archive"""
+        key = kwdict.get(self._cache_key) or self.keygen(kwdict)
+        self.cursor.execute(
+            "INSERT OR IGNORE INTO archive (entry) VALUES (?)", (key,))
+
+    def check(self, kwdict):
+        """Return True if the item described by 'kwdict' exists in archive"""
+        key = kwdict[self._cache_key] = self.keygen(kwdict)
+        self.cursor.execute(
+            "SELECT 1 FROM archive WHERE entry=? LIMIT 1", (key,))
+        return self.cursor.fetchone()
+
+    def finalize(self):
+        pass
+
+
+class DownloadArchiveMemory(DownloadArchive):
+
+    def __init__(self, path, format_string, pragma=None,
+                 cache_key="_archive_key"):
+        DownloadArchive.__init__(self, path, format_string, pragma, cache_key)
+        self.keys = set()
+
+    def add(self, kwdict):
+        self.keys.add(
+            kwdict.get(self._cache_key) or
+            self.keygen(kwdict))
+
+    def check(self, kwdict):
+        key = kwdict[self._cache_key] = self.keygen(kwdict)
+        if key in self.keys:
+            return True
+        self.cursor.execute(
+            "SELECT 1 FROM archive WHERE entry=? LIMIT 1", (key,))
+        return self.cursor.fetchone()
+
+    def finalize(self):
+        if not self.keys:
+            return
+
+        cursor = self.cursor
+        with self.connection:
+            try:
+                cursor.execute("BEGIN")
+            except sqlite3.OperationalError:
+                pass
+
+            stmt = "INSERT OR IGNORE INTO archive (entry) VALUES (?)"
+            if len(self.keys) < 100:
+                for key in self.keys:
+                    cursor.execute(stmt, (key,))
+            else:
+                cursor.executemany(stmt, ((key,) for key in self.keys))