[soundwave] Add pandas_sqlite.InsertOrReplaceRecords

Pandas DataFrame.to_sql method has limitation of not being able to
"insert or replace" records, see e.g:
pandas-dev/pandas#14553

Using pandas.io.sql primitives, however, it's not too hard to implement
such a functionality (for the SQLite case only).

Assuming that index columns of the frame have names, this method will
use those columns as the PRIMARY KEY of the table.

Bug: catapult:#4382
Change-Id: I5f70a04c18b998590b1e77e5ff6b89b2d27138af
Reviewed-on: https://chromium-review.googlesource.com/1035266
Commit-Queue: Juan Antonio Navarro Pérez <[email protected]>
Reviewed-by: Charlie Andrews <[email protected]>

Author: perezju

experimental/soundwave/soundwave/commands.py

@@ -6,6 +6,7 @@
 import sqlite3
 
 from soundwave import dashboard_api
+from soundwave import pandas_sqlite
 from soundwave import tables
 
 
@@ -16,19 +17,13 @@ def FetchAlertsData(args):
     alerts = tables.alerts.DataFrameFromJson(
         api.GetAlertData(args.benchmark, args.days))
     print '%d alerts found!' % len(alerts)
-    # TODO: Make this update rather than replace the existing table.
-    # Note that if_exists='append' does not work since there is no way to
-    # specify in pandas' |to_sql| a primary key or, more generally, uniqueness
-    # constraints on columns. So this would lead to duplicate entries for
-    # alerts with the same |key|.
-    alerts.to_sql('alerts', conn, if_exists='replace')
+    pandas_sqlite.InsertOrReplaceRecords(alerts, 'alerts', conn)
 
     bug_ids = set(alerts['bug_id'].unique())
     bug_ids.discard(0)  # A bug_id of 0 means untriaged.
     print '%d bugs found!' % len(bug_ids)
     bugs = tables.bugs.DataFrameFromApi(api, bug_ids)
-    # TODO: Ditto. Make this update rather than replace the existing table.
-    bugs.to_sql('bugs', conn, if_exists='replace')
+    pandas_sqlite.InsertOrReplaceRecords(bugs, 'bugs', conn)
   finally:
     conn.close()
 

experimental/soundwave/soundwave/pandas_sqlite.py

@@ -0,0 +1,43 @@
+# Copyright 2018 The Chromium Authors. All rights reserved.
+# Use of this source code is governed by a BSD-style license that can be
+# found in the LICENSE file.
+
+"""
+Helper methods for dealing with a SQLite database with pandas.
+"""
+import pandas.io.sql  # pylint: disable=import-error
+
+
+def _InsertOrReplaceStatement(name, keys):
+  columns = ','.join(keys)
+  values = ','.join('?' for _ in keys)
+  return 'INSERT OR REPLACE INTO %s(%s) VALUES (%s)' % (name, columns, values)
+
+
+def InsertOrReplaceRecords(frame, name, conn):
+  """Insert or replace records from a DataFrame into a SQLite database.
+
+  Assumes that index columns of the frame have names, and those are used as to
+  set the PRIMARY KEY of the table when creating anew. If the table already
+  exists, any new records with a matching PRIMARY KEY will replace existing
+  records.
+
+  Args:
+    frame: DataFrame with records to write.
+    name: Name of SQL table.
+    conn: A sqlite connection object.
+  """
+  db = pandas.io.sql.SQLiteDatabase(conn)
+  if db.has_table(name):
+    table = pandas.io.sql.SQLiteTable(
+        name, db, frame=frame, index=True, if_exists='append')
+    keys, data = table.insert_data()
+    insert_statement = _InsertOrReplaceStatement(name, keys)
+    with db.run_transaction() as c:
+      c.executemany(insert_statement, zip(*data))
+  else:
+    table = pandas.io.sql.SQLiteTable(
+        name, db, frame=frame, index=True, keys=frame.index.names,
+        if_exists='fail')
+    table.create()
+    table.insert()

experimental/soundwave/soundwave/pandas_sqlite_test.py

@@ -0,0 +1,54 @@
+# Copyright 2018 The Chromium Authors. All rights reserved.
+# Use of this source code is governed by a BSD-style license that can be
+# found in the LICENSE file.
+
+import pandas  # pylint: disable=import-error
+import sqlite3
+import unittest
+
+from soundwave import pandas_sqlite
+
+
+class TestPandasSQLite(unittest.TestCase):
+  def testInsertOrReplaceRecords_newTable(self):
+    columns = ('bug_id', 'summary', 'status')
+    df1 = pandas.DataFrame.from_records(
+        [(123, 'Some bug', 'Started'), (456, 'Another bug', 'Assigned')],
+        columns=columns, index=columns[0])
+    conn = sqlite3.connect(':memory:')
+    try:
+      # Write new table to database, read back and check they are equal.
+      pandas_sqlite.InsertOrReplaceRecords(df1, 'bugs', conn)
+      df = pandas.read_sql('SELECT * FROM bugs', conn, index_col=columns[0])
+      self.assertTrue(df.equals(df1))
+    finally:
+      conn.close()
+
+  def testInsertOrReplaceRecords_existingTable(self):
+    columns = ('bug_id', 'summary', 'status')
+    df1 = pandas.DataFrame.from_records(
+        [(123, 'Some bug', 'Started'), (456, 'Another bug', 'Assigned')],
+        columns=columns, index=columns[0])
+    df2 = pandas.DataFrame.from_records(
+        [(123, 'Some bug', 'Fixed'), (789, 'A new bug', 'Untriaged')],
+        columns=columns, index=columns[0])
+    conn = sqlite3.connect(':memory:')
+    try:
+      # Write first data frame to database.
+      pandas_sqlite.InsertOrReplaceRecords(df1, 'bugs', conn)
+      df = pandas.read_sql('SELECT * FROM bugs', conn, index_col=columns[0])
+      self.assertEqual(len(df), 2)
+      self.assertEqual(df.loc[123]['status'], 'Started')
+
+      # Write second data frame to database.
+      pandas_sqlite.InsertOrReplaceRecords(df2, 'bugs', conn)
+      df = pandas.read_sql('SELECT * FROM bugs', conn, index_col=columns[0])
+      self.assertEqual(len(df), 3)  # Only one extra record added.
+      self.assertEqual(df.loc[123]['status'], 'Fixed')  # Bug is now fixed.
+      self.assertItemsEqual(df.index, (123, 456, 789))
+    finally:
+      conn.close()
+
+
+if __name__ == '__main__':
+  unittest.main()

experimental/soundwave/soundwave/tables/bugs_test.py

@@ -34,7 +34,7 @@ def testDataFrameFromApi(self):
     bugs = tables.bugs.DataFrameFromApi(api, [12345])
     self.assertEqual(len(bugs), 1)
 
-    bug = bugs.loc[12345]
+    bug = bugs.loc[12345]  # Get bug by id.
     self.assertEqual(bug['published'], datetime.datetime(
         year=2018, month=4, day=9, hour=17, minute=1, second=9))
     self.assertEqual(bug['status'], 'Fixed')