alex/nova-api
1
0
Fork 0
mirror of https://github.com/NovaOSS/nova-api.git synced 2024-11-25 20:43:56 +01:00
Code Issues Packages Projects Releases Wiki Activity

Add model support

Browse source
This commit is contained in:
Game_Time 2023-06-24 22:26:36 +05:00 committed by GitHub
parent 6ca84147f2
commit 280e5d9a58
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
1 changed files with 18 additions and 15 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

33
api/keys.py
View file

@ -12,23 +12,25 @@ class Keys:
# --- END OF CONFIG --- # --- END OF CONFIG ---
locked_keys = set() locked_keys = set()
cache = set() cache = {}
# Initialize MongoDB # Initialize MongoDB
client = MongoClient(MONGO_URI) client = MongoClient(MONGO_URI)
db = client.get_database('keys_db') db = client.get_database('keys_db')
collection = db['keys'] collection = db['keys']
def __init__(self, key: str): def __init__(self, key: str, model: str):
self.key = key self.key = key
self.model = model
if not Keys.cache: if not Keys.cache:
self._load_keys() self._load_keys()
def _load_keys(self) -> None: def _load_keys(self) -> None:
cursor = self.collection.find({}, {'_id': 0, 'key_value': 1}) cursor = Keys.collection.find({}, {'_id': 0, 'key_value': 1, 'model': 1})
for doc in cursor: for doc in cursor:
key_value = doc['key_value'] key_value = doc['key_value']
self.cache.add(key_value) model = doc['model']
Keys.cache.setdefault(model, set()).add(key_value)
def lock(self) -> None: def lock(self) -> None:
self.locked_keys.add(self.key) self.locked_keys.add(self.key)
@ -39,34 +41,35 @@ class Keys:
def is_locked(self) -> bool: def is_locked(self) -> bool:
return self.key in self.locked_keys return self.key in self.locked_keys
def get() -> str: @staticmethod
key_candidates = list(Keys.cache) def get(model: str) -> str:
key_candidates = list(Keys.cache.get(model, set()))
random.shuffle(key_candidates) random.shuffle(key_candidates)
for key_candidate in key_candidates: for key_candidate in key_candidates:
key = Keys(key_candidate) key = Keys(key_candidate, model)
if not key.is_locked(): if not key.is_locked():
key.lock() key.lock()
return key.key return key.key
print("[WARN] No unlocked keys found in get keys request!") print(f"[WARN] No unlocked keys found for model '{model}' in get keys request!")
def delete(self) -> None: def delete(self) -> None:
self.collection.delete_one({'key_value': self.key}) Keys.collection.delete_one({'key_value': self.key, 'model': self.model})
# Update cache # Update cache
try: try:
Keys.cache.remove(self.key) Keys.cache[self.model].remove(self.key)
except KeyError: except KeyError:
print("[WARN] Tried to remove a key from cache which was not present: " + self.key) print(f"[WARN] Tried to remove a key from cache which was not present: {self.key}")
def save(self) -> None: def save(self) -> None:
self.collection.insert_one({'key_value': self.key}) Keys.collection.insert_one({'key_value': self.key, 'model': self.model})
# Update cache # Update cache
Keys.cache.add(self.key) Keys.cache.setdefault(self.model, set()).add(self.key)
# Usage example: # Usage example:
# os.environ['MONGO_URI'] = "mongodb://localhost:27017" # os.environ['MONGO_URI'] = "mongodb://localhost:27017"
# key_instance = Keys("example_key") # key_instance = Keys("example_key", "example_model")
# key_instance.save() # key_instance.save()
# key_value = Keys.get() # key_value = Keys.get("example_model")