from threading import Event
import time
from cassandra.query import SimpleStatement
from swh.storage import get_storage
storage = get_storage('cassandra', {
'hosts': [
'128.93.66.190',
'128.93.66.191',
'128.93.66.187',
'128.93.64.42',
],
'keyspace': 'swh_test',
'objstorage': {
'cls': 'memory',
'args': {},
},
})
session = storage._proxy._session
class PagedResultHandler(object):
def __init__(self, query, callback):
self.callback = callback
self.total_rows = 0
self.error = None
self.finished_event = Event()
statement = SimpleStatement(query, fetch_size=10000)
self.start_time = time.time()
self.future = session.execute_async(statement)
self.future.add_callbacks(
callback=self.handle_page,
errback=self.handle_error)
def handle_page(self, rows):
if self.future.has_more_pages:
self.future.start_fetching_next_page()
self.total_rows += len(rows)
print('total rows: %s (%dk/s)' % (
self.total_rows,
self.total_rows/(time.time()-self.start_time)/1000))
for row in rows:
self.callback(row)
if not self.future.has_more_pages:
self.finished_event.set()
def handle_error(self, exc):
print('%s' % exc)
self.finished_event.set()
with open('nodes.csv', 'w') as nodes_fd, open('edges.csv', 'w') as edges_fd:
def revision_callback(row):
rev_id = 'swh:1:rev:%s' % row.id.hex()
nodes_fd.write(rev_id + '\n')
edges_fd.write('%s swh:1:dir:%s\n' % (rev_id, row.directory.hex()))
for parent in row.parents:
edges_fd.write('%s swh:1:rev:%s\n' % (rev_id, parent.hex()))
handler = PagedResultHandler("SELECT *token(id) as tok, id, directory, parents FROM revision", revision_callback)
handler.finished_event.wait()