Add shared storage option

libretranslate/storage.py

@@ -0,0 +1,147 @@
+import redis
+
+storage = None
+def get_storage():
+    return storage
+
+class Storage:
+    def set_bool(self, key, value):
+        raise Exception("not implemented")
+    def get_bool(self, key):
+        raise Exception("not implemented")
+
+    def set_int(self, key, value):
+        raise Exception("not implemented")
+    def get_int(self, key):
+        raise Exception("not implemented")
+
+    def set_str(self, key, value):
+        raise Exception("not implemented")
+    def get_str(self, key):
+        raise Exception("not implemented")
+
+    def set_hash_value(self, ns, key, value):
+        raise Exception("not implemented")
+    def get_hash_int(self, ns, key):
+        raise Exception("not implemented")
+    def inc_hash_int(self, ns, key):
+        raise Exception("not implemented")
+    def dec_hash_int(self, ns, key):
+        raise Exception("not implemented")
+    
+    def get_hash_keys(self, ns):
+        raise Exception("not implemented")
+    def del_hash(self, ns, key):
+        raise Exception("not implemented")
+    
+class MemoryStorage(Storage):
+    def __init__(self):
+        self.store = {}
+
+    def set_bool(self, key, value):
+        self.store[key] = bool(value)
+    
+    def get_bool(self, key):
+        return bool(self.store[key])
+
+    def set_int(self, key, value):
+        self.store[key] = int(value)
+
+    def get_int(self, key):
+        return int(self.store.get(key, 0))
+    
+    def set_str(self, key, value):
+        self.store[key] = value
+    
+    def get_str(self, key):
+        return str(self.store.get(key, ""))
+
+    def get_hash_int(self, ns, key):
+        d = self.store.get(ns, {})
+        return int(d.get(key, 0))
+
+    def inc_hash_int(self, ns, key):
+        if ns not in self.store:
+            self.store[ns] = {}
+        
+        if key not in self.store[ns]:
+            self.store[ns][key] = 0
+        else:
+            self.store[ns][key] += 1
+
+    def dec_hash_int(self, ns, key):
+        if ns not in self.store:
+            self.store[ns] = {}
+        
+        if key not in self.store[ns]:
+            self.store[ns][key] = 0
+        else:
+            self.store[ns][key] -= 1
+
+    def get_all_hash_int(self, ns):
+        return [{str(k): int(v)} for k,v in self.store[ns].items()]
+
+    def del_hash(self, ns, key):
+        del self.store[ns][key]
+
+
+class RedisStorage(Storage):
+    def __init__(self, redis_uri):
+        self.conn = redis.from_url(redis_uri)
+        self.conn.ping()
+
+    def set_bool(self, key, value):
+        self.conn.set(key, "1" if value else "0")
+    
+    def get_bool(self, key):
+        return bool(self.conn.get(key))
+
+    def set_int(self, key, value):
+        self.conn.set(key, str(value))
+
+    def get_int(self, key):
+        v = self.conn.get(key)
+        if v is None:
+            return 0
+        else:
+            return v
+    
+    def set_str(self, key, value):
+        self.conn.set(key, value)
+    
+    def get_str(self, key):
+        v = self.conn.get(key)
+        if v is None:
+            return ""
+        else:
+            return v.decode('utf-8')
+
+    def get_hash_int(self, ns, key):
+        v = self.conn.hget(ns, key)
+        if v is None:
+            return 0
+        else:
+            return int(v)
+
+    def inc_hash_int(self, ns, key):
+        return int(self.conn.hincrby(ns, key))
+
+    def dec_hash_int(self, ns, key):
+        return int(self.conn.hincrby(ns, key, -1))
+
+    def get_all_hash_int(self, ns):
+        return [{k.decode("utf-8"): int(v)} for k,v in self.conn.hgetall(ns).items()]
+
+    def del_hash(self, ns, key):
+        conn.hdel(ns, key)
+
+def setup(storage_uri):
+    global storage
+    if storage_uri.startswith("memory://"):
+        storage = MemoryStorage()
+    elif storage_uri.startswith("redis://"):
+        storage = RedisStorage(storage_uri)
+    else:
+        raise Exception("Invalid storage URI: " + storage_uri)
+
+    return storage