Merge pull request #27 from eyeseast/rowid-bug

Fix error with rowid tables

Author: eyeseast

Diff for: geocode_sqlite/cli.py

@@ -159,13 +159,13 @@ def geocode(ctx, geocoder):
     )
 
     with click.progressbar(gen, length=count, label=f"{count} rows") as bar:
-        for row, success in bar:
-            pks = [row[pk] for pk in table.pks]
+        for pk, row, success in bar:
+            # pks = [row[pk] for pk in table.pks]
             if success:
-                table.update(pks, row)
+                table.update(pk, row)
                 done += 1
             else:
-                errors.append(pks)
+                errors.append(pk)
 
     click.echo("Geocoded {} rows".format(done))
     if errors:

Diff for: geocode_sqlite/utils.py

@@ -63,12 +63,12 @@ def geocode_table(
 
     count = 0
     log.info(f"Geocoding {todo} rows from {table.name}")
-    for row in rows:
+    for pk, row in rows:
         result = geocode_row(geocode, query_template, row, **kwargs)
         if result:
-            pks = [row[pk] for pk in table.pks]
+            # pks = [row[pk] for pk in table.pks]
             table.update(
-                pks,
+                pk,
                 {
                     latitude_column: result.latitude,
                     longitude_column: result.longitude,
@@ -95,20 +95,21 @@ def geocode_list(
     """
     Geocode an arbitrary list of rows, returning a generator.
     This does not query or save geocoded results into a table.
-    If geocoding succeeds, it will yield a two-tuple:
+    If geocoding succeeds, it will yield a three-tuple:
+     - the primary key of the row (rowid or actual PK)
      - the row with latitude and longitude columns set
      - and True
 
     If geocoding fails, it will yield the original row and False.
     """
-    for row in rows:
+    for pk, row in rows:
         result = geocode_row(geocode, query_template, row, **kwargs)
         if result:
             row[longitude_column] = result.longitude
             row[latitude_column] = result.latitude
             row["geocoder"] = get_geocoder_class(geocode)
 
-        yield row, bool(result)
+        yield pk, row, bool(result)
 
 
 def geocode_row(geocode, query_template, row, **kwargs):
@@ -135,7 +136,9 @@ def select_ungeocoded(
     if count:
         count = count[0]
 
-    rows = table.rows_where(f"{latitude_column} IS NULL OR {longitude_column} IS NULL")
+    rows = table.pks_and_rows_where(
+        f"{latitude_column} IS NULL OR {longitude_column} IS NULL"
+    )
 
     return rows, count
 

Diff for: tests/test_geocode_sqlite.py

@@ -24,12 +24,14 @@
 
 
 @pytest.fixture
-def db():
+def db(request):
     db = Database(DB_PATH)
     table = db[TABLE_NAME]
 
+    pk = getattr(request, "param", "id")
+
     # load csv data, which will be geocoded
-    table.insert_all(csv.DictReader(open(CSV_DATA)), alter=True, pk="id")
+    table.insert_all(csv.DictReader(open(CSV_DATA)), alter=True, pk=pk)
 
     # load our geojson data, for our fake geocoder
     fc = json.load(open(GEOJSON_DATA))
@@ -56,6 +58,7 @@ def test_version():
         assert result.output.startswith("cli, version ")
 
 
+@pytest.mark.parametrize("db", ["id", None], indirect=True)
 def test_cli_geocode_table(db, geocoder):
     runner = CliRunner()
     table = db[TABLE_NAME]
@@ -243,7 +246,7 @@ def test_resume_table(db, geocoder):
 def test_geocode_list(db, geocoder):
     table = db[TABLE_NAME]
 
-    utah = list(table.rows_where('"state" = "UT"'))
+    utah = list(table.pks_and_rows_where('"state" = "UT"'))
     assert len(utah) == 10
 
     gen = geocode_list(utah, geocoder.geocode, "{id}")
@@ -255,9 +258,9 @@ def test_geocode_list(db, geocoder):
     # geocode the whole table, to cheeck results
     geocode_table(db, TABLE_NAME, geocoder, "{id}")
 
-    for row, success in done:
+    for pk, row, success in done:
         assert success
-        assert row == table.get(row["id"])
+        assert row == table.get(pk)
 
 
 def test_label_results(db, geocoder):