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()