Task 5. Nạp metadata nguồn vào MongoDB bằng Spark

import json
import os
import re
import subprocess
import sys
import uuid
from datetime import datetime, timezone
from pathlib import Path

PROJECT_ROOT = Path.cwd()
if PROJECT_ROOT.name == 'lab04-book':
    PROJECT_ROOT = PROJECT_ROOT.parent
SRC_ROOT = PROJECT_ROOT / 'src'
if str(SRC_ROOT) not in sys.path:
    sys.path.insert(0, str(SRC_ROOT))
os.chdir(PROJECT_ROOT)

def load_local_env(path: Path) -> None:
    if not path.exists():
        raise FileNotFoundError('Missing .env. Copy .env.example to .env first.')
    for raw_line in path.read_text(encoding='utf-8').splitlines():
        line = raw_line.strip()
        if not line or line.startswith('#') or '=' not in line:
            continue
        name, value = line.split('=', 1)
        os.environ.setdefault(name.strip(), value.strip().strip('\"').strip("'"))

load_local_env(PROJECT_ROOT / '.env')
RUN_ID = os.environ.get('TASK5_NOTEBOOK_RUN_ID') or datetime.now(timezone.utc).strftime('%Y%m%dT%H%M%SZ') + '-' + uuid.uuid4().hex[:8]
RUN_ROOT = PROJECT_ROOT / 'workspace' / 'tmp' / 'task5-notebook' / RUN_ID
CHECKPOINT_DIR = PROJECT_ROOT / 'workspace' / 'checkpoints' / 'task5-notebook' / RUN_ID
if RUN_ROOT.exists() or CHECKPOINT_DIR.exists():
    raise RuntimeError(f'Run ID already exists: {RUN_ID}. Choose a new TASK5_NOTEBOOK_RUN_ID.')

KAFKA_BOOTSTRAP = os.environ['KAFKA_BOOTSTRAP_SERVERS']
MONGO_URI = os.environ['MONGODB_URI']
MONGO_DATABASE = os.environ['MONGODB_DATABASE']
MONGO_COLLECTION = os.environ['MONGODB_COLLECTION']
MASKED_MONGO_URI = re.sub(r'//([^:]+):[^@]+@', r'//\1:***@', MONGO_URI)
print(f'RUN_ID={RUN_ID}')
print('KAFKA_TOPOLOGY=KRaft (no ZooKeeper)')
print(f'KAFKA_BOOTSTRAP_SERVERS={KAFKA_BOOTSTRAP}')
print(f'MONGODB_URI={MASKED_MONGO_URI}')
print(f'CHECKPOINT_DIR={CHECKPOINT_DIR.relative_to(PROJECT_ROOT).as_posix()}')

compose = subprocess.run(
    ['docker', 'compose', '--env-file', '.env', '-f', 'infra/docker-compose.yml', 'ps'],
    capture_output=True, text=True, check=True, cwd=PROJECT_ROOT,
)
print(compose.stdout)
assert 'cpg-kafka' in compose.stdout and 'cpg-mongodb' in compose.stdout

required_topics = [os.environ['TOPIC_NODES'], os.environ['TOPIC_EDGES'], os.environ['TOPIC_METADATA'], os.environ['TOPIC_ERRORS']]
for topic_name in required_topics:
    subprocess.run(
        ['docker', 'exec', 'cpg-kafka', 'kafka-topics', '--bootstrap-server', 'kafka:29092', '--create', '--if-not-exists', '--topic', topic_name, '--partitions', '1', '--replication-factor', '1'],
        capture_output=True, text=True, check=True,
    )
topic = subprocess.run(
    ['docker', 'exec', 'cpg-kafka', 'kafka-topics', '--bootstrap-server', 'kafka:29092', '--describe', '--topic', os.environ['TOPIC_METADATA']],
    capture_output=True, text=True, check=True,
)
print(topic.stdout)
assert 'PartitionCount: 1' in topic.stdout

mongo_indexes = subprocess.run(
    ['docker', 'exec', 'cpg-mongodb', 'mongosh', '--quiet', '--username', os.environ['MONGO_ROOT_USERNAME'], '--password', os.environ['MONGO_ROOT_PASSWORD'], '--authenticationDatabase', 'admin', '--eval', "db=db.getSiblingDB('cpg_metadata'); printjson(db.file_statistics.getIndexes())"],
    capture_output=True, text=True, check=True,
)
print(mongo_indexes.stdout)
assert "name: 'file_id_1'" in mongo_indexes.stdout
assert "name: 'repository_id_1_file_path_1'" in mongo_indexes.stdout
assert mongo_indexes.stdout.count('unique: true') >= 2

from application.services.process_file import ProcessFileService
from domain.models import SourceFile
from infrastructure.filesystem.git_source_repository import GitSourceRepository
from infrastructure.messaging.event_validator import EventValidator
from infrastructure.messaging.kafka_producer import KafkaEventProducer
from infrastructure.state.sqlite_state_store import SqliteStateStore
from parsing.cpg_parser import CpgParser
from parsing.identifiers import IdentifierGenerator

source_root = RUN_ROOT / 'source'
state_db = RUN_ROOT / 'state.sqlite3'
target_file = Path('src/task5_runtime.py')
repository_id = f'notebook/task5/{RUN_ID}'
initial_source = b'def normalize(value):\n    total = value + 1\n    return total\n'
modified_source = b'def normalize(value):\n    total = value + 1\n    if total > 10:\n        return total * 2\n    return total\n\nclass Normalizer:\n    def apply(self, value):\n        return normalize(value)\n'
source_path = source_root / target_file
source_path.parent.mkdir(parents=True, exist_ok=True)
source_path.write_bytes(initial_source)
repo_adapter = GitSourceRepository(source_root, '', None)
validator = EventValidator(PROJECT_ROOT / 'schemas')
state_store = SqliteStateStore(state_db, repository_id=repository_id)
producer = KafkaEventProducer(bootstrap_servers=KAFKA_BOOTSTRAP)
parser = CpgParser(repository_id=repository_id)
process_service = ProcessFileService(
    repo_adapter, parser, state_store, validator, producer,
    topic_nodes=os.environ['TOPIC_NODES'],
    topic_edges=os.environ['TOPIC_EDGES'],
    topic_metadata=os.environ['TOPIC_METADATA'],
    topic_errors=os.environ['TOPIC_ERRORS'],
)

def make_source_file() -> SourceFile:
    return SourceFile(repository_id, str(source_root), target_file.as_posix(), f'notebook-{RUN_ID}', source_path.stat().st_size)

def kafka_latest_offset() -> int:
    result = subprocess.run(['docker', 'exec', 'cpg-kafka', 'kafka-run-class', 'kafka.tools.GetOffsetShell', '--broker-list', 'kafka:29092', '--topic', 'source.metadata', '--time', '-1'], capture_output=True, text=True, check=True)
    return int(result.stdout.strip().rsplit(':', 1)[1])

baseline_latest_offset = kafka_latest_offset()
first_result = process_service.execute(make_source_file())
file_id = first_result.file_id
first_content_hash = first_result.content_hash
assert first_result.status.value == 'SUCCESS'
assert first_result.emitted_event_counts.get('source.metadata') == 1
print(json.dumps({'repository_id': repository_id, 'file_id': file_id, 'content_hash': first_content_hash, 'parse_status': 'SUCCESS', 'event_offset': baseline_latest_offset, 'emitted_event_counts': first_result.emitted_event_counts}, indent=2))

def read_metadata_record(offset: int) -> tuple[str, dict]:
    result = subprocess.run(
        ['docker', 'exec', 'cpg-kafka', 'kafka-console-consumer', '--bootstrap-server', 'kafka:29092', '--topic', 'source.metadata', '--partition', '0', '--offset', str(offset), '--max-messages', '1', '--property', 'print.key=true', '--property', 'print.value=true'],
        capture_output=True, text=True, check=True,
    )
    key, raw_value = result.stdout.strip().split('\t', 1)
    event = json.loads(raw_value)
    return key, event

event_key, actual_event_v1 = read_metadata_record(baseline_latest_offset)
validator.validate('FILE_METADATA_UPSERT', actual_event_v1)
assert event_key == file_id
assert actual_event_v1['file_id'] == file_id
assert actual_event_v1['repository_id'] == repository_id
assert actual_event_v1['content_hash'] == first_content_hash
assert actual_event_v1['metadata']['parse_status'] == 'SUCCESS'
print(json.dumps({'kafka_key': event_key, 'event_id': actual_event_v1['event_id'], 'schema_valid': True, 'file_id_matches_key': True, 'metadata': actual_event_v1['metadata']}, indent=2))

NETWORK = subprocess.check_output(['docker', 'inspect', '-f', '{{range $k, $v := .NetworkSettings.Networks}}{{$k}}{{end}}', 'cpg-kafka'], text=True).strip()
INTERNAL_MONGO_URI = MONGO_URI.replace('localhost', 'cpg-mongodb')
SPARK_PACKAGES = 'org.mongodb.spark:mongo-spark-connector_2.12:10.1.1,org.apache.spark:spark-sql-kafka-0-10_2.12:3.3.0'
STARTING_OFFSETS = json.dumps({'source.metadata': {'0': baseline_latest_offset}})

def run_spark(starting_offsets: str | None = None) -> tuple[int, str]:
    command = ['docker', 'run', '--rm', '--network', NETWORK, '-v', f'{PROJECT_ROOT}:/opt/project', '-w', '/opt/project', 'apache/spark-py:v3.3.0', '/opt/spark/bin/spark-submit', '--master', 'local[2]', '--conf', 'spark.jars.ivy=/tmp/ivy', '--packages', SPARK_PACKAGES, 'spark_jobs/metadata_to_mongodb.py', '--bootstrap-servers', 'kafka:29092', '--mongodb-uri', INTERNAL_MONGO_URI, '--checkpoint-dir', CHECKPOINT_DIR.relative_to(PROJECT_ROOT).as_posix(), '--app-name', f'cpg-task5-{RUN_ID}']
    if starting_offsets is not None:
        command.extend(['--starting-offsets', starting_offsets])
    command.append('--available-now')
    result = subprocess.run(command, capture_output=True, text=True, timeout=300)
    return result.returncode, result.stdout + result.stderr

def input_rows(logs: str) -> list[int]:
    return [int(value) for value in re.findall(r'"numInputRows"\s*:\s*(\d+)', logs)]

initial_exit, initial_logs = run_spark(STARTING_OFFSETS)
initial_input_rows = input_rows(initial_logs)
assert initial_exit == 0
assert initial_input_rows and initial_input_rows[-1] == 1
print(json.dumps({'spark_exit': initial_exit, 'num_input_rows': initial_input_rows[-1], 'starting_offsets': STARTING_OFFSETS, 'checkpoint': CHECKPOINT_DIR.as_posix()}, indent=2))

def mongo_json(javascript: str) -> dict | list:
    result = subprocess.run(['docker', 'exec', 'cpg-mongodb', 'mongosh', '--quiet', '--username', os.environ['MONGO_ROOT_USERNAME'], '--password', os.environ['MONGO_ROOT_PASSWORD'], '--authenticationDatabase', 'admin', '--eval', javascript], capture_output=True, text=True, check=True)
    return json.loads(result.stdout.strip().splitlines()[-1])

def document_state() -> dict:
    return mongo_json(
        f"db=db.getSiblingDB('{MONGO_DATABASE}'); const d=db.{MONGO_COLLECTION}.findOne({{file_id:'{file_id}'}}); print(JSON.stringify({{document_count:db.{MONGO_COLLECTION}.countDocuments({{file_id:'{file_id}'}}), content_hash:d ? d.content_hash : null, repository_id:d ? d.repository_id : null, file_path:d ? d.file_path : null}}));"
    )

baseline_document = document_state()
baseline_checkpoint_files = len([path for path in CHECKPOINT_DIR.rglob('*') if path.is_file()])
print(json.dumps({'mongo': baseline_document, 'checkpoint_file_count': baseline_checkpoint_files}, indent=2))
assert baseline_document['document_count'] == 1
assert baseline_document['content_hash'] == first_content_hash
assert baseline_checkpoint_files > 0

restart_exit, restart_logs = run_spark()
restart_input_rows = input_rows(restart_logs)
restart_document = document_state()
assert restart_exit == 0
assert restart_input_rows and restart_input_rows[-1] == 0
assert restart_document == baseline_document
print(json.dumps({'spark_exit': restart_exit, 'num_input_rows': restart_input_rows[-1], 'mongo_after_restart': restart_document, 'document_unchanged': True}, indent=2))

before_replay_offset = kafka_latest_offset()
source_path.write_bytes(modified_source)
replay_result = process_service.execute(make_source_file())
after_replay_offset = kafka_latest_offset()
assert replay_result.status.value == 'SUCCESS'
assert after_replay_offset == before_replay_offset + 1
second_event_offset = after_replay_offset - 1
second_key, actual_event_v2 = read_metadata_record(second_event_offset)
validator.validate('FILE_METADATA_UPSERT', actual_event_v2)
assert second_key == file_id
assert actual_event_v2['file_id'] == file_id
assert actual_event_v2['content_hash'] != first_content_hash
print(json.dumps({'replay_status': replay_result.status.value, 'old_content_hash': first_content_hash, 'new_content_hash': actual_event_v2['content_hash'], 'kafka_key_matches_file_id': second_key == file_id, 'schema_valid': True}, indent=2))

upsert_exit, upsert_logs = run_spark()
upsert_input_rows = input_rows(upsert_logs)
final_document = document_state()
assert upsert_exit == 0
assert upsert_input_rows and upsert_input_rows[-1] == 1
assert final_document['document_count'] == 1
assert final_document['content_hash'] == actual_event_v2['content_hash']
print(json.dumps({'spark_exit': upsert_exit, 'num_input_rows': upsert_input_rows[-1], 'mongo_after_upsert': final_document, 'upsert_kept_one_document': True}, indent=2))

duplicates = mongo_json(
    f"db=db.getSiblingDB('{MONGO_DATABASE}'); const by_file_id=db.{MONGO_COLLECTION}.aggregate([{{$group:{{_id:'$file_id', count:{{$sum:1}}}}}}, {{$match:{{count:{{$gt:1}}}}}}]).toArray(); const by_repository_path=db.{MONGO_COLLECTION}.aggregate([{{$group:{{_id:{{repository_id:'$repository_id', file_path:'$file_path'}}, count:{{$sum:1}}}}}}, {{$match:{{count:{{$gt:1}}}}}}]).toArray(); print(JSON.stringify({{by_file_id:by_file_id, by_repository_path:by_repository_path}}));"
)
assert duplicates['by_file_id'] == []
assert duplicates['by_repository_path'] == []
print(json.dumps({'duplicate_by_file_id': duplicates['by_file_id'], 'duplicate_by_repository_path': duplicates['by_repository_path'], 'duplicate_checks_passed': True}, indent=2))