def testInsertOrReplaceRecords_existingRecords(self):
column_types = (('bug_id', int), ('summary', str), ('status', str))
rows1 = [(123, 'Some bug', 'Started'),
(456, 'Another bug', 'Assigned')]
df1 = pandas_sqlite.DataFrame(column_types, index='bug_id', rows=rows1)
rows2 = [(123, 'Some bug', 'Fixed'), (789, 'A new bug', 'Untriaged')]
df2 = pandas_sqlite.DataFrame(column_types, index='bug_id', rows=rows2)
con = sqlite3.connect(':memory:')
try:
pandas_sqlite.CreateTableIfNotExists(con, 'bugs', df1)
# Write first data frame to database.
pandas_sqlite.InsertOrReplaceRecords(con, 'bugs', df1)
df = pandas.read_sql('SELECT * FROM bugs', con, index_col='bug_id')
self.assertEqual(len(df), 2)
self.assertEqual(df.loc[123]['status'], 'Started')
# Write second data frame to database.
pandas_sqlite.InsertOrReplaceRecords(con, 'bugs', df2)
df = pandas.read_sql('SELECT * FROM bugs', con, index_col='bug_id')
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:
con.close()
def Get(con, bug_id):
"""Find the record for a bug_id in the given database connection.
Returns:
A pandas.Series with the record if found, or None otherwise.
"""
df = pandas.read_sql('SELECT * FROM %s WHERE id=?' % TABLE_NAME,
con,
params=(bug_id, ),
index_col=INDEX,
parse_dates=DATE_COLUMNS)
return df.loc[bug_id] if len(df) else None
def GetTimeSeries(con, test_path, extra_cond=None):
"""Get the records for all data points on the given test_path.
Returns:
A pandas.DataFrame with all records found.
"""
config = _ParseConfigFromTestPath(test_path)
params = tuple(config[c] for c in INDEX[:-1])
query = _QUERY_TIME_SERIES
if extra_cond is not None:
query = ' '.join([query, extra_cond])
return pandas.read_sql(query,
con,
params=params,
parse_dates=['timestamp'])
def testWorkerPoolRun(self):
tempdir = tempfile.mkdtemp()
try:
args = argparse.Namespace()
args.database_file = os.path.join(tempdir, 'test.db')
args.processes = 3
schema = pandas_sqlite.DataFrame([('item', int)])
items = range(20) # We'll write these in the database.
con = sqlite3.connect(args.database_file)
try:
pandas_sqlite.CreateTableIfNotExists(con, 'items', schema)
with open(os.devnull, 'w') as devnull:
worker_pool.Run('Processing:',
TestWorker,
args,
items,
stream=devnull)
df = pandas.read_sql('SELECT * FROM items', con)
# Check all of our items were written.
self.assertItemsEqual(df['item'], items)
finally:
con.close()
finally:
shutil.rmtree(tempdir)