Merge pull request #41 from mcilya/next

Next

README.md

@@ -364,7 +364,10 @@ Enter this commands if prompt starts with `Search@DB>`
 | export_path     | String (Path)                               | Where to store exported images                                              |
 | time_wait       | Integer (seconds)                           | How long thread can stay alive                                              |
 | ids_file        | String (Path/Filename)                      | Name for tempfile (No need to change)                                       |
-| db_name         | String (Path/Filename)                      | Where to store DB file                                                      |
+| db_name         | String                                      | Defines name for database file In case of using mysql defines schema name   |
+| use_mysql       | Bool (True/False)                           | Defines if mysql connection should be used                                  |
+| mysql_user      | String                                      | Defines mysql user                                                          |
+| mysql_password  | String                                      | Defines mysql password                                                      |
 | thread_cap      | Integer (number)                            | Defines maximum running threads before blocking creating new threads        |
 | sleep_time      | Integer (seconds)                           | Defines time to wait before creating new thread after thread cap is reached |
 | enable_polling  | Bool                                        | Setting for enabling/disabling polling for changes in settings_file.py      |

dermod/db.py

@@ -1,4 +1,5 @@
 import os
+from os import remove
 import sqlite3
 import sys
 import codecs
@@ -7,57 +8,82 @@
 from dermod import input_parser as ip
 import settings_file
 
+try:
+    import mysql.connector
+except ImportError:
+    pass
+
 if settings_file.suppress_errors:
     logging.raiseExceptions = False
 
 def precomp():
-    global tag_col, tag_col_full, tag_col_serv
+    global tag_col, tag_col_full, tag_col_serv, rand_func
     tag_col = 'tags'
-    tag_col_full = 'fname, tags, height, width, ratio, source_link, prefix, id'
-    tag_col_serv = 'height, width, ratio, source_link, prefix, id'
-    init_db()
-    t = [x[1] for x in cursor.execute("select * from sqlite_master").fetchall()]
+
+    if settings_file.use_mysql:
+        tag_col_full = '`fname` VARCHAR(256), `tags` MEDIUMTEXT, `height` INT, `width` INT, `ratio` TEXT, `source_link` MEDIUMTEXT, `prefix` VARCHAR(32), `id` INT'
+        tag_col_serv = '`height` INT, `width` INT, `ratio` TEXT, `source_link` MEDIUMTEXT, `prefix` VARCHAR(32), `id` INT'
+        rand_func = "RAND()"
+    else:
+        tag_col_full = 'fname, tags, height, width, ratio, source_link, prefix, id INT'
+        tag_col_serv = 'height, width, ratio, source_link, prefix, id INT'
+        rand_func = "RANDOM()"
+
+    if settings_file.use_mysql:
+        try:
+            conn = mysql.connector.connect(user=settings_file.mysql_user, password=settings_file.mysql_password, db=settings_file.db_name)
+        except:
+            conn = mysql.connector.connect(user=settings_file.mysql_user, password=settings_file.mysql_password)
+            cur = conn.cursor()
+            cur.execute("create schema `{}`".format(settings_file.db_name))
+            conn.commit()
+            del cur
+        conn.close()
+        del conn
+    conn, cursor = init_db()
+    if settings_file.use_mysql:
+        cursor.execute("SHOW TABLES")
+        t = [x[0] for x in cursor.fetchall()]
+    else:
+        t = [x[1] for x in cursor.execute("select * from sqlite_master").fetchall()]
     if "images" not in t:
         mkdb("images")
 
 
 def init_db():
-    global cursor
-    global conn
-    conn = sqlite3.connect(settings_file.db_name)
+    if settings_file.use_mysql:
+        conn = mysql.connector.connect(user=settings_file.mysql_user, password=settings_file.mysql_password, db=settings_file.db_name)
+    else:
+        conn = sqlite3.connect(settings_file.db_name)
     cursor = conn.cursor()
+    return conn, cursor
 
 
 def total_found():
-    init_db()
-    print("Images in DataBase =>", cursor.execute("SELECT count(*) FROM images where fname").fetchone()[0])
+    conn, cursor = init_db()
+    cursor.execute("SELECT count(*) FROM images where fname")
+    print("Images in DataBase =>", cursor.fetchone()[0])
 
 
 def get_all_entries():
-    init_db()
-    result = list(cursor.execute(
-        "SELECT tags from images"))
+    conn, cursor = init_db()
+    cursor.execute("SELECT tags from images")
+    result = list(cursor.fetchall())
     return result
 
 
 def mkdb(table_name):
-    init_db()
+    conn, cursor = init_db()
     cursor.execute('drop table IF EXISTS {}'.format(table_name))
-    sample = """CREATE TABLE {}({} INT)""".format(table_name, tag_col_full)
+    sample = """CREATE TABLE {}({}, PRIMARY KEY (prefix, id))""".format(table_name, tag_col_full)
     cursor.execute(sample)
     conn.commit()
-
-def mk_tdb(table_name):
-    init_db()
-    cursor.execute('drop table IF EXISTS {}'.format(table_name))
-    sample = """CREATE TEMP TABLE {}({} INT)""".format(table_name, tag_col_full)
-    cursor.execute(sample)
+    cursor.execute(f"CREATE UNIQUE INDEX `prefix_id_idx` on {table_name}(id, prefix)")
     conn.commit()
 
-
 def fill_db(file=settings_file.ids_file):
     print("\nFilling DB")
-    init_db()
+    conn, cursor = init_db()
     unparsed = open(file).read()
     halfparsed = unparsed.strip("\n").split("\n")
     cnt = 0
@@ -68,30 +94,36 @@ def fill_db(file=settings_file.ids_file):
         k = i[6]
         k = str(k).strip("[]").replace('" ', '"').replace(
             ' "', '"').replace('\' ', '\'').replace(' \'', '\'')
-        j = "INSERT OR REPLACE INTO images VALUES ('{}.{}', '{}', '{}', '{}', '{}', '{}', '{}', {})".format(
+        j = "REPLACE INTO images VALUES ('{}.{}', '{}', '{}', '{}', '{}', '{}', '{}', {})".format(
             i[0], i[1], k, i[3], i[4], i[5], i[2], i[7], i[0]).replace('\\', '')
-        cursor.execute(j)
+        try:
+            cursor.execute(j)
+        except:
+            j = "INSERT OR REPLACE INTO images VALUES ('{}.{}', '{}', '{}', '{}', '{}', '{}', '{}', {})".format(
+                i[0], i[1], k, i[3], i[4], i[5], i[2], i[7], i[0]).replace('\\', '')
+            cursor.execute(j)
         if cnt == 10:
             conn.commit()
             cnt = 0
     conn.commit()
-    cursor.execute("delete from images where rowid not in (select min(rowid) from images group by fname)")
-    conn.commit()
-    conn.execute("VACUUM")
-    conn.commit()
+    if not settings_file.use_mysql:
+        cursor.execute("delete from images where rowid not in (select min(rowid) from images group by fname)")
+        conn.commit()
+        conn.execute("VACUUM")
+        conn.commit()
 
 
 def count_tag(tag_to_count):
-    init_db()
+    conn, cursor = init_db()
 
     sample = """select count(*) FROM images where tags like '%,,{},,%'""".format(tag_to_count)
     output = list(cursor.execute(sample))
     print(str(output[0]).strip("()").replace(",", "") +
           " images tagged {}".format(tag_to_count))
 
 
-def search(list_search, list_remove, page=0):
-    init_db()
+def search(list_search, list_remove, page=0, return_query=False):
+    conn, cursor = init_db()
 
     # Handle special tags
     specials = []
@@ -115,8 +147,6 @@ def search(list_search, list_remove, page=0):
         specials = " where " + (" and ".join(specials))
     else:
         specials = ""
-
-    mk_tdb('temp1')
 
     if len(list_search) != 0:
         autogen_template = "tags like '%,,{},,%'"
@@ -134,80 +164,98 @@ def search(list_search, list_remove, page=0):
             else:
                 queries.append(autogen_template.format(i))
 
-        autogen_query = "SELECT * from images where {}".format(" and ".join(queries))
-        query = "INSERT INTO temp1 SELECT DISTINCT * from ({})".format(autogen_query)
-        cursor.execute(query)
+        autogen_query = "SELECT DISTINCT * from images where {}".format(" and ".join(queries))
+        #query = "INSERT INTO temp1 SELECT DISTINCT * from ({})".format(autogen_query)
+        #cursor.execute(query)
     else:
-        sample = "INSERT INTO temp1 SELECT DISTINCT * FROM images"
-        cursor.execute(sample)
+        autogen_query = "SELECT DISTINCT * FROM images"
+        # cursor.execute(sample)
 
 
     if len(list_remove) == 0:
         pass
     else:
+        remove_queries = []
         for i in list_remove:
-            cursor.execute(
-                "DELETE FROM temp1 WHERE {} like '%,,{},,%'".format(tag_col, i.replace("*", "%")))
-            conn.commit()
-
-    final_autogen = "SELECT * from temp1 {specials} order by id DESC limit {imgs_amount} offset {offset}".format(
-        imgs_amount=settings_file.showing_imgs, offset=settings_file.showing_imgs*page, specials=specials)
-    results = list(cursor.execute(final_autogen))
+            remove_queries.append("{} not like '%,,{},,%'".format(tag_col, i.replace("*", "%")))
+        remove_autogen = " and ".join(remove_queries)
+
+        if 'where' in autogen_query:
+            autogen_query = autogen_query + " and " + remove_autogen
+        else:
+            autogen_query = autogen_query + " where " + remove_autogen
+
+    if return_query:
+        final_autogen = "{autogen_query} {specials}".format(
+            autogen_query=autogen_query, specials=specials)
+        return final_autogen
+    else:
+        final_autogen = "{autogen_query} {specials} order by id DESC limit {imgs_amount} offset {offset}".format(
+            autogen_query=autogen_query, imgs_amount=settings_file.showing_imgs, offset=settings_file.showing_imgs*page, specials=specials)
+    cursor.execute(final_autogen)
+    results = list(cursor.fetchall())
 
-    total = cursor.execute("SELECT COUNT(*) FROM temp1 {}".format(specials)).fetchone()
+    cursor.execute("SELECT COUNT(*) FROM ({autogen_query}) as T {}".format(specials, autogen_query=autogen_query))
+    total = cursor.fetchone()
     conn.commit()
 
     return results, total
+
 
 
 def search_by_id(img_id, prefix="%"):
-    init_db()
+    conn, cursor = init_db()
     sql = "SELECT * FROM images WHERE id = {} and prefix like '{}_'".format(img_id, prefix)
-    result = list(cursor.execute(sql))
+    cursor.execute(sql)
+    result = list(cursor.fetchall())
     if len(result) != 0:
         return result[0]
     else:
         return []
 
 
 def random_img():
-    init_db()
-    result = list(cursor.execute(
-        "SELECT * FROM images ORDER BY RANDOM() LIMIT 1"))
+    conn, cursor = init_db()
+    cursor.execute("SELECT * FROM images ORDER BY {rand} LIMIT 1".format(rand=rand_func))
+    result = list(cursor.fetchall())
     return result
 
 
 def tagged_random(tag):
-    init_db()
-    search(tag['search'], tag['remove'])
-    result = cursor.execute("select * from temp1 order by random() Limit 1").fetchone()
-
+    conn, cursor = init_db()
+    query = search(tag['search'], tag['remove'], return_query=True)
+    cursor.execute("select * from ({}) as T order by {rand} Limit 1".format(query, rand=rand_func))
+    result = cursor.fetchone()
     return result
 
 def get_prev(id):
-    init_db()
-    result = list(cursor.execute("select * from images where (id<{}) order by id desc limit 1".format(int(id))).fetchall())[0]
+    conn, cursor = init_db()
+    cursor.execute("select * from images where (id<{}) order by id desc limit 1".format(int(id)))
+    result = list(cursor.fetchall())[0]
     return result
 
 def get_next(id):
-    init_db()
-    result = list(cursor.execute("select * from images where (id>{}) order by id asc limit 1".format(int(id))).fetchall())[0]
+    conn, cursor = init_db()
+    cursor.execute("select * from images where (id>{}) order by id asc limit 1".format(int(id)))
+    result = list(cursor.fetchall())[0]
     return result
 
 def tagged_get_prev(id, tag):
-    init_db()
-    search(tag['search'], tag['remove'])
-    result = list(cursor.execute("select * from temp1 where (id<{}) order by id desc limit 1".format(int(id))).fetchall())[0]
+    conn, cursor = init_db()
+    query = search(tag['search'], tag['remove'], return_query=True)
+    cursor.execute("select * from ({}) as T where (id<{}) order by id desc limit 1".format(query, int(id)))
+    result = list(cursor.fetchall())[0]
     return result
 
 def tagged_get_next(id, tag):
-    init_db()
-    search(tag['search'], tag['remove'])
-    result = list(cursor.execute("select * from temp1 where (id>{}) order by id asc limit 1".format(int(id))).fetchall())[0]
+    conn, cursor = init_db()
+    query = search(tag['search'], tag['remove'], return_query=True)
+    cursor.execute("select * from ({}) as T where (id>{}) order by id asc limit 1".format(query, int(id)))
+    result = list(cursor.fetchall())[0]
     return result
 
 def remove_entry(imgid, prefix):
-    init_db()
+    conn, cursor = init_db()
     sql = f"delete from images where (id = {int(imgid)}) and (prefix = '{prefix}')"
     print(sql)
     cursor.execute(sql)

dermod/imgloader.py

@@ -38,12 +38,12 @@ def run(self):
             self.get_raw_image()
         except:
             self.readiness = 1
-            quit(0)
+            return
         print(f"Wiritng {self.id}") if self.format == '' else ''
         self.writer()
         self.readiness = 1
         del self.raw_data
-        quit(0)
+        return
 
     def get_raw_image(self):
         with requests.Session() as s:
@@ -60,7 +60,7 @@ def get_raw_image(self):
             if self.tmp.status_code >= 400:
                 global is_error_code
                 is_error_code = True
-                quit(1)
+                return
             else:
                 self.raw_data = self.tmp.content
 

dermod/listloader.py

@@ -32,7 +32,7 @@ def get_data(self):
         global is_error_code
         with requests.Session() as s:
             if is_error_code is True:
-                quit(1)
+                return
             s.headers = {
                 'User-Agent': 'DBooru/2.0 (Api checker module)(github.com/mcilya/DBooru)'}
             if settings_file.enable_proxy is False:
@@ -97,7 +97,7 @@ def run(self):
             self.get_data()
         except:
             self.readiness = 1
-            quit()
+            return
         else:
             if re.match("{}".format(self.module.empty_page), self.raw_data) is not None:
                 global empties
@@ -112,7 +112,7 @@ def run(self):
         #     if is_error_code == False:
         #         self.run()
             self.readiness = 1
-            quit()
+            return
 
 
 def run(module, pages_num=0, file=settings_file.ids_file, endwith="\r"):

requirements.txt

@@ -9,3 +9,4 @@ markupsafe
 click
 itsdangerous
 termcolor
+mysql-connector-python