{
 "cells": [
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {
    "id": "2IrD-ZoJQAPD"
   },
   "source": [
    "[![Binder](https://mybinder.org/badge_logo.svg)](https://mybinder.org/v2/gh/daviddavo/pocket2omnivore/HEAD?labpath=pocket2omnivore.ipynb)\n",
    "\n",
    "> This notebook is part of the [How to export your Pocket data and migrate to Omnivore](https://blog.ddavo.me/posts/tutorials/pocket-to-omnivore) tutorial\n",
    "\n",
    "# 1. Upload, parse, and store the Pocket File. \n",
    "\n",
    "First, let upload the `ril_export.html` file generated in https://getpocket.com/export\n",
    "\n",
    "The html has the following extructure:\n",
    "\n",
    "- `<h1>` Unread\n",
    "- `<ul>` with list items of `<a>`. The href is the link to the article, and the anchor text is the title. It also has a `tags` and `time_added` attributes.\n",
    "- `<h1>` Read\n",
    "- Another `<ul>` like the one above\n",
    "\n",
    "We will transform this into a dict of: \n",
    "- *read*: Boolean on wether the article has been read\n",
    "- *time_added*: The time the item was added\n",
    "- *tags*: An array of strings\n",
    "- *href*: The url\n",
    "- *title*: The title of the article"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 73
    },
    "id": "ez79reGjPokG",
    "outputId": "13422d0c-3dff-40ff-b0aa-c43fad089867",
    "tags": []
   },
   "outputs": [],
   "source": [
    "from ipywidgets import FileUpload\n",
    "from IPython.display import display\n",
    "from datetime import datetime\n",
    "from pathlib import Path\n",
    "import logging\n",
    "\n",
    "upload = FileUpload(accept='.html', multiple=False)\n",
    "\n",
    "def save_file():\n",
    "    for v in upload.value:\n",
    "        content = v['content']\n",
    "        with open(v['name'], 'wb') as f:\n",
    "            f.write(bytes(content))\n",
    "\n",
    "upload.observe(save_file, names='value')\n",
    "\n",
    "display(upload)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "assert Path(\"ril_export.html\").exists(), \"Upload the file before continue running\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "wEUX9qenQIBM",
    "tags": []
   },
   "outputs": [],
   "source": [
    "from bs4 import BeautifulSoup\n",
    "import os\n",
    "\n",
    "OMNIVORE_API_URL = \"https://api-prod.omnivore.app/api/graphql\"\n",
    "# The API key will have the following format \"00000000-0000-0000-0000-000000000000\"\n",
    "OMNIVORE_API_KEY = os.environ.get('OMNIVORE_API_KEY')\n",
    "SCHEMA_URL = \"https://raw.githubusercontent.com/omnivore-app/omnivore/c9fcbe72ddc6f40dd06e7073b8ffe3c1e71bd650/packages/api/src/generated/schema.graphql\"\n",
    "REQUESTS_SLEEP_TIME = 60 # Number of seconds\n",
    "PARALLEL_API_CALL_SIZE = 4\n",
    "\n",
    "if not OMNIVORE_API_KEY:\n",
    "    OMNIVORE_API_KEY=input('Enter your omnivore API key (should have a format similar to 00000000-0000-0000-0000-000000000000)')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "ka7QucTeQUFs",
    "outputId": "c85f5ad9-a75e-4a47-fa7b-e26c7c1313f5"
   },
   "outputs": [],
   "source": [
    "with open('ril_export.html', 'r') as f:\n",
    "    soup = BeautifulSoup(f, 'html.parser')\n",
    "\n",
    "soup.title"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Extract the articles and tags from the HTML doc."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "IEYZUjJFRshV"
   },
   "outputs": [],
   "source": [
    "def process_list(h1):\n",
    "    ul = h1.find_next_sibling('ul')\n",
    "    print(len(ul), h1.text, 'articles')\n",
    "    read = h1.text != 'Unread'\n",
    "\n",
    "    items = []\n",
    "    for a in ul.findAll('a', href=True):\n",
    "        items.append({\n",
    "            'read': read,\n",
    "            'time_added': datetime.fromtimestamp(int(a['time_added'])),\n",
    "            'href': a['href'],\n",
    "            'tags': a['tags'].split(',') if a['tags'] else [],\n",
    "            'title': a.text,\n",
    "        })\n",
    "\n",
    "    return items\n",
    "\n",
    "articles = [item for sublist in [process_list(h1) for h1 in soup.findAll('h1')] for item in sublist]\n",
    "labels = set([item for sublist in [article['tags'] for article in articles] for item in sublist])"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 2. Store the articles and tags in a SQLLite Database\n",
    "\n",
    "We want to be able to track our process, as the API for Omnivore has rate limiting, and takes a while to upload the files. For this we will use a SQL Database.\n",
    "\n",
    "We will store the Articles in the format: \n",
    "- *read*: Boolean on wether the article has been read\n",
    "- *time_added*: The time the item was added\n",
    "- *tags*: An array of strings\n",
    "- *href*: The url\n",
    "- *id*: nullable field that stores the omnivore file.\n",
    "\n",
    "and the tags with:\n",
    "- *name*: Name of the tag \n",
    "- *id*: if the tag has been stored. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from enum import Enum\n",
    "import peewee\n",
    "from peewee import SqliteDatabase, Model\n",
    "from playhouse.sqlite_ext import JSONField\n",
    "from playhouse.reflection import print_table_sql\n",
    "\n",
    "class StrEnum(str, Enum):\n",
    "    pass\n",
    "\n",
    "db = peewee.SqliteDatabase('omnivore.db')\n",
    "\n",
    "class BaseModel(peewee.Model):\n",
    "    class Meta:\n",
    "        database = db\n",
    "\n",
    "class ArticleStatus(StrEnum):\n",
    "    PROCESSING = \"PROCESSING\"\n",
    "    SUCCEEDED = \"SUCCEEDED\"\n",
    "    RETRY = \"_RETRY\"\n",
    "    DONE = \"_DONE\"\n",
    "\n",
    "class Article(BaseModel):\n",
    "    id = peewee.TextField(null=True)\n",
    "    read = peewee.BooleanField()\n",
    "    time_added = peewee.DateTimeField()\n",
    "    tags = JSONField(default=[])\n",
    "    href = peewee.TextField(primary_key=True)\n",
    "    status = peewee.TextField(null=True, choices=ArticleStatus)\n",
    "    slug = peewee.TextField(null=True)\n",
    "\n",
    "class Tag(BaseModel):\n",
    "    id = peewee.TextField(null=True)\n",
    "    name = peewee.TextField(primary_key=True)\n",
    "\n",
    "print_table_sql(Article)\n",
    "print_table_sql(Tag)\n",
    "\n",
    "# Create the Article Table \n",
    "db.connect()\n",
    "db.create_tables([Article, Tag])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "for label in labels:\n",
    "    Tag.get_or_create(name=label)\n",
    "# insert_tag_sql = f\"\"\"INSERT OR IGNORE into tags (name) values (?)\"\"\"\n",
    "# cursor.executemany(insert_tag_sql, [(label,) for label in labels])\n",
    "\n",
    "for article in articles:\n",
    "    Article.get_or_create(\n",
    "        read = article['read'],\n",
    "        time_added = article['time_added'],\n",
    "        href = article['href'],\n",
    "        tags = article['tags'],\n",
    "    )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "Ei9PTNIzawUo",
    "outputId": "e0adcbcd-278e-4d3f-d5c2-c5db0c9c3de7"
   },
   "outputs": [],
   "source": [
    "import requests\n",
    "\n",
    "with requests.get(SCHEMA_URL) as r:\n",
    "    r.raise_for_status()\n",
    "    schema = r.text\n",
    "\n",
    "    assert schema is not None\n",
    "\n",
    "print(schema[:100])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "vIQr1xQ9fq-W"
   },
   "outputs": [],
   "source": [
    "from gql import gql, Client\n",
    "from gql.transport.requests import RequestsHTTPTransport\n",
    "\n",
    "def create_client():\n",
    "    transport = RequestsHTTPTransport(\n",
    "       url=OMNIVORE_API_URL,\n",
    "        headers = {\n",
    "            'authorization': OMNIVORE_API_KEY,\n",
    "        }\n",
    "    )\n",
    "    return Client(transport=transport, schema=schema, fetch_schema_from_transport=False, execute_timeout=None)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "sJNbK-UaBr0-",
    "outputId": "555ac6ee-8e64-431f-928c-475f8a1a8d02"
   },
   "outputs": [],
   "source": [
    "# Doing a \"test query\" to check if everything is correct\n",
    "\n",
    "with create_client() as session: \n",
    "    r = session.execute(gql(\"\"\"\n",
    "    query Viewer {\n",
    "        me {\n",
    "            id\n",
    "            name\n",
    "            profile {\n",
    "                username\n",
    "            }\n",
    "        }\n",
    "    }\n",
    "    \"\"\"))\n",
    "\n",
    "    result = r\n",
    "    USERNAME = result['me']['profile']['username']\n",
    "\n",
    "    print(f\"Hello {result['me']['name']} ({USERNAME})!\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def getExistingTags():\n",
    "  with create_client() as session: \n",
    "    r = session.execute(gql(\"\"\"\n",
    "    query Labels {\n",
    "        labels {\n",
    "              ...on LabelsSuccess { \n",
    "                  labels { name, id }\n",
    "              }\n",
    "          }\n",
    "    }\n",
    "    \"\"\"))\n",
    "\n",
    "    result = r\n",
    "    return result['labels']['labels']\n",
    "\n",
    "#Then remove all the tags from the ones we created before\n",
    "def saveTags(tagName): \n",
    "    with create_client() as client: \n",
    "      mutation = f\"\"\"\n",
    "      mutation {{\n",
    "        createLabel(input: {{color: \"#F00\", name: \"{tagName}\" }}) {{\n",
    "          ... on CreateLabelSuccess {{\n",
    "            label {{\n",
    "              id\n",
    "              name\n",
    "              color\n",
    "              description\n",
    "              createdAt\n",
    "            }}\n",
    "          }}\n",
    "          ... on CreateLabelError {{\n",
    "            errorCodes\n",
    "          }}\n",
    "        }}\n",
    "      }}\n",
    "      \"\"\"\n",
    "\n",
    "      r = client.execute(gql(mutation), {'name': str(tagName)})\n",
    "      print(r)\n",
    "      return r['createLabel']['label']['id']\n",
    "\n",
    "server_tags = getExistingTags()\n",
    "unsaved_tags = list(Tag.select().where(Tag.id.is_null()))\n",
    "\n",
    "# tagIds = {f\"{dictionary['name']}\": dictionary[\"id\"] for dictionary in presaved_tags + server_tags}\n",
    "name2id = { t['name']:t['id'] for t in server_tags }\n",
    "for tag in unsaved_tags:\n",
    "    if tag.name not in name2id:\n",
    "      try:\n",
    "          id = saveTags(tag.name)\n",
    "          name2id[tag] = id\n",
    "      except Exception as e:\n",
    "          print(\"An error occurred:\", e)\n",
    "\n",
    "tags = [ Tag(name=k, id=v) for k,v in name2id.items() ]\n",
    "\n",
    "# query = \"UPDATE tags SET id = ? where name = ?\"\n",
    "# cursor.executemany(query, [(value, key) for key, value in tagIds.items()])\n",
    "Tag.bulk_update(tags, fields=[Tag.id])\n",
    "\n",
    "all_tags = list(Tag.select())\n",
    "all_tags"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "kxdnsdP4hYA8",
    "outputId": "16a47c1a-dee1-4fdf-d6d3-74ee2add63b1"
   },
   "outputs": [],
   "source": [
    "import backoff\n",
    "\n",
    "createArticle = gql(\"\"\"\n",
    "  mutation CreateArticleSavingRequest($url: String!) {\n",
    "    createArticleSavingRequest(input: {url: $url}) {\n",
    "      ... on CreateArticleSavingRequestSuccess {\n",
    "        articleSavingRequest {\n",
    "          id\n",
    "          status\n",
    "          slug\n",
    "          createdAt\n",
    "          updatedAt\n",
    "          url\n",
    "          errorCode\n",
    "        }\n",
    "      }\n",
    "      ... on CreateArticleSavingRequestError {\n",
    "        errorCodes\n",
    "      }\n",
    "    }\n",
    "  }\n",
    "\"\"\")\n",
    "\n",
    "setLabels = gql(\"\"\"\n",
    "mutation SetLabel($articleId: ID!, $labelIds: [ID!]!) { \n",
    "    setLabels(input: {pageId: $articleId, labelIds: $labelIds}) {\n",
    "        ...on SetLabelsSuccess { \n",
    "            labels { \n",
    "                id\n",
    "            }\n",
    "        }\n",
    "    }\n",
    "}\n",
    "\"\"\")\n",
    "\n",
    "updatePageSavedDate =  gql(\"\"\"\n",
    "mutation UpdatePageDate($id: ID!, $date: Date!) {\n",
    "    updatePage(input: {pageId: $id, savedAt: $date, publishedAt: $date}) {\n",
    "        ... on UpdatePageSuccess {\n",
    "            updatedPage {\n",
    "                id\n",
    "                savedAt\n",
    "                publishedAt\n",
    "                title\n",
    "            }\n",
    "        }\n",
    "        ...on UpdatePageError {\n",
    "            errorCodes\n",
    "        }\n",
    "    }\n",
    "}\n",
    "\"\"\")\n",
    "\n",
    "archivePage = gql(\"\"\"\n",
    "mutation ArchivePage($id: ID!) {\n",
    "    setLinkArchived (input: {linkId: $id, archived: true}) {\n",
    "        ... on ArchiveLinkSuccess {\n",
    "            linkId\n",
    "            message\n",
    "        }\n",
    "        ... on ArchiveLinkError {\n",
    "            message\n",
    "            errorCodes\n",
    "        }\n",
    "    }\n",
    "}\n",
    "\"\"\")\n",
    "\n",
    "createTag = gql(\"\"\"\n",
    "mutation CreateLabel($nam: String!, $col: String, $desc: String) {\n",
    "  createLabel(input: {name: $nam, color: $col, description: $desc}) {\n",
    "    ... on CreateLabelSuccess {\n",
    "      label {\n",
    "        id\n",
    "        name\n",
    "        color\n",
    "        description\n",
    "        createdAt\n",
    "      }\n",
    "    }\n",
    "    ... on CreateLabelError {\n",
    "      errorCodes\n",
    "    }\n",
    "  }\n",
    "}\n",
    "\"\"\")\n",
    "\n",
    "getArticleSavingRequest = gql(\"\"\"\n",
    "query ArticleSavingRequest($id: ID!) {\n",
    "    articleSavingRequest(id: $id) {\n",
    "        ... on ArticleSavingRequestSuccess {\n",
    "            articleSavingRequest {\n",
    "              id\n",
    "              status\n",
    "              slug\n",
    "              createdAt\n",
    "              updatedAt\n",
    "              url\n",
    "              errorCode\n",
    "            }\n",
    "        }\n",
    "        ... on ArticleSavingRequestError {\n",
    "            errorCodes\n",
    "        }\n",
    "    }\n",
    "}\n",
    "\"\"\")\n",
    "\n",
    "getArticle = gql(\"\"\"\n",
    "query GetArticle($slug: String!, $username: String!) {\n",
    "    article(slug: $slug, username: $username) {\n",
    "        ...on ArticleSuccess {\n",
    "            article {\n",
    "                id\n",
    "                title\n",
    "                slug\n",
    "                isArchived\n",
    "                savedAt\n",
    "            }\n",
    "        }\n",
    "        ...on ArticleError {\n",
    "            errorCodes\n",
    "        }\n",
    "    }\n",
    "}\n",
    "\"\"\")\n",
    "                \n",
    "@backoff.on_predicate(\n",
    "    backoff.runtime,\n",
    "    predicate=lambda r: isinstance(r, RequestsHTTPTransport),\n",
    "    value=lambda r: int(r.response_headers[\"RateLimit-Reset\"]) + 1,\n",
    "    jitter=None,\n",
    ")\n",
    "def archiveArticle(articleId): \n",
    "  with create_client() as client: \n",
    "    try: \n",
    "      res = client.execute(archivePage, { 'id': articleId })\n",
    "      return res\n",
    "    except Exception as e:\n",
    "          if (hasattr(e, 'code') and e.code == 429): \n",
    "            return session.transport\n",
    "          raise\n",
    "\n",
    "@backoff.on_predicate(\n",
    "    backoff.runtime,\n",
    "    predicate=lambda r: isinstance(r, RequestsHTTPTransport),\n",
    "    value=lambda r: int(r.response_headers[\"RateLimit-Reset\"]) + 1,\n",
    "    jitter=None,\n",
    ")\n",
    "def saveLabels(articleId, labels): \n",
    "    with create_client() as client: \n",
    "      try:\n",
    "        return client.execute(setLabels, {'articleId': articleId, 'labelIds': labels})\n",
    "      except Exception as e:\n",
    "          if (hasattr(e, 'code') and e.code == 429): \n",
    "            return session.transport\n",
    "          raise\n",
    "\n",
    "@backoff.on_predicate(\n",
    "    backoff.runtime,\n",
    "    predicate=lambda r: isinstance(r, RequestsHTTPTransport),\n",
    "    value=lambda r: int(r.response_headers[\"RateLimit-Reset\"]) + 1,\n",
    "    jitter=None,\n",
    ")\n",
    "def saveArticle(article):\n",
    "    with create_client() as client: \n",
    "      try: \n",
    "        logging.info(f\"Saving article {article.href}\")\n",
    "        url = article.href\n",
    "        tags = article.tags\n",
    "        # First createArticleSavingRequest\n",
    "        r = client.execute(createArticle, variable_values={'url': url})\n",
    "        article.id = r['createArticleSavingRequest']['articleSavingRequest']['id']\n",
    "        article.status = r['createArticleSavingRequest']['articleSavingRequest']['status']\n",
    "        \n",
    "        if len(tags) != 0:\n",
    "            saveLabels(article.id, tags)\n",
    "                            \n",
    "        # Return the article with the id of the saved document\n",
    "        return article\n",
    "      except Exception as e:\n",
    "          if (hasattr(e, 'code') and e.code == 429): \n",
    "            return session.transport\n",
    "          # I don't know why this happens and I will figure it out later.\n",
    "          raise\n",
    "\n",
    "@backoff.on_predicate(\n",
    "    backoff.runtime,\n",
    "    predicate=lambda r: isinstance(r, RequestsHTTPTransport),\n",
    "    value=lambda r: int(r.response_headers[\"RateLimit-Reset\"]) + 1,\n",
    "    jitter=None,\n",
    ")\n",
    "def updateArticleTimeAfterProcessing(articleId, date: datetime = None):\n",
    "    with create_client() as client: \n",
    "        try: \n",
    "            if date is not None: \n",
    "              # Wait a bit, it seems there's a race condition.\n",
    "              res = client.execute(updatePageSavedDate, {\n",
    "                  'id': articleId,\n",
    "                  'date': date.isoformat(),\n",
    "              })\n",
    "              return res\n",
    "\n",
    "        except Exception as e:\n",
    "          if (hasattr(e, 'code') and e.code == 429): \n",
    "            return session.transport\n",
    "          \n",
    "          # I don't know why this happens and I will figure it out later.\n",
    "          raise\n",
    "\n",
    "@backoff.on_predicate(\n",
    "    backoff.runtime,\n",
    "    predicate=lambda r: isinstance(r, RequestsHTTPTransport),\n",
    "    value=lambda r: int(r.response_headers[\"RateLimit-Reset\"]) + 1,\n",
    "    jitter=None,\n",
    ")\n",
    "def articleSavingRequest(articleId):\n",
    "    with create_client() as client:\n",
    "        try: \n",
    "            # Wait a bit, it seems there's a race condition.\n",
    "            res = client.execute(getArticleSavingRequest, {\n",
    "                'id': articleId,\n",
    "            })\n",
    "\n",
    "            return res['articleSavingRequest']['articleSavingRequest']\n",
    "        except Exception as e:\n",
    "          if (hasattr(e, 'code') and e.code == 429): \n",
    "            return session.transport\n",
    "          \n",
    "          # I don't know why this happens and I will figure it out later.\n",
    "          raise\n",
    "\n",
    "@backoff.on_predicate(\n",
    "    backoff.runtime,\n",
    "    predicate=lambda r: isinstance(r, RequestsHTTPTransport),\n",
    "    value=lambda r: int(r.response_headers[\"RateLimit-Reset\"]) + 1,\n",
    "    jitter=None,\n",
    ")\n",
    "def getArticleInfo(slug):\n",
    "    with create_client() as client:\n",
    "        try: \n",
    "            # Wait a bit, it seems there's a race condition.\n",
    "            res = client.execute(getArticle, {\n",
    "                'slug': slug,\n",
    "                'username': USERNAME,\n",
    "            })\n",
    "\n",
    "            return res['article']['article']\n",
    "        except Exception as e:\n",
    "          if (hasattr(e, 'code') and e.code == 429): \n",
    "            return session.transport\n",
    "          \n",
    "          # I don't know why this happens and I will figure it out later.\n",
    "          raise"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Thread Pool executors fall a bit behind on our usage, so I will create a custom one\n",
    "\n",
    "It has a number of `worker` threads that just spin on a queue of functions executing the function."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import datetime as dt\n",
    "import pytz\n",
    "\n",
    "def checkTask(article):\n",
    "    \"\"\" Runs a check agains Omnivore, if the article is processed updates its date,\n",
    "    Returns the new article object\n",
    "    \"\"\"\n",
    "    asr = articleSavingRequest(article.id)\n",
    "    article.status = asr['status']\n",
    "    article.slug = asr['slug']\n",
    "\n",
    "    if article.status != 'SUCCEEDED':\n",
    "        print(f\"Retrying article\", asr)\n",
    "        return article\n",
    "\n",
    "    updateArticleTimeAfterProcessing(article.id, article.time_added)\n",
    "    if article.read:\n",
    "        archiveArticle(article.id)\n",
    "    \n",
    "    omnArt = getArticleInfo(asr['slug'])\n",
    "\n",
    "    dtexpected = article.time_added.replace(tzinfo=pytz.UTC)\n",
    "    dtgot = datetime.fromisoformat(omnArt['savedAt'])\n",
    "    \n",
    "    c1 = article.read == omnArt['isArchived']\n",
    "    c2 = abs((dtexpected - dtgot).total_seconds()) < 3600\n",
    "    if not c1 or not c2:\n",
    "        logstr = f\"Retrying article {omnArt['slug']}, \"\n",
    "        if not c1:\n",
    "            logstr += f\"Incorrect archive status. expected {article.read}, got {omnArt['isArchived']}. \"\n",
    "        if not c2:\n",
    "            logstr += f\"datetime not changed. expected {dtexpected}, got {dtgot}\"\n",
    "        logging.info(logstr)\n",
    "        \n",
    "        article.status = '_RETRY'\n",
    "        return article\n",
    "    \n",
    "    return article"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "Article.select()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from concurrent.futures import ThreadPoolExecutor, as_completed\n",
    "from multiprocessing.pool import ThreadPool\n",
    "from tqdm.notebook import tqdm\n",
    "from threading import Lock\n",
    "\n",
    "TIMEOUT = 120 # seconds\n",
    "\n",
    "unsaved_articles = list(Article.select().where(Article.id.is_null()))\n",
    "submitted_articles = list(Article.select().where(Article.id.is_null(False) & (Article.status != ArticleStatus.DONE)))\n",
    "done_articles = list(Article.select().where(Article.status == ArticleStatus.DONE))\n",
    "all_articles = unsaved_articles + submitted_articles + done_articles\n",
    "\n",
    "assert len(all_articles) == len(Article.select())\n",
    "tqdm_custom = lambda desc: tqdm(desc=desc, total=len(all_articles))\n",
    "\n",
    "with (ThreadPoolExecutor(PARALLEL_API_CALL_SIZE) as executor,\n",
    "      tqdm_custom('Submit') as pbar1,\n",
    "      tqdm_custom('Done') as pbar2):\n",
    "\n",
    "    try:\n",
    "        unchecked_articles = {}\n",
    "\n",
    "        # TODO: Load from db and repopulate unsaved and unchecked articles\n",
    "        fmap1 = { executor.submit(saveArticle, article) : article for article in unsaved_articles }\n",
    "\n",
    "        print(f\"Skipping submit of {len(submitted_articles)} articles\")\n",
    "        pbar1.update(len(submitted_articles))\n",
    "\n",
    "        print(f\"Skipping update of {len(done_articles)} previously done articles\")\n",
    "        pbar1.update(len(done_articles))\n",
    "        pbar2.update(len(done_articles))\n",
    "        \n",
    "        for f in as_completed(fmap1, timeout=TIMEOUT):\n",
    "            try:\n",
    "                newarticle = f.result()\n",
    "                pbar1.update()\n",
    "                unchecked_articles[executor.submit(checkTask, newarticle)] = newarticle\n",
    "                newarticle.save()\n",
    "                # Save state to db if needed\n",
    "            except Exception as e:\n",
    "                logging.error(repr(e))\n",
    "\n",
    "        for a in submitted_articles:\n",
    "            unchecked_articles[executor.submit(checkTask, a)] = a\n",
    "\n",
    "        pbar1.refresh()\n",
    "        pbar2.refresh()\n",
    "        while unchecked_articles:\n",
    "            try:\n",
    "                for f in as_completed(unchecked_articles, timeout=TIMEOUT*2):\n",
    "                    newarticle = f.result()\n",
    "                    if newarticle.status == 'SUCCEEDED': # Task completed\n",
    "                        newarticle.status = ArticleStatus.DONE\n",
    "                        newarticle.save()\n",
    "                        pbar2.update()\n",
    "                    elif newarticle.status == 'PROCESSING' or newarticle.status == '_RETRY':\n",
    "                        unchecked_articles[executor.submit(checkTask, newarticle)] = newarticle\n",
    "                    else:\n",
    "                        raise ValueError(f\"Unknown article status: {newarticle.status}\")\n",
    "                    unchecked_articles.pop(f)\n",
    "            except exception as e:\n",
    "                logging.error(repr(e))\n",
    "    except KeyboardInterrupt:\n",
    "        executor.shutdown(cancel_futures=True)\n",
    "        raise"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "colab": {
   "provenance": []
  },
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}