Evict old cached usernames in case of collision

This commit is contained in:
Lonami Exo 2020-10-01 13:20:29 +02:00
parent 233daafd96
commit 67b87a0ea0

View File

@ -1,5 +1,6 @@
import datetime
import os
import time
from telethon.tl import types
from .memory import MemorySession, _SentFileType
@ -17,7 +18,7 @@ except ImportError as e:
sqlite3_err = type(e)
EXTENSION = '.session'
CURRENT_VERSION = 6 # database version
CURRENT_VERSION = 7 # database version
class SQLiteSession(MemorySession):
@ -84,7 +85,8 @@ class SQLiteSession(MemorySession):
hash integer not null,
username text,
phone integer,
name text
name text,
date integer
)"""
,
"""sent_files (
@ -148,6 +150,9 @@ class SQLiteSession(MemorySession):
# hashes for User and Channel are wrong, so drop them off.
old += 1
c.execute('delete from entities')
if old == 6:
old += 1
c.execute("alter table entities add column date integer")
c.close()
@ -264,10 +269,9 @@ class SQLiteSession(MemorySession):
# Entity processing
def process_entities(self, tlo):
"""Processes all the found entities on the given TLObject,
unless .enabled is False.
Returns True if new input entities were added.
"""
Processes all the found entities on the given TLObject,
unless .save_entities is False.
"""
if not self.save_entities:
return
@ -278,8 +282,10 @@ class SQLiteSession(MemorySession):
c = self._cursor()
try:
now_tup = (int(time.time()),)
rows = [row + now_tup for row in rows]
c.executemany(
'insert or replace into entities values (?,?,?,?,?)', rows)
'insert or replace into entities values (?,?,?,?,?,?)', rows)
finally:
c.close()
@ -288,8 +294,25 @@ class SQLiteSession(MemorySession):
'select id, hash from entities where phone = ?', phone)
def get_entity_rows_by_username(self, username):
return self._execute(
'select id, hash from entities where username = ?', username)
c = self._cursor()
try:
results = c.execute(
'select id, hash, date from entities where username = ?',
(username,)
).fetchall()
if not results:
return None
# If there is more than one result for the same username, evict the oldest one
if len(results) > 1:
results.sort(key=lambda t: t[2] or 0)
c.executemany('update entities set username = null where id = ?',
[(t[0],) for t in results[:-1]])
return results[-1][0], results[-1][1]
finally:
c.close()
def get_entity_rows_by_name(self, name):
return self._execute(