def RemoveSnapshotFromArticle(self, request, context):
article_id = request.article_id
snapshot_id = request.snapshot_id
with self.db_fn() as db:
db.execute(
"DELETE FROM snapshots WHERE snapshots.uuid = ? AND snapshots.article = ?",
(snapshot_id, article_id),
)
return co.Empty()
def TryLogin(self, request, context):
token = request.token
try:
if self.auth.request_login(token):
return co.Empty()
else:
context.abort(StatusCode.DEADLINE_EXCEEDED, "Confirmation timeout")
except AuthException:
context.abort(StatusCode.INVALID_ARGUMENT, "Token is in use")
def DeleteArticle(self, request, context):
openid = context.openid
article_id = request.article_id
with self.db_fn() as db:
db.execute(
"DELETE FROM articles WHERE articles.id = ? AND articles.user = ?",
(article_id, openid),
)
return co.Empty()
def ChangeArticleTitle(self, request, context):
openid = context.openid
article_id = request.article_id
title = request.title
with self.db_fn() as db:
db.execute(
"UPDATE articles SET title = ? WHERE articles.id = ? AND articles.user = ?",
(title, article_id, openid),
)
return co.Empty()
def Capture(self, request, context):
openid = context.openid
urls = list(set(request.urls))
article_id = request.article_id
worker = self.worker_stub()
tasks = {} # url -> task_id
with self.db_fn() as db:
for url in urls:
task_id = uuid()
tasks[url] = task_id
db.execute(
"INSERT INTO tasks VALUES (?,?,?,?)",
(task_id, openid, url, article_id),
)
def _async_action():
successful_urls = set()
for res in worker.Crawl(wo.CrawlRequest(urls=urls)):
url = res.url
task_id = tasks[url]
content = res.content
logging.info(f'Capture succeeded for {url} of type {content.type}')
timestamp = unix_time()
_id = uuid()
ledger_key = ledger.add(content.hash)
with self.db_fn() as db:
# If the snapshot already exists(because this snapshot has multiple pieces of data attached), ignore.
db.execute(
"INSERT OR IGNORE INTO snapshots VALUES (?,?,?,?,?,?)",
(_id, article_id, url, timestamp, False, None),
)
db.execute(
"INSERT INTO data VALUES (?,?,?,?,?)",
(_id, content.type, content.data.decode(), content.hash, ledger_key),
)
db.execute('DELETE FROM tasks WHERE id = ?', (task_id,))
successful_urls.add(url)
self.task_event.notify()
# Worker hang up
for url, task_id in tasks.items():
if url not in successful_urls:
logging.info(f'Capture failed for {url}')
self._add_notification(openid, f"拍摄快照失败：{url}", is_error=True)
with self.db_fn() as db:
db.execute('DELETE FROM tasks WHERE id = ?', (task_id,))
self.task_event.notify()
threading.Thread(target=_async_action).start()
self.task_event.notify()
return co.Empty()
def ConfirmLogin(self, request, context):
token = request.token
openid = request.openid
username = request.name
try:
self.auth.confirm_login(token, openid)
except AuthException:
context.abort(StatusCode.INVALID_ARGUMENT, "Non-existing token")
with self.db_fn() as db:
cnt = db.execute(
"SELECT COUNT(*) FROM users WHERE users.openid = ?", (openid,)
).fetchone()[0]
if cnt == 0:
logging.info(f"Creating new user: username={username}, openid={openid}")
db.execute("INSERT INTO users VALUES (?,?)", (openid, username))
return co.Empty()
def inner():
for tasks in servicer.GetActiveTasks(co.Empty(), ctx):
if canceled:
return
results.append(tasks.tasks)
def test_login(servicer, ctx, login):
r = servicer.GetUserData(co.Empty(), ctx)
assert r.username == username
def RegisterWorker(self, request, context):
addr = request.addr
port = request.port
self.workers.add(f"{addr}:{port}")
return co.Empty()
def Report(self, req, ctx):
with self.db_fn() as db:
db.execute('UPDATE snapshots SET reported = 1, report_reason = ? WHERE uuid = ?', (req.reason, req.snapshot_id))
return co.Empty()
def MarkAllAsRead(self, req, ctx):
with self.db_fn() as db:
db.execute('UPDATE notifications SET has_read = 1 WHERE user = ?', (ctx.openid,))
return co.Empty()
def ClearTasks(self, req, ctx):
with self.db_fn() as db:
db.execute('DELETE FROM tasks WHERE user = ?', (ctx.openid,))
return co.Empty()
