Optimization: Code

"""Confirm that imdb.db built correctly, read-only.

Usage: python3 check_database.py [--database imdb.db]

Opens the database with mode=ro, so this never creates or changes a file.
Prints the table list, a row count for each table, and one joined row, so you
can compare the counts against what import_imdb.py reported.
"""
import argparse
from pathlib import Path
import sqlite3


def check(database):
    print(Path(database).resolve())
    connection = sqlite3.connect(f"file:{database}?mode=ro", uri=True)
    print(connection.execute(
        "select name from sqlite_master where type = 'table' order by name"
    ).fetchall())
    for table in ("titles", "ratings"):
        count = connection.execute(f"select count(*) from {table}").fetchone()[0]
        print(table, count)

    print(connection.execute("""
        select t.tconst, t.primary_title, t.start_year,
               r.average_rating, r.num_votes
        from titles t left join ratings r on r.tconst = t.tconst
        where t.tconst = ?
    """, ("tt0000001",)).fetchone())
    connection.close()


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Check that imdb.db built correctly.")
    parser.add_argument("--database", default="imdb.db")
    args = parser.parse_args()
    check(args.database)

"""Compare doing the work in Python with asking the database to do it.

Usage: python3 compare_python_sql.py [--database imdb.db]

Two comparisons on imdb.db. The first looks up 2,000 titles one query at a
time and then with one join. The second counts movies per decade, first by
fetching every year into Python and then with GROUP BY.

SQLite runs inside the Python process, so it hides the cost of sending rows
over a network. With a database server the differences would be larger.
"""
import argparse
from collections import Counter
import sqlite3
import time

INDEXES = [
    "create index if not exists idx_ratings_votes on ratings(num_votes)",
    "create index if not exists idx_titles_type_year on titles(title_type, start_year)",
]


def timed(function):
    started = time.perf_counter()
    result = function()
    return result, time.perf_counter() - started


def one_query_per_title(connection, ids):
    total = 0
    for tconst in ids:
        row = connection.execute(
            "select primary_title from titles where tconst = ?", (tconst,)).fetchone()
        total += len(row[0])
    return total


def one_join(connection, count):
    rows = connection.execute(
        "select t.primary_title from ratings r join titles t on t.tconst = r.tconst "
        "order by r.num_votes desc limit ?", (count,)).fetchall()
    return sum(len(row[0]) for row in rows)


def movies_per_decade_in_python(connection):
    counts = Counter()
    for (year,) in connection.execute(
            "select start_year from titles where title_type = 'movie' and start_year is not null"):
        counts[year // 10 * 10] += 1
    return dict(sorted(counts.items()))


def movies_per_decade_in_sql(connection):
    return dict(connection.execute(
        "select start_year / 10 * 10 as decade, count(*) from titles "
        "where title_type = 'movie' and start_year is not null "
        "group by decade order by decade"))


def main():
    parser = argparse.ArgumentParser(description="Compare Python and SQL.")
    parser.add_argument("--database", default="imdb.db")
    args = parser.parse_args()
    connection = sqlite3.connect(args.database)
    for statement in INDEXES:
        connection.execute(statement)
    connection.commit()

    count = 2000
    ids = [row[0] for row in connection.execute(
        "select tconst from ratings order by num_votes desc limit ?", (count,))]
    _, slow = timed(lambda: one_query_per_title(connection, ids))
    _, fast = timed(lambda: one_join(connection, count))
    print(f"Titles for the {count:,} most-voted ratings")
    print(f"  one query per title: {slow * 1000:8.1f} ms  ({count:,} queries)")
    print(f"  one join:            {fast * 1000:8.1f} ms  (1 query)")

    in_python, python_time = timed(lambda: movies_per_decade_in_python(connection))
    in_sql, sql_time = timed(lambda: movies_per_decade_in_sql(connection))
    print("\nMovies per decade")
    print(f"  fetch every year, count in Python: {python_time * 1000:8.1f} ms  ({sum(in_python.values()):,} rows fetched)")
    print(f"  GROUP BY in SQL:                   {sql_time * 1000:8.1f} ms  ({len(in_sql)} rows fetched)")
    print(f"  same answer: {in_python == in_sql}")
    for decade in list(in_sql)[-4:]:
        print(f"    {decade}s: {in_sql[decade]:,}")
    connection.close()


if __name__ == "__main__":
    main()

"""Watch the query planner and the clock as indexes are added to imdb.db.

Usage: python3 explore_indexes.py [--database imdb.db]

Each step shows a query, the plan SQLite chose (EXPLAIN QUERY PLAN), and the
best of three run times. Steps that add an index show the time to build it and
the growth of the file. The last two steps show what ANALYZE changes. The script removes every index it creates (names that
start with idx_) and any statistics from ANALYZE before it starts, so it can be
run again.
"""
import argparse
import sqlite3
import time

TITLE = "The Godfather"

# (heading, sql, parameters, index to create first or None)
STEPS = [
    ("Lookup by the primary key (tconst)",
     "select primary_title, start_year from titles where tconst = ?",
     ("tt0068646",), None),
    ("Lookup by a column with no index",
     "select tconst, start_year from titles where primary_title = ?",
     (TITLE,), None),
    ("The same query after an index on primary_title",
     "select tconst, start_year from titles where primary_title = ?",
     (TITLE,), "create index idx_titles_primary_title on titles(primary_title)"),
    ("LIKE with a prefix, against a case-sensitive index",
     "select tconst, primary_title from titles where primary_title like ?",
     ("The Godfather%",), None),
    ("The same LIKE after an index that ignores case",
     "select tconst, primary_title from titles where primary_title like ?",
     ("The Godfather%",),
     "create index idx_titles_primary_title_nocase on titles(primary_title collate nocase)"),
    ("LIKE with a wildcard at the front",
     "select tconst, primary_title from titles where primary_title like ?",
     ("%Godfather%",), None),
    ("Two conditions, no index that covers both",
     "select count(*) from titles where title_type = ? and start_year = ?",
     ("movie", 1994), None),
    ("The same query after an index on start_year",
     "select count(*) from titles where title_type = ? and start_year = ?",
     ("movie", 1994), "create index idx_titles_year on titles(start_year)"),
    ("The same query after a two-column index",
     "select count(*) from titles where title_type = ? and start_year = ?",
     ("movie", 1994),
     "create index idx_titles_type_year on titles(title_type, start_year)"),
    ("The second column alone, with the single-column index gone",
     "select count(*) from titles where start_year = ?",
     (1994,), "drop index idx_titles_year"),
    ("The first column alone",
     "select count(*) from titles where title_type = ?",
     ("movie",), None),
    ("Returning a column that is not in the index",
     "select primary_title from titles where title_type = ? and start_year = ?",
     ("movie", 1994), None),
    ("The same query after a covering index",
     "select primary_title from titles where title_type = ? and start_year = ?",
     ("movie", 1994),
     "create index idx_titles_type_year_title on titles(title_type, start_year, primary_title)"),
    ("Two OR-connected conditions, one of them with no index",
     "select tconst from titles where start_year = ? or primary_title = ?",
     (1894, TITLE), None),
    ("Two OR-connected conditions, each with an index",
     "select tconst from titles where start_year = ? or primary_title = ?",
     (1894, TITLE), "create index idx_titles_year on titles(start_year)"),
    ("The most-voted titles, with no index on num_votes",
     "select tconst, num_votes from ratings order by num_votes desc limit 10",
     (), None),
    ("The same query after an index on num_votes",
     "select tconst, num_votes from ratings order by num_votes desc limit 10",
     (), "create index idx_ratings_votes on ratings(num_votes)"),
    ("A join on the primary keys",
     "select t.primary_title, r.num_votes from ratings r "
     "join titles t on t.tconst = r.tconst "
     "order by r.num_votes desc limit 10",
     (), None),
    ("The second column alone, before ANALYZE",
     "select count(*) from titles where start_year = ?",
     (1994,), "drop index idx_titles_year"),
    ("The second column alone, after ANALYZE",
     "select count(*) from titles where start_year = ?",
     (1994,), "analyze"),
]


def drop_our_indexes(connection):
    names = connection.execute(
        "select name from sqlite_master where type = 'index' and name like 'idx_%'"
    ).fetchall()
    for (name,) in names:
        connection.execute(f"drop index {name}")
    has_statistics = connection.execute(
        "select count(*) from sqlite_master where name = 'sqlite_stat1'").fetchone()[0]
    if has_statistics:
        connection.execute("delete from sqlite_stat1")
    connection.commit()


def plan(connection, sql, parameters):
    rows = connection.execute("explain query plan " + sql, parameters).fetchall()
    return [row[3] for row in rows]


def best_time(connection, sql, parameters, repeats=3):
    best = None
    for _ in range(repeats):
        started = time.perf_counter()
        connection.execute(sql, parameters).fetchall()
        elapsed = time.perf_counter() - started
        best = elapsed if best is None else min(best, elapsed)
    return best


def size_mb(connection):
    """Megabytes of the file in use. Pages freed by DROP INDEX do not count."""
    page_size = connection.execute("pragma page_size").fetchone()[0]
    pages = connection.execute("pragma page_count").fetchone()[0]
    free = connection.execute("pragma freelist_count").fetchone()[0]
    return (pages - free) * page_size / 1_000_000


def write_cost(rows=500_000):
    """Time the same inserts into a scratch table with 0, 1, and 3 indexes."""
    print("\n== What indexes cost when writing")
    for count in (0, 1, 3):
        connection = sqlite3.connect(":memory:")
        connection.execute("create table scratch (a integer, b integer, c text)")
        for name in ("a", "b", "c")[:count]:
            connection.execute(f"create index idx_scratch_{name} on scratch({name})")
        data = [(i, (i * 7919) % 10007, f"name{(i * 31) % 50021}") for i in range(rows)]
        started = time.perf_counter()
        connection.executemany("insert into scratch values (?, ?, ?)", data)
        connection.commit()
        print(f"   {rows:,} inserts with {count} index(es): {time.perf_counter() - started:.2f} seconds")
        connection.close()


def main():
    parser = argparse.ArgumentParser(description="Explore indexes on imdb.db.")
    parser.add_argument("--database", default="imdb.db")
    args = parser.parse_args()

    connection = sqlite3.connect(args.database)
    drop_our_indexes(connection)
    connection.close()
    connection = sqlite3.connect(args.database)
    print(f"SQLite version {sqlite3.sqlite_version}")
    print(f"Database size with no extra indexes: {size_mb(connection):,.0f} MB")

    for heading, sql, parameters, create in STEPS:
        print(f"\n== {heading}")
        if create:
            print(f"   {create}", flush=True)
            before = size_mb(connection)
            started = time.perf_counter()
            connection.execute(create)
            connection.commit()
            elapsed = time.perf_counter() - started
            grew = size_mb(connection) - before
            if create.startswith("create"):
                print(f"   built in {elapsed:.1f} seconds; the database grew by {grew:,.0f} MB")
            elif create == "analyze":
                print(f"   analyzed in {elapsed:.1f} seconds")
        print(f"   {' '.join(sql.split())}")
        for line in plan(connection, sql, parameters):
            print(f"   plan: {line}")
        print(f"   time: {best_time(connection, sql, parameters) * 1000:,.1f} ms")
    connection.close()
    write_cost()


if __name__ == "__main__":
    main()

"""The small FruitsForSale table from the SQLite query planner document.

Usage: python3 fruits.py

Builds the seven-row table in memory, adds the indexes one at a time, and prints
the plan SQLite chooses for each query. The chapter's figures show the same
steps. With seven rows the times mean nothing, but the plans are the ones a
large table would get.
"""
import sqlite3

ROWS = [
    (1, "Orange", "FL", 0.85),
    (2, "Apple", "NC", 0.45),
    (4, "Peach", "SC", 0.60),
    (5, "Grape", "CA", 0.80),
    (18, "Lemon", "FL", 1.25),
    (19, "Strawberry", "NC", 2.45),
    (23, "Orange", "CA", 1.05),
]

PEACH = "select price from fruitsforsale where fruit = 'Peach'"
ORANGE = "select price from fruitsforsale where fruit = 'Orange'"
CA_ORANGE = "select price from fruitsforsale where fruit = 'Orange' and state = 'CA'"
EITHER = "select price from fruitsforsale where fruit = 'Orange' or state = 'CA'"
ROWID = "select price from fruitsforsale where rowid = 4"
SORTED = "select * from fruitsforsale order by fruit"
ORANGES_BY_STATE = "select price from fruitsforsale where fruit = 'Orange' order by state"
PARTIAL = "select * from fruitsforsale order by fruit, price"

# (heading, index to create first or None, queries to plan)
STEPS = [
    ("No indexes", None, [PEACH, ROWID, SORTED]),
    ("Idx1 on fruit", "create index Idx1 on fruitsforsale(fruit)", [PEACH, ORANGE, SORTED]),
    ("Idx2 on state", "create index Idx2 on fruitsforsale(state)", [CA_ORANGE]),
    ("Idx3 on fruit, state", "create index Idx3 on fruitsforsale(fruit, state)",
     [CA_ORANGE, PEACH, ORANGES_BY_STATE]),
    ("Idx4 on fruit, state, price (a covering index)",
     "create index Idx4 on fruitsforsale(fruit, state, price)",
     [CA_ORANGE, SORTED, PARTIAL]),
]


def plan(connection, sql):
    return [row[3] for row in connection.execute("explain query plan " + sql)]


def build():
    connection = sqlite3.connect(":memory:")
    connection.execute("create table fruitsforsale (fruit text, state text, price real)")
    connection.executemany(
        "insert into fruitsforsale (rowid, fruit, state, price) values (?, ?, ?, ?)", ROWS)
    return connection


def main():
    connection = build()
    for heading, index, queries in STEPS:
        print(f"== {heading}")
        if index:
            connection.execute(index)
        for sql in queries:
            print(f"   {sql}")
            for line in plan(connection, sql):
                print(f"      {line}")
    connection.close()


if __name__ == "__main__":
    main()

https://docs.google.com/presentation/d/1lYhfjyB6GiZFI8mV64K0Q4Q-jDStWnEVar2s0wG9xpg/edit?usp=sharing

"""Load IMDb title.basics and title.ratings into SQLite, in chunks.

Usage: python3 import_imdb.py [--data-dir data] [--database imdb.db] [--replace]

Reads the compressed files (title.basics.tsv.gz, title.ratings.tsv.gz) directly,
one row at a time, so memory use stays small however large the files are.
Rows go to the database in chunks of 50,000. The tables are created from
schema.sql with no secondary indexes, so the load does not have to maintain
them row by row. (The primary keys are indexed automatically.) The
chapter adds the indexes afterward.
"""
import argparse
import csv
import gzip
from itertools import islice
import os
from pathlib import Path
import sqlite3
import time

HERE = Path(__file__).resolve().parent
CHUNK_SIZE = 50_000

FILES = [
    ("title.basics.tsv.gz", "titles",
     ["text", "text", "text", "text", "int", "int", "int", "int", "text"]),
    ("title.ratings.tsv.gz", "ratings", ["text", "real", "int"]),
]


def convert(value, kind):
    """IMDb writes a missing value as \\N. Convert it to None, and numbers to numbers."""
    if value == "\\N":
        return None
    if kind == "int":
        return int(value)
    if kind == "real":
        return float(value)
    return value


def rows(path, kinds):
    with gzip.open(path, "rt", encoding="utf-8", newline="") as handle:
        reader = csv.reader(handle, delimiter="\t", quoting=csv.QUOTE_NONE)
        next(reader)
        for record in reader:
            yield tuple(convert(value, kind) for value, kind in zip(record, kinds))


def load(connection, path, table, kinds):
    marks = ", ".join("?" for _ in kinds)
    statement = f"insert into {table} values ({marks})"
    source = rows(path, kinds)
    total = 0
    while True:
        chunk = list(islice(source, CHUNK_SIZE))
        if not chunk:
            break
        connection.executemany(statement, chunk)
        connection.commit()
        before = total
        total += len(chunk)
        if total // 1_000_000 > before // 1_000_000:
            print(f"  {table}: {total:,} rows", flush=True)
    print(f"  {table}: {total:,} rows in total")
    return total


def import_all(data_dir, database):
    connection = sqlite3.connect(database)
    try:
        connection.executescript((HERE / "schema.sql").read_text())
        for filename, table, kinds in FILES:
            started = time.perf_counter()
            print(f"Loading {filename}")
            load(connection, Path(data_dir) / filename, table, kinds)
            print(f"  {time.perf_counter() - started:.1f} seconds")
    finally:
        connection.close()


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Load IMDb files into SQLite.")
    parser.add_argument("--data-dir", default="data")
    parser.add_argument("--database", default="imdb.db")
    parser.add_argument("--replace", action="store_true",
                        help="replace the database file if it already exists")
    args = parser.parse_args()
    if Path(args.database).exists():
        if not args.replace:
            raise SystemExit(f"{args.database} already exists. Use --replace to rebuild it.")
        os.remove(args.database)
    import_all(args.data_dir, args.database)

-- Two of the IMDb files, loaded with no secondary indexes yet.
-- tconst is the text identifier IMDb uses for a title, such as tt0000001.
create table titles (
    tconst text primary key,
    title_type text,
    primary_title text,
    original_title text,
    is_adult integer,
    start_year integer,
    end_year integer,
    runtime_minutes integer,
    genres text
);

create table ratings (
    tconst text primary key,
    average_rating real,
    num_votes integer
);

-- A stored procedure, for PostgreSQL. SQLite has no CREATE PROCEDURE.
-- Run with:  psql -d yourdatabase -f stored_procedure.sql

drop table if exists adoption_log;
drop table if exists pets;

create table pets (
    id serial primary key,
    name text not null,
    owner text not null
);

create table adoption_log (
    id serial primary key,
    pet_id integer not null references pets(id),
    old_owner text not null,
    new_owner text not null,
    changed_at timestamptz not null default now()
);

insert into pets (name, owner) values ('Buddy', 'Kim'), ('Whiskers', 'Lee');

-- One call from the application replaces a read, an update, and an insert.
create or replace procedure change_owner(pet integer, new_owner text)
language plpgsql
as $$
declare
    previous text;
begin
    select owner into previous from pets where id = pet for update;
    if not found then
        raise exception 'No pet with id %', pet;
    end if;
    update pets set owner = new_owner where id = pet;
    insert into adoption_log (pet_id, old_owner, new_owner)
    values (pet, previous, new_owner);
end;
$$;

call change_owner(1, 'Sam');

select p.name, p.owner, l.old_owner, l.new_owner
from pets p join adoption_log l on l.pet_id = p.id;

"""Run with python3 -m unittest -v. Checks the behavior the chapter describes.

Every test uses temporary files and a small made-up file in the IMDb format,
so nothing here needs the real IMDb download.
"""
import contextlib
import gzip
import io
from pathlib import Path
import random
import sqlite3
import subprocess
import sys
import tempfile
import unittest

import compare_python_sql
import explore_indexes
import fruits
import import_imdb

HERE = Path(__file__).resolve().parent
TYPES = ["movie", "short", "tvEpisode", "tvSeries"]


def write_gz(path, text):
    with gzip.open(path, "wt", encoding="utf-8", newline="") as handle:
        handle.write(text)


def make_imdb_files(directory, count=3000):
    """Write small title.basics and title.ratings files with IMDb's layout."""
    generator = random.Random(1)
    basics = ["tconst\ttitleType\tprimaryTitle\toriginalTitle\tisAdult\tstartYear\t"
              "endYear\truntimeMinutes\tgenres"]
    ratings = ["tconst\taverageRating\tnumVotes"]
    for number in range(1, count + 1):
        tconst = f"tt{number:07d}"
        kind = generator.choice(TYPES)
        year = generator.choice([r"\N"] + [str(y) for y in range(1990, 2000)])
        title = f"Title {generator.randrange(500)}"
        basics.append(f"{tconst}\t{kind}\t{title}\t{title}\t0\t{year}\t\\N\t90\tDrama")
        if number % 2 == 0:
            ratings.append(f"{tconst}\t{generator.randrange(10, 100) / 10}\t"
                           f"{generator.randrange(5, 5000)}")
    basics.append("tt0068646\tmovie\tThe Godfather\tThe Godfather\t0\t1972\t\\N\t175\tCrime,Drama")
    ratings.append("tt0068646\t9.2\t2000000")
    directory = Path(directory)
    write_gz(directory / "title.basics.tsv.gz", "\n".join(basics) + "\n")
    write_gz(directory / "title.ratings.tsv.gz", "\n".join(ratings) + "\n")


class FruitsTests(unittest.TestCase):
    def plans(self):
        connection = fruits.build()
        self.addCleanup(connection.close)
        found = {}
        for heading, index, queries in fruits.STEPS:
            if index:
                connection.execute(index)
            for sql in queries:
                found[(heading, sql)] = " | ".join(fruits.plan(connection, sql))
        return found

    def test_the_seven_rows_are_there(self):
        connection = fruits.build()
        self.addCleanup(connection.close)
        self.assertEqual(connection.execute("select count(*) from fruitsforsale").fetchone()[0], 7)

    def test_each_index_changes_the_plan_as_the_chapter_says(self):
        plans = self.plans()
        self.assertIn("SCAN", plans[("No indexes", fruits.PEACH)])
        self.assertIn("INTEGER PRIMARY KEY", plans[("No indexes", fruits.ROWID)])
        self.assertIn("TEMP B-TREE FOR ORDER BY", plans[("No indexes", fruits.SORTED)])
        self.assertIn("SEARCH fruitsforsale USING INDEX Idx1",
                      plans[("Idx1 on fruit", fruits.PEACH)])
        self.assertNotIn("TEMP B-TREE", plans[("Idx1 on fruit", fruits.SORTED)])
        self.assertIn("USING INDEX Idx2 (state=?)", plans[("Idx2 on state", fruits.CA_ORANGE)])
        self.assertIn("Idx3 (fruit=? AND state=?)",
                      plans[("Idx3 on fruit, state", fruits.CA_ORANGE)])
        self.assertNotIn("TEMP B-TREE",
                         plans[("Idx3 on fruit, state", fruits.ORANGES_BY_STATE)])
        covering = "Idx4 on fruit, state, price (a covering index)"
        self.assertIn("COVERING INDEX Idx4", plans[(covering, fruits.CA_ORANGE)])
        self.assertIn("LAST TERM OF ORDER BY", plans[(covering, fruits.PARTIAL)])

    def test_indexes_never_change_the_answer(self):
        connection = fruits.build()
        self.addCleanup(connection.close)
        before = sorted(connection.execute(fruits.EITHER).fetchall())
        for _, index, _ in fruits.STEPS:
            if index:
                connection.execute(index)
        self.assertEqual(sorted(connection.execute(fruits.EITHER).fetchall()), before)
        self.assertEqual(before, [(0.8,), (0.85,), (1.05,)])


class ImdbTests(unittest.TestCase):
    def setUp(self):
        self.temp = tempfile.TemporaryDirectory()
        self.addCleanup(self.temp.cleanup)
        self.directory = Path(self.temp.name)
        make_imdb_files(self.directory)
        self.database = str(self.directory / "imdb.db")
        with contextlib.redirect_stdout(io.StringIO()):
            import_imdb.import_all(self.directory, self.database)
        self.connection = sqlite3.connect(self.database)
        self.addCleanup(self.connection.close)

    def one(self, sql, *parameters):
        return self.connection.execute(sql, parameters).fetchone()[0]

    def test_all_rows_arrive(self):
        self.assertEqual(self.one("select count(*) from titles"), 3001)
        self.assertEqual(self.one("select count(*) from ratings"), 1501)

    def test_missing_values_become_null_and_numbers_become_numbers(self):
        self.assertGreater(self.one("select count(*) from titles where start_year is null"), 0)
        self.assertEqual(self.one("select count(*) from titles where start_year = '\\N'"), 0)
        self.assertEqual(self.one("select typeof(start_year) from titles where tconst = 'tt0068646'"),
                         "integer")
        self.assertEqual(self.one("select typeof(average_rating) from ratings where tconst = 'tt0068646'"),
                         "real")

    def test_loading_uses_more_than_one_chunk(self):
        original = import_imdb.CHUNK_SIZE
        import_imdb.CHUNK_SIZE = 700
        self.addCleanup(setattr, import_imdb, "CHUNK_SIZE", original)
        other = str(self.directory / "chunked.db")
        with contextlib.redirect_stdout(io.StringIO()):
            import_imdb.import_all(self.directory, other)
        connection = sqlite3.connect(other)
        self.addCleanup(connection.close)
        self.assertEqual(connection.execute("select count(*) from titles").fetchone()[0], 3001)

    def test_the_importer_will_not_overwrite_a_database(self):
        result = subprocess.run(
            [sys.executable, str(HERE / "import_imdb.py"), "--data-dir", str(self.directory),
             "--database", self.database], capture_output=True, text=True)
        self.assertNotEqual(result.returncode, 0)
        self.assertIn("already exists", result.stderr)

    def plan(self, sql, parameters=()):
        return " | ".join(explore_indexes.plan(self.connection, sql, parameters))

    def test_the_steps_in_the_chapter_produce_the_plans_it_describes(self):
        seen = {}
        for heading, sql, parameters, create in explore_indexes.STEPS:
            if create:
                self.connection.execute(create)
            seen[heading] = self.plan(sql, parameters)
        self.assertIn("sqlite_autoindex_titles_1", seen["Lookup by the primary key (tconst)"])
        self.assertTrue(seen["Lookup by a column with no index"].startswith("SCAN titles"))
        self.assertIn("SEARCH titles USING INDEX idx_titles_primary_title",
                      seen["The same query after an index on primary_title"])
        self.assertTrue(seen["LIKE with a prefix, against a case-sensitive index"].startswith("SCAN titles"))
        self.assertIn("idx_titles_primary_title_nocase",
                      seen["The same LIKE after an index that ignores case"])
        self.assertTrue(seen["LIKE with a wildcard at the front"].startswith("SCAN titles"))
        self.assertIn("idx_titles_year", seen["The same query after an index on start_year"])
        self.assertIn("COVERING INDEX idx_titles_type_year (title_type=? AND start_year=?)",
                      seen["The same query after a two-column index"])
        self.assertIn("SCAN titles USING COVERING INDEX idx_titles_type_year",
                      seen["The second column alone, with the single-column index gone"])
        self.assertIn("SEARCH titles USING INDEX idx_titles_type_year ",
                      seen["Returning a column that is not in the index"])
        self.assertIn("COVERING INDEX idx_titles_type_year_title",
                      seen["The same query after a covering index"])
        self.assertTrue(seen["Two OR-connected conditions, one of them with no index"]
                        .startswith("SCAN titles"))
        self.assertIn("MULTI-INDEX OR", seen["Two OR-connected conditions, each with an index"])
        self.assertIn("USE TEMP B-TREE FOR ORDER BY",
                      seen["The most-voted titles, with no index on num_votes"])
        self.assertNotIn("TEMP B-TREE", seen["The same query after an index on num_votes"])
        self.assertNotIn("SEARCH", seen["The second column alone, before ANALYZE"])
        self.assertIn("ANY(title_type)", seen["The second column alone, after ANALYZE"])

    def test_a_query_gives_the_same_rows_with_and_without_an_index(self):
        sql = "select tconst from titles where primary_title = ? order by tconst"
        before = self.connection.execute(sql, ("Title 7",)).fetchall()
        self.connection.execute("create index idx_titles_primary_title on titles(primary_title)")
        self.assertEqual(self.connection.execute(sql, ("Title 7",)).fetchall(), before)
        self.assertGreater(len(before), 0)

    def test_dropping_our_indexes_leaves_the_primary_key_index(self):
        self.connection.execute("create index idx_titles_year on titles(start_year)")
        self.connection.execute("analyze")
        explore_indexes.drop_our_indexes(self.connection)
        self.assertEqual(self.one("select count(*) from sqlite_stat1"), 0)
        names = [row[0] for row in self.connection.execute(
            "select name from sqlite_master where type = 'index'")]
        self.assertNotIn("idx_titles_year", names)
        self.assertIn("sqlite_autoindex_titles_1", names)

    def test_python_and_sql_give_the_same_answers(self):
        for statement in compare_python_sql.INDEXES:
            self.connection.execute(statement)
        ids = [row[0] for row in self.connection.execute(
            "select tconst from ratings order by num_votes desc limit 50")]
        self.assertEqual(compare_python_sql.one_query_per_title(self.connection, ids),
                         compare_python_sql.one_join(self.connection, 50))
        self.assertEqual(compare_python_sql.movies_per_decade_in_python(self.connection),
                         compare_python_sql.movies_per_decade_in_sql(self.connection))


if __name__ == "__main__":
    unittest.main()

"""Time one query before and after adding an index, by hand.

Usage: python3 try_an_index.py [--database imdb.db]

A small, readable version of the first step in explore_indexes.py: look up a
title by primary_title, add an index on that column, and look it up again.
Prints the plan and the time for both, and checks that the index did not
change the answer. Leaves idx_titles_primary_title in place afterward.
"""
import argparse
from time import perf_counter
import sqlite3

QUERY = "select tconst, start_year from titles where primary_title = ?"
PARAMETERS = ("The Godfather",)


def plan(connection, sql, parameters=()):
    for row in connection.execute("explain query plan " + sql, parameters):
        print(row[3])


def timed(connection, sql, parameters=()):
    started = perf_counter()
    rows = connection.execute(sql, parameters).fetchall()
    return rows, (perf_counter() - started) * 1000


def main(database):
    connection = sqlite3.connect(f"file:{database}?mode=rw", uri=True)
    connection.execute("drop index if exists idx_titles_primary_title")

    plan(connection, QUERY, PARAMETERS)
    before, before_ms = timed(connection, QUERY, PARAMETERS)
    print("Before:", before_ms, "ms")

    connection.execute("create index idx_titles_primary_title on titles(primary_title)")
    connection.commit()

    plan(connection, QUERY, PARAMETERS)
    after, after_ms = timed(connection, QUERY, PARAMETERS)
    print("After:", after_ms, "ms")

    print("Same rows:", sorted(before) == sorted(after))
    connection.close()


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Time a query before and after an index.")
    parser.add_argument("--database", default="imdb.db")
    args = parser.parse_args()
    main(args.database)