1- import json
21from collections import defaultdict
2+ from typing import Any
33
4- import redis
4+ from sqlalchemy import create_engine , text
5+ from sqlalchemy .engine import Engine
6+ from sqlalchemy .orm import Session , sessionmaker
57
6- from mgallery .utils .settings import REDIS_URL
8+ from mgallery .utils .settings import DATABASE_URL
79
810
911class Database :
10- def __init__ (self ):
11- self .client = redis .from_url (REDIS_URL )
12+ def __init__ (self , engine : Engine | None = None ):
13+ self ._engine = engine or create_engine (
14+ DATABASE_URL ,
15+ pool_size = 5 ,
16+ max_overflow = 5 ,
17+ pool_pre_ping = True ,
18+ )
19+ self ._session_factory = sessionmaker (bind = self ._engine )
20+
21+ def _session (self ) -> Session :
22+ return self ._session_factory ()
1223
13- def get (self , key : str | bytes ) -> dict :
14- if data := self .client .get (key ):
15- return json .loads (data )
24+ def get (self , key : str | bytes ) -> dict [str , Any ]:
25+ key_str = key .decode () if isinstance (key , bytes ) else key
26+ parts = key_str .split ("-" , 1 )
27+ if len (parts ) != 2 :
28+ return {}
29+ phash , rest = parts
30+ if "/" not in rest :
31+ return {}
32+ path , name = rest .rsplit ("/" , 1 )
33+ if phash == "None" :
34+ with self ._session () as session :
35+ row = session .execute (
36+ text ("SELECT path, name, phash, width, height, size FROM images WHERE phash IS NULL AND path = :path AND name = :name" ),
37+ {"path" : path , "name" : name },
38+ ).mappings ().first ()
39+ if row :
40+ return dict (row )
41+ else :
42+ with self ._session () as session :
43+ row = session .execute (
44+ text ("SELECT path, name, phash, width, height, size FROM images WHERE phash = :phash AND path = :path AND name = :name" ),
45+ {"phash" : phash , "path" : path , "name" : name },
46+ ).mappings ().first ()
47+ if row :
48+ return dict (row )
1649 return {}
1750
18- def all (self , pattern : str = "*" ) -> list :
19- return [self .get (key ) for key in self .client .keys (pattern )]
51+ def all (self , pattern : str = "*" ) -> list [dict [str , Any ]]:
52+ if pattern == "*" :
53+ with self ._session () as session :
54+ rows = session .execute (text ("SELECT path, name, phash, width, height, size FROM images" )).mappings ().all ()
55+ else :
56+ with self ._session () as session :
57+ rows = session .execute (
58+ text ("SELECT path, name, phash, width, height, size FROM images WHERE path LIKE :pattern" ),
59+ {"pattern" : f"%{ pattern } },
60+ ).mappings ().all ()
61+ return [dict (row ) for row in rows ]
2062
21- def duplicates (self ) -> dict [str , list ]:
22- duplicates = defaultdict (list )
23- for item in self .all ():
63+ def duplicates (self ) -> dict [str , list [dict [str , Any ]]]:
64+ with self ._session () as session :
65+ rows = session .execute (
66+ text ("""
67+ SELECT path, name, phash, width, height, size
68+ FROM images
69+ WHERE phash IN (
70+ SELECT phash FROM images WHERE phash IS NOT NULL
71+ GROUP BY phash
72+ HAVING COUNT(*) > 1
73+ )
74+ """ ),
75+ ).mappings ().all ()
76+ result : dict [str , list [dict [str , Any ]]] = defaultdict (list )
77+ for row in rows :
78+ item = dict (row )
2479 if item ["phash" ]:
25- duplicates [item ["phash" ]].append (item )
26- return { k : v for k , v in duplicates . items () if len ( v ) > 1 }
80+ result [item ["phash" ]].append (item )
81+ return result
2782
2883 def create (
2984 self ,
@@ -33,20 +88,63 @@ def create(
3388 width : int | None = None ,
3489 height : int | None = None ,
3590 size : int | None = None ,
91+ session : Session | None = None ,
3692 ):
37- key = f"{ phash } { path } { name }
38- value = json .dumps (
39- {
40- "path" : path ,
41- "name" : name ,
42- "phash" : phash ,
43- "width" : width ,
44- "height" : height ,
45- "size" : size ,
46- }
47- )
48- self .client .set (key , value )
93+ external_session = session is not None
94+ session = session or self ._session ()
95+ try :
96+ session .execute (
97+ text ("""
98+ INSERT INTO images (path, name, phash, width, height, size)
99+ VALUES (:path, :name, :phash, :width, :height, :size)
100+ ON CONFLICT (path, name) DO UPDATE SET
101+ phash = COALESCE(EXCLUDED.phash, images.phash),
102+ width = COALESCE(EXCLUDED.width, images.width),
103+ height = COALESCE(EXCLUDED.height, images.height),
104+ size = COALESCE(EXCLUDED.size, images.size)
105+ """ ),
106+ {"path" : path , "name" : name , "phash" : phash , "width" : width , "height" : height , "size" : size },
107+ )
108+ if not external_session :
109+ session .commit ()
110+ finally :
111+ if not external_session :
112+ session .close ()
113+
114+ def create_batch (self , items : list [dict [str , Any ]], session : Session | None = None ) -> None :
115+ external_session = session is not None
116+ session = session or self ._session ()
117+ try :
118+ for item in items :
119+ session .execute (
120+ text ("""
121+ INSERT INTO images (path, name, phash, width, height, size)
122+ VALUES (:path, :name, :phash, :width, :height, :size)
123+ ON CONFLICT (path, name) DO UPDATE SET
124+ phash = COALESCE(EXCLUDED.phash, images.phash),
125+ width = COALESCE(EXCLUDED.width, images.width),
126+ height = COALESCE(EXCLUDED.height, images.height),
127+ size = COALESCE(EXCLUDED.size, images.size)
128+ """ ),
129+ {
130+ "path" : item ["path" ],
131+ "name" : item ["name" ],
132+ "phash" : item .get ("phash" ),
133+ "width" : item .get ("width" ),
134+ "height" : item .get ("height" ),
135+ "size" : item .get ("size" ),
136+ },
137+ )
138+ if not external_session :
139+ session .commit ()
140+ finally :
141+ if not external_session :
142+ session .close ()
49143
50144 def delete (self , path : str , name : str ):
51- for key in self .client .keys (f"*{ path } { name } ):
52- self .client .delete (key )
145+ with self ._session () as session :
146+ session .execute (
147+ text ("DELETE FROM images WHERE path = :path AND name = :name" ),
148+ {"path" : path , "name" : name },
149+ )
150+ session .commit ()
0 commit comments