#!/usr/bin/env python3
# Copyright (C) 2021 The Software Heritage developers
# See the AUTHORS file at the top-level directory of this distribution
# License: GNU General Public License version 3, or any later version
# See top-level LICENSE file for more information
import functools
import psycopg2
import psycopg2.extras
from swh.journal.client import get_journal_client
total_revisions_processed = 0
def process_journal_objects(messages, *, conn):
global total_revisions_processed
assert set(messages) == {"revision"}, set(messages)
revisions = messages["revision"]
rows = []
for revision in revisions:
if revision.get("extra_headers"):
rows.append((revision["id"], revision["extra_headers"]))
else:
eh = (revision.get("metadata") or {}).get("extra_headers")
if eh:
print(eh)
elif revision["type"] != "tar" and revision.get("metadata"):
print(revision["metadata"])
cur = conn.cursor()
psycopg2.extras.execute_values(cur, """
select ids.id from (VALUES %s) AS ids (id, eh)
inner join revision on (ids.id=revision.id)
where
(revision.extra_headers = ARRAY[]::bytea[]
OR revision.extra_headers IS NULL)""", rows)
try:
for res in cur:
id_ = bytes(res[0])
print("update needed:", id_.hex())
print([row[1] for row in rows if row[0] == id_])
except psycopg2.ProgrammingError:
print("update needed:", res[0].hex())pass
psycopg2.extras.execute_values(
cur,
"""
UPDATE revision
SET extra_headers = data.extra_headers
FROM (VALUES %s) AS data (id, extra_headers)
WHERE
revision.id=data.id
AND (
-- Don't unnecessarily update rows that already have their
-- 'extra_headers' cell populated
revision.extra_headers = ARRAY[]::bytea[]
OR revision.extra_headers IS NULL
)
""",
rows,
)
total_revisions_processed += len(revisions)
print(f"processed {len(revisions)} revisions,\t{len(rows)} updates")
def main():
client = get_journal_client(
cls="kafka",
prefix="swh.journal.objects",
object_types=["revision"],
brokers=[f"kafka{i+1}.internal.softwareheritage.org:9092" for i in range(4)],
group_id="vlorentz-T2564-migrate-extra-headers",
#batch_size=1000,
#batch_size=10000, # Sends mostly batches of 1000 to 3000 rows to postgresql
)
conn = psycopg2.connect(<redacted>)
worker_fn = functools.partial(process_journal_objects, conn=conn)
try:
client.process(worker_fn)
except KeyboardInterrupt:
pass
finally:
print(f"Total revisions processed: {total_revisions_processed}")
if __name__ == "__main__":
main()