Repository: dogsheep/twitter-to-sqlite

Issue:
Extract "source" into a separate lookup table
It's pretty bulky and ugly at the moment:

<img width="334" alt="trump__tweets__1_820_rows" src="https://user-images.githubusercontent.com/9599/66264630-df23a080-e7bd-11e9-9154-403c2e69f841.png">


Comments:

Patch:
diff --git a/twitter_to_sqlite/cli.py b/twitter_to_sqlite/cli.py
index 56ba167..ad44655 100644
--- a/twitter_to_sqlite/cli.py
+++ b/twitter_to_sqlite/cli.py
@@ -5,7 +5,7 @@ import pathlib
 import time
 
 import click
-import sqlite_utils
+
 from twitter_to_sqlite import archive
 from twitter_to_sqlite import utils
 
@@ -102,7 +102,7 @@ def followers(db_path, auth, user_id, screen_name, silent):
     "Save followers for specified user (defaults to authenticated user)"
     auth = json.load(open(auth))
     session = utils.session_for_auth(auth)
-    db = sqlite_utils.Database(db_path)
+    db = utils.open_database(db_path)
     fetched = []
     # Get the follower count, so we can have a progress bar
     count = 0
@@ -152,7 +152,7 @@ def favorites(db_path, auth, user_id, screen_name, stop_after):
     "Save tweets favorited by specified user"
     auth = json.load(open(auth))
     session = utils.session_for_auth(auth)
-    db = sqlite_utils.Database(db_path)
+    db = utils.open_database(db_path)
     profile = utils.get_profile(db, session, user_id, screen_name)
     with click.progressbar(
         utils.fetch_favorites(session, user_id, screen_name, stop_after),
@@ -193,7 +193,7 @@ def user_timeline(db_path, auth, stop_after, user_id, screen_name, since, since_
         raise click.ClickException("Use either --since or --since_id, not both")
     auth = json.load(open(auth))
     session = utils.session_for_auth(auth)
-    db = sqlite_utils.Database(db_path)
+    db = utils.open_database(db_path)
     profile = utils.get_profile(db, session, user_id, screen_name)
     expected_length = profile["statuses_count"]
 
@@ -209,7 +209,9 @@ def user_timeline(db_path, auth, stop_after, user_id, screen_name, since, since_
             pass
 
     with click.progressbar(
-        utils.fetch_user_timeline(session, user_id, screen_name, stop_after, since_id=since_id),
+        utils.fetch_user_timeline(
+            session, user_id, screen_name, stop_after, since_id=since_id
+        ),
         length=expected_length,
         label="Importing tweets",
         show_pos=True,
@@ -253,7 +255,7 @@ def home_timeline(db_path, auth, since, since_id):
         raise click.ClickException("Use either --since or --since_id, not both")
     auth = json.load(open(auth))
     session = utils.session_for_auth(auth)
-    db = sqlite_utils.Database(db_path)
+    db = utils.open_database(db_path)
     profile = utils.get_profile(db, session)
     expected_length = 800
     if since and db["timeline_tweets"].exists:
@@ -310,7 +312,7 @@ def users_lookup(db_path, identifiers, attach, sql, auth, ids):
     "Fetch user accounts"
     auth = json.load(open(auth))
     session = utils.session_for_auth(auth)
-    db = sqlite_utils.Database(db_path)
+    db = utils.open_database(db_path)
     identifiers = utils.resolve_identifiers(db, identifiers, attach, sql)
     for batch in utils.fetch_user_batches(session, identifiers, ids):
         utils.save_users(db, batch)
@@ -338,7 +340,7 @@ def statuses_lookup(db_path, identifiers, attach, sql, auth, skip_existing, sile
     "Fetch tweets by their IDs"
     auth = json.load(open(auth))
     session = utils.session_for_auth(auth)
-    db = sqlite_utils.Database(db_path)
+    db = utils.open_database(db_path)
     identifiers = utils.resolve_identifiers(db, identifiers, attach, sql)
     if skip_existing:
         existing_ids = set(
@@ -381,7 +383,7 @@ def list_members(db_path, identifiers, auth, ids):
     "Fetch lists - accepts one or more screen_name/list_slug identifiers"
     auth = json.load(open(auth))
     session = utils.session_for_auth(auth)
-    db = sqlite_utils.Database(db_path)
+    db = utils.open_database(db_path)
     for identifier in identifiers:
         utils.fetch_and_save_list(db, session, identifier, ids)
 
@@ -477,7 +479,7 @@ def track(db_path, track, auth, verbose):
     "Experimental: Save tweets matching these keywords in real-time"
     auth = json.load(open(auth))
     session = utils.session_for_auth(auth)
-    db = sqlite_utils.Database(db_path)
+    db = utils.open_database(db_path)
     for tweet in utils.stream_filter(session, track=track):
         if verbose:
             print(json.dumps(tweet, indent=2))
@@ -505,7 +507,7 @@ def follow(db_path, identifiers, attach, sql, ids, auth, verbose):
     "Experimental: Follow these Twitter users and save tweets in real-time"
     auth = json.load(open(auth))
     session = utils.session_for_auth(auth)
-    db = sqlite_utils.Database(db_path)
+    db = utils.open_database(db_path)
     identifiers = utils.resolve_identifiers(db, identifiers, attach, sql)
     # Make sure we have saved these users to the database
     for batch in utils.fetch_user_batches(session, identifiers, ids):
@@ -528,7 +530,7 @@ def _shared_friends_ids_followers_ids(
 ):
     auth = json.load(open(auth))
     session = utils.session_for_auth(auth)
-    db = sqlite_utils.Database(db_path)
+    db = utils.open_database(db_path)
     identifiers = utils.resolve_identifiers(db, identifiers, attach, sql)
     for identifier in identifiers:
         # Make sure this user is saved
@@ -568,7 +570,7 @@ def import_(db_path, paths):
     Import data from a Twitter exported archive. Input can be the path to a zip
     file, a directory full of .js files or one or more direct .js files.
     """
-    db = sqlite_utils.Database(db_path)
+    db = utils.open_database(db_path)
     for filepath in paths:
         path = pathlib.Path(filepath)
         if path.suffix == ".zip":
diff --git a/twitter_to_sqlite/migrations.py b/twitter_to_sqlite/migrations.py
new file mode 100644
index 0000000..13d7c65
--- /dev/null
+++ b/twitter_to_sqlite/migrations.py
@@ -0,0 +1,22 @@
+from .utils import extract_and_save_source
+
+MIGRATIONS = []
+
+
+def migration(fn):
+    MIGRATIONS.append(fn)
+    return fn
+
+
+@migration
+def convert_source_column(db):
+    tables = set(db.table_names())
+    if "tweets" not in tables:
+        return
+    # Now we extract any '<a href=...' records from the source
+    for id, source in db.conn.execute(
+        "select id, source from tweets where source like '<%'"
+    ).fetchall():
+        db["tweets"].update(id, {"source": extract_and_save_source(db, source)})
+    db["tweets"].create_index(["source"])
+    db["tweets"].add_foreign_key("source")
diff --git a/twitter_to_sqlite/utils.py b/twitter_to_sqlite/utils.py
index 6f2a44e..1755c51 100644
--- a/twitter_to_sqlite/utils.py
+++ b/twitter_to_sqlite/utils.py
@@ -2,16 +2,46 @@ import datetime
 import html
 import json
 import pathlib
+import re
 import time
 import urllib.parse
 import zipfile
 
 from dateutil import parser
 from requests_oauthlib import OAuth1Session
+import sqlite_utils
 
 # Twitter API error codes
 RATE_LIMIT_ERROR_CODE = 88
 
+source_re = re.compile('<a href="(?P<url>.*?)".*?>(?P<name>.*?)</a>')
+
+
+def open_database(db_path):
+    db = sqlite_utils.Database(db_path)
+    # Only run migrations if this is an existing DB (has tables)
+    if db.tables:
+        migrate(db)
+    return db
+
+
+def migrate(db):
+    from twitter_to_sqlite.migrations import MIGRATIONS
+
+    if "migrations" not in db.table_names():
+        db["migrations"].create({"name": str, "applied": str}, pk="name")
+    applied_migrations = {
+        m[0] for m in db.conn.execute("select name from migrations").fetchall()
+    }
+    for migration in MIGRATIONS:
+        name = migration.__name__
+        if name in applied_migrations:
+            continue
+        migration(db)
+        db["migrations"].insert(
+            {"name": name, "applied": datetime.datetime.utcnow().isoformat()}
+        )
+
 
 def session_for_auth(auth):
     return OAuth1Session(
@@ -186,6 +216,8 @@ def ensure_tables(db):
     table_names = set(db.table_names())
     if "places" not in table_names:
         db["places"].create({"id": str}, pk="id")
+    if "sources" not in table_names:
+        db["sources"].create({"id": str, "name": str, "url": str}, pk="id")
     if "users" not in table_names:
         db["users"].create(
             {
@@ -210,9 +242,14 @@ def ensure_tables(db):
                 "retweeted_status": int,
                 "quoted_status": int,
                 "place": str,
+                "source": str,
             },
             pk="id",
-            foreign_keys=(("user", "users", "id"), ("place", "places", "id")),
+            foreign_keys=(
+                ("user", "users", "id"),
+                ("place", "places", "id"),
+                ("source", "sources", "id"),
+            ),
         )
         db["tweets"].enable_fts(["full_text"], create_triggers=True)
         db["tweets"].add_foreign_key("retweeted_status", "tweets")
@@ -235,6 +272,7 @@ def save_tweets(db, tweets, favorited_by=None):
         user = tweet.pop("user")
         transform_user(user)
         tweet["user"] = user["id"]
+        tweet["source"] = extract_and_save_source(db, tweet["source"])
         if tweet.get("place"):
             db["places"].upsert(tweet["place"], pk="id", alter=True)
             tweet["place"] = tweet["place"]["id"]
@@ -472,3 +510,9 @@ def read_archive_js(filepath):
     for zi in zf.filelist:
         if zi.filename.endswith(".js"):
             yield zi.filename, zf.open(zi.filename).read()
+
+
+def extract_and_save_source(db, source):
+    m = source_re.match(source)
+    details = m.groupdict()
+    return db["sources"].upsert(details, hash_id="id").last_pk


Test patch:
diff --git a/tests/test_migrations.py b/tests/test_migrations.py
new file mode 100644
index 0000000..5877d34
--- /dev/null
+++ b/tests/test_migrations.py
@@ -0,0 +1,49 @@
+import sqlite_utils
+from click.testing import CliRunner
+import sqlite_utils
+from twitter_to_sqlite import cli, migrations
+
+from .test_import import zip_contents_path
+
+
+def test_no_migrations_on_first_run(tmpdir, zip_contents_path):
+    output = str(tmpdir / "output.db")
+    args = ["import", output, str(zip_contents_path / "follower.js")]
+    result = CliRunner().invoke(cli.cli, args)
+    assert 0 == result.exit_code, result.stdout
+    db = sqlite_utils.Database(output)
+    assert ["archive_follower"] == db.table_names()
+    # Re-running the command again should also run the migrations
+    result = CliRunner().invoke(cli.cli, args)
+    db = sqlite_utils.Database(output)
+    assert {"archive_follower", "migrations"} == set(db.table_names())
+
+
+def test_convert_source_column():
+    db = sqlite_utils.Database(memory=True)
+    db["tweets"].insert_all(
+        [
+            {"id": 1, "source": '<a href="URL">NAME</a>'},
+            {"id": 2, "source": '<a href="URL2">NAME2</a>'},
+            {"id": 3, "source": "d3c1d39c57fecfc09202f20ea5e2db30262029fd"},
+        ],
+        pk="id",
+    )
+    migrations.convert_source_column(db)
+    assert [
+        {
+            "id": "d3c1d39c57fecfc09202f20ea5e2db30262029fd",
+            "url": "URL",
+            "name": "NAME",
+        },
+        {
+            "id": "000e4c4db71278018fb8c322f070d051e76885b1",
+            "url": "URL2",
+            "name": "NAME2",
+        },
+    ] == list(db["sources"].rows)
+    assert [
+        {"id": 1, "source": "d3c1d39c57fecfc09202f20ea5e2db30262029fd"},
+        {"id": 2, "source": "000e4c4db71278018fb8c322f070d051e76885b1"},
+        {"id": 3, "source": "d3c1d39c57fecfc09202f20ea5e2db30262029fd"},
+    ] == list(db["tweets"].rows)
diff --git a/tests/test_save_tweets.py b/tests/test_save_tweets.py
index 4bbe7ae..dc87a71 100644
--- a/tests/test_save_tweets.py
+++ b/tests/test_save_tweets.py
@@ -20,6 +20,7 @@ def db(tweets):
 
 def test_tables(db):
     assert {
+        "sources",
         "users_fts_idx",
         "users_fts_data",
         "tweets_fts",
@@ -182,9 +183,9 @@ def test_tweets(db):
             "retweeted_status": None,
             "quoted_status": None,
             "place": None,
+            "source": "e6528b505bcfd811fdd40ff2d46665dbccba2024",
             "truncated": 0,
             "display_text_range": "[0, 139]",
-            "source": '<a href="http://itunes.apple.com/us/app/twitter/id409789998?mt=12" rel="nofollow">Twitter for Mac</a>',
             "in_reply_to_status_id": None,
             "in_reply_to_user_id": None,
             "in_reply_to_screen_name": None,
@@ -207,9 +208,9 @@ def test_tweets(db):
             "retweeted_status": None,
             "quoted_status": 861696799362478100,
             "place": None,
+            "source": "1f89d6a41b1505a3071169f8d0d028ba9ad6f952",
             "truncated": 0,
             "display_text_range": "[0, 239]",
-            "source": '<a href="https://mobile.twitter.com" rel="nofollow">Twitter Web App</a>',
             "in_reply_to_status_id": None,
             "in_reply_to_user_id": None,
             "in_reply_to_screen_name": None,
@@ -232,9 +233,9 @@ def test_tweets(db):
             "retweeted_status": None,
             "quoted_status": None,
             "place": "01a9a39529b27f36",
+            "source": "95f3aaaddaa45937ac94765e0ddb68ba2be92d20",
             "truncated": 0,
             "display_text_range": "[45, 262]",
-            "source": '<a href="http://twitter.com/download/iphone" rel="nofollow">Twitter for iPhone</a>',
             "in_reply_to_status_id": "1169079390577320000",
             "in_reply_to_user_id": "82016165",
             "in_reply_to_screen_name": "scientiffic",
@@ -257,9 +258,9 @@ def test_tweets(db):
             "retweeted_status": None,
             "quoted_status": None,
             "place": None,
+            "source": "942cfc2bf9f290ddbe3d78f1907dc084a00ed23f",
             "truncated": 0,
             "display_text_range": "[0, 235]",
-            "source": '<a href="http://www.voxmedia.com" rel="nofollow">Vox Media</a>',
             "in_reply_to_status_id": None,
             "in_reply_to_user_id": None,
             "in_reply_to_screen_name": None,
@@ -282,9 +283,9 @@ def test_tweets(db):
             "retweeted_status": 1169242008432644000,
             "quoted_status": None,
             "place": None,
+            "source": "95f3aaaddaa45937ac94765e0ddb68ba2be92d20",
             "truncated": 0,
             "display_text_range": "[0, 143]",
-            "source": '<a href="http://twitter.com/download/iphone" rel="nofollow">Twitter for iPhone</a>',
             "in_reply_to_status_id": None,
             "in_reply_to_user_id": None,
             "in_reply_to_screen_name": None,
@@ -302,6 +303,32 @@ def test_tweets(db):
     ] == tweet_rows
 
 
+def test_sources(db):
+    source_rows = list(db["sources"].rows)
+    assert [
+        {
+            "id": "942cfc2bf9f290ddbe3d78f1907dc084a00ed23f",
+            "name": "Vox Media",
+            "url": "http://www.voxmedia.com",
+        },
+        {
+            "id": "95f3aaaddaa45937ac94765e0ddb68ba2be92d20",
+            "name": "Twitter for iPhone",
+            "url": "http://twitter.com/download/iphone",
+        },
+        {
+            "id": "1f89d6a41b1505a3071169f8d0d028ba9ad6f952",
+            "name": "Twitter Web App",
+            "url": "https://mobile.twitter.com",
+        },
+        {
+            "id": "e6528b505bcfd811fdd40ff2d46665dbccba2024",
+            "name": "Twitter for Mac",
+            "url": "http://itunes.apple.com/us/app/twitter/id409789998?mt=12",
+        },
+    ] == source_rows
+
+
 def test_places(db):
     place_rows = list(db["places"].rows)
     assert [