sashatrask
1006 lines
45 KiB
Python
1006 lines
45 KiB
Python
import json
|
|
import os
|
|
from pathlib import Path
|
|
import sys
|
|
import tempfile
|
|
import unittest
|
|
from unittest import mock
|
|
|
|
|
|
ROOT = Path(__file__).resolve().parents[1]
|
|
APP_DIR = ROOT / 'app'
|
|
sys.path.insert(0, str(APP_DIR))
|
|
|
|
import scanner_db
|
|
import console_runner
|
|
import supervisor
|
|
from scanner_db import ScannerDB
|
|
|
|
|
|
class Cursor:
|
|
def __init__(self, rows=None, rowcount=0):
|
|
self.rows = list(rows or [])
|
|
self.rowcount = rowcount
|
|
|
|
def fetchall(self):
|
|
return list(self.rows)
|
|
|
|
def fetchone(self):
|
|
return self.rows[0] if self.rows else None
|
|
|
|
|
|
class ClaimConnection:
|
|
is_postgres = True
|
|
is_sqlite = False
|
|
|
|
def __init__(self, exhausted=False):
|
|
self.exhausted = exhausted
|
|
self.statements = []
|
|
self.claim = None
|
|
|
|
def execute(self, sql, params=None):
|
|
params = tuple(params or ())
|
|
self.statements.append((sql, params))
|
|
normalized = ' '.join(str(sql).split())
|
|
if "status = 'in_progress'" in normalized and 'lease_owner = ? AND claim_batch = ?' in normalized:
|
|
return Cursor()
|
|
if 'SELECT q.id FROM candidates c' in normalized:
|
|
return Cursor([{'id': 99}] if self.exhausted else [])
|
|
if 'SELECT q.id, q.target FROM candidates c' in normalized:
|
|
return Cursor([{'id': 7, 'target': 'https://github.com/example/repo'}])
|
|
if normalized.startswith("UPDATE target_queue SET status = 'failed'"):
|
|
return Cursor(rowcount=1)
|
|
if normalized.startswith("UPDATE target_queue SET status = 'in_progress'"):
|
|
self.claim = {
|
|
'id': int(params[-1]),
|
|
'target': 'https://github.com/example/repo',
|
|
'normalized_target': 'https://github.com/example/repo',
|
|
'attempts': 1,
|
|
'lease_owner': params[0],
|
|
'lease_token': params[1],
|
|
'claim_batch': params[2],
|
|
}
|
|
return Cursor(rowcount=1)
|
|
if 'FROM target_queue WHERE id IN (' in normalized:
|
|
return Cursor([self.claim])
|
|
raise AssertionError(normalized)
|
|
|
|
def commit(self):
|
|
return None
|
|
|
|
def rollback(self):
|
|
return None
|
|
|
|
|
|
class ProductionQueryShapeTests(unittest.TestCase):
|
|
@staticmethod
|
|
def claim_db(connection):
|
|
db = ScannerDB.__new__(ScannerDB)
|
|
db.conn = connection
|
|
db._last_claim_expectation = None
|
|
db.last_error = ''
|
|
return db
|
|
|
|
def test_postgres_claim_uses_status_branches_and_skip_locked(self):
|
|
connection = ClaimConnection()
|
|
rows = self.claim_db(connection).claim_targets(
|
|
'source', 'github', 1, 'owner', 60,
|
|
max_attempts=0, return_rows=True, claim_batch='batch',
|
|
)
|
|
self.assertEqual([row['id'] for row in rows], [7])
|
|
candidate_sql = next(
|
|
sql for sql, _ in connection.statements if 'SELECT q.id, q.target FROM candidates c' in sql
|
|
)
|
|
self.assertEqual(candidate_sql.count('UNION ALL'), 2)
|
|
self.assertIn("status = 'pending'", candidate_sql)
|
|
self.assertIn("status = 'deferred'", candidate_sql)
|
|
self.assertIn("status = 'in_progress'", candidate_sql)
|
|
self.assertIn('FOR UPDATE OF q SKIP LOCKED', candidate_sql)
|
|
|
|
def test_docker_depth_dispatch_is_bounded_skip_locked_and_workload_first(self):
|
|
source = (APP_DIR / 'scanner_db.py').read_text(encoding='utf-8')
|
|
candidate_start = source.index(
|
|
' def _claim_docker_depth_experiment_target_locked('
|
|
)
|
|
candidate_end = source.index(
|
|
'\n def _docker_depth_binding_for_reservation_locked(', candidate_start,
|
|
)
|
|
candidate = source[candidate_start:candidate_end]
|
|
self.assertIn("target.state = 'pending'", candidate)
|
|
self.assertIn('prior.dispatch_wave < target.dispatch_wave', candidate)
|
|
self.assertIn("prior.state NOT IN ('done','failed')", candidate)
|
|
self.assertIn('binding.attempt = prior.reservation_count', candidate)
|
|
self.assertIn("binding.state = 'completed'", candidate)
|
|
self.assertIn("binding.state = 'failed'", candidate)
|
|
self.assertIn('ORDER BY target.dispatch_order, target.id', candidate)
|
|
self.assertIn('LIMIT 1 FOR UPDATE OF target, queue SKIP LOCKED', candidate)
|
|
self.assertIn('docker_depth_experiment_selections selection', candidate)
|
|
|
|
reserve_start = source.index(' def reserve_and_claim_target(')
|
|
reserve_end = source.index('\n @staticmethod\n def _result_reservation_claim', reserve_start)
|
|
reserve = source[reserve_start:reserve_end]
|
|
authority_lock = reserve.index(
|
|
'self._locked_docker_depth_experiment_authority('
|
|
)
|
|
workload_lock = reserve.index(
|
|
'self._claim_docker_depth_experiment_target_locked('
|
|
)
|
|
capacity_lock = reserve.index(
|
|
"'SELECT * FROM pipeline_capacity WHERE id = 1 FOR UPDATE'"
|
|
)
|
|
reservation_insert = reserve.index('INSERT INTO result_reservations(')
|
|
binding_insert = reserve.index(
|
|
'INSERT INTO docker_depth_experiment_scan_bindings('
|
|
)
|
|
self.assertLess(authority_lock, workload_lock)
|
|
self.assertLess(workload_lock, capacity_lock)
|
|
self.assertLess(capacity_lock, reservation_insert)
|
|
self.assertLess(reservation_insert, binding_insert)
|
|
|
|
def test_remote_dispatch_gate_precedes_every_new_admission_lock(self):
|
|
source = (APP_DIR / 'scanner_db.py').read_text(encoding='utf-8')
|
|
reserve_start = source.index(' def reserve_and_claim_target(')
|
|
reserve_end = source.index(
|
|
'\n @staticmethod\n def _result_reservation_claim', reserve_start,
|
|
)
|
|
reserve = source[reserve_start:reserve_end]
|
|
intent_lock = reserve.index(
|
|
'SELECT * FROM admission_intents WHERE reservation_token = ? FOR UPDATE'
|
|
)
|
|
dispatch_gate = reserve.index(
|
|
'self._locked_runtime_control_state(shared=True)'
|
|
)
|
|
experiment_lock = reserve.index(
|
|
'self._locked_docker_depth_experiment_authority('
|
|
)
|
|
remote_user_lock = reserve.index('FROM remote_worker_users')
|
|
workload_lock = reserve.index('LIMIT 1 FOR UPDATE SKIP LOCKED')
|
|
capacity_lock = reserve.index(
|
|
"'SELECT * FROM pipeline_capacity WHERE id = 1 FOR UPDATE'"
|
|
)
|
|
reservation_insert = reserve.index('INSERT INTO result_reservations(')
|
|
self.assertLess(intent_lock, dispatch_gate)
|
|
self.assertLess(dispatch_gate, experiment_lock)
|
|
self.assertLess(dispatch_gate, remote_user_lock)
|
|
self.assertLess(dispatch_gate, workload_lock)
|
|
self.assertLess(dispatch_gate, capacity_lock)
|
|
self.assertLess(dispatch_gate, reservation_insert)
|
|
self.assertIn("if remote and intent['state'] == 'pending':", reserve)
|
|
self.assertIn("resolution_detail = 'dispatch_gate_closed'", reserve)
|
|
|
|
def test_multisource_plan_validation_precedes_capacity_and_survives_recovery(self):
|
|
source = (APP_DIR / 'scanner_db.py').read_text(encoding='utf-8')
|
|
reserve_start = source.index(' def reserve_and_claim_target(')
|
|
reserve_end = source.index(
|
|
'\n @staticmethod\n def _result_reservation_claim', reserve_start,
|
|
)
|
|
reserve = source[reserve_start:reserve_end]
|
|
target_lock = reserve.index('LIMIT 1 FOR UPDATE SKIP LOCKED')
|
|
contextual_validation = reserve.index('remote_assignment_execution_plan({')
|
|
capacity_lock = reserve.index(
|
|
"'SELECT * FROM pipeline_capacity WHERE id = 1 FOR UPDATE'"
|
|
)
|
|
self.assertLess(target_lock, contextual_validation)
|
|
self.assertLess(contextual_validation, capacity_lock)
|
|
self.assertIn('remote_assignment_execution_plan(existing)', reserve)
|
|
|
|
reconcile_start = source.index(
|
|
' def reconcile_remote_assignment_request('
|
|
)
|
|
reconcile_end = source.index(
|
|
'\n def renew_result_claim(', reconcile_start,
|
|
)
|
|
reconcile = source[reconcile_start:reconcile_end]
|
|
self.assertIn('FOR UPDATE OF r, q', reconcile)
|
|
self.assertIn('execution_plan = remote_assignment_execution_plan(', reconcile)
|
|
self.assertIn("'execution_plan': execution_plan", reconcile)
|
|
|
|
ready_start = source.index(' def mark_result_bundle_ready(')
|
|
ready_end = source.index(
|
|
'\n def recover_expired_result_bundle_ready(', ready_start,
|
|
)
|
|
ready = source[ready_start:ready_end]
|
|
self.assertIn('validate_remote_result_execution_plan(', ready)
|
|
self.assertNotIn("platform']).lower() in ('github', 'gitlab')", ready)
|
|
|
|
ingest_start = source.index(' def ingest_result_bundle(')
|
|
ingest_end = source.index('\n def ', ingest_start + 10)
|
|
ingest = source[ingest_start:ingest_end]
|
|
self.assertIn('validate_remote_result_execution_plan(', ingest)
|
|
|
|
def test_drain_reconciliation_locks_control_before_unlocked_blocker_counts(self):
|
|
source = (APP_DIR / 'scanner_db.py').read_text(encoding='utf-8')
|
|
blockers_start = source.index(' def _runtime_drain_blockers_locked(')
|
|
blockers_end = source.index('\n def runtime_drain_progress(', blockers_start)
|
|
blockers = source[blockers_start:blockers_end]
|
|
reconcile_start = source.index(' def reconcile_runtime_drain(')
|
|
reconcile_end = source.index('\n def final_cutover_status(', reconcile_start)
|
|
reconcile = source[reconcile_start:reconcile_end]
|
|
commit_start = source.index(
|
|
' def _commit_runtime_control_transition_locked('
|
|
)
|
|
commit_end = source.index('\n def _runtime_control_mutation(', commit_start)
|
|
commit = source[commit_start:commit_end]
|
|
|
|
control_lock = reconcile.index('self._locked_runtime_control_state()')
|
|
blocker_query = reconcile.index('self._runtime_drain_blockers_locked()')
|
|
transition = reconcile.index('self._commit_runtime_control_transition_locked(')
|
|
self.assertLess(control_lock, blocker_query)
|
|
self.assertLess(blocker_query, transition)
|
|
self.assertIn("assignment_kind = 'remote'", blockers)
|
|
self.assertIn('remote_resolved_at IS NULL', blockers)
|
|
self.assertIn("state IN ('ready', 'ingesting')", blockers)
|
|
self.assertNotIn('FOR UPDATE', blockers)
|
|
self.assertIn("state['drain_state'] != 'draining'", reconcile)
|
|
self.assertIn("action='control.drain.complete'", reconcile)
|
|
self.assertIn('WHERE id = 1 AND revision = ?', commit)
|
|
|
|
def test_discovery_retry_claim_is_allowlisted_bounded_and_skip_locked(self):
|
|
class DiscoveryConnection:
|
|
is_postgres = True
|
|
is_sqlite = False
|
|
|
|
def __init__(self):
|
|
self.statements = []
|
|
|
|
def execute(self, sql, params=None):
|
|
params = tuple(params or ())
|
|
self.statements.append((sql, params))
|
|
normalized = ' '.join(str(sql).split())
|
|
if normalized == 'SELECT * FROM runtime_operations_control WHERE id = 1 FOR SHARE':
|
|
return Cursor([{
|
|
'revision': 0,
|
|
'discovery_paused': 0,
|
|
'dispatch_paused': 0,
|
|
'drain_state': 'normal',
|
|
'actor': 'system:migration',
|
|
'operation_id': None,
|
|
'created_at': '2026-09-19T00:00:00+00:00',
|
|
'updated_at': '2026-09-19T00:00:00+00:00',
|
|
}])
|
|
if normalized.startswith("UPDATE discovery_retry_queue SET status = 'held'"):
|
|
return Cursor(rowcount=0)
|
|
if normalized.startswith('SELECT id, work_key, source, query, policy_sha256'):
|
|
return Cursor([{
|
|
'id': 7,
|
|
'work_key': 'b' * 64,
|
|
'source': 'dockerhub',
|
|
'query': 'exact-query',
|
|
'source_cycle_id': None,
|
|
'policy_sha256': 'a' * 64,
|
|
'pass_kind': 'ordinary',
|
|
'work_kind': 'page',
|
|
'page_start': 2,
|
|
'page_end': 2,
|
|
'next_page': 2,
|
|
'attempts': 1,
|
|
'last_error_category': 'network',
|
|
}])
|
|
if normalized.startswith("UPDATE discovery_retry_queue SET status = 'leased'"):
|
|
return Cursor([{'attempts': 2}], rowcount=1)
|
|
raise AssertionError(normalized)
|
|
|
|
def commit(self):
|
|
return None
|
|
|
|
def rollback(self):
|
|
return None
|
|
|
|
connection = DiscoveryConnection()
|
|
db = self.claim_db(connection)
|
|
with mock.patch.object(scanner_db.secrets, 'token_urlsafe', return_value='lease-token'):
|
|
rows = db.claim_discovery_retries(
|
|
'dockerhub', {'exact-query': 'a' * 64}, 'owner',
|
|
limit=1, lease_seconds=60,
|
|
)
|
|
self.assertEqual(rows[0]['lease_token'], 'lease-token')
|
|
self.assertEqual(rows[0]['attempts'], 2)
|
|
self.assertEqual(
|
|
' '.join(connection.statements[0][0].split()),
|
|
'SELECT * FROM runtime_operations_control WHERE id = 1 FOR SHARE',
|
|
)
|
|
candidate_sql, candidate_params = next(
|
|
(sql, params) for sql, params in connection.statements
|
|
if 'ORDER BY COALESCE(available_after' in sql
|
|
)
|
|
self.assertIn("status = 'pending'", candidate_sql)
|
|
self.assertIn("status = 'leased'", candidate_sql)
|
|
self.assertIn('lease_expires_at <= ?', candidate_sql)
|
|
self.assertIn('FOR UPDATE SKIP LOCKED', candidate_sql)
|
|
self.assertTrue(candidate_sql.rstrip().endswith('FOR UPDATE SKIP LOCKED'))
|
|
self.assertIn('exact-query', candidate_params)
|
|
self.assertNotIn('exact-query', candidate_sql)
|
|
self.assertEqual(candidate_params[-1], 1)
|
|
claim_sql = next(
|
|
sql for sql, _ in connection.statements
|
|
if "UPDATE discovery_retry_queue SET status = 'leased'" in sql
|
|
)
|
|
self.assertIn('attempts + 1', claim_sql)
|
|
self.assertIn('RETURNING attempts', claim_sql)
|
|
|
|
def test_discovery_page_persistence_propagates_database_failure(self):
|
|
class FailedConnection:
|
|
is_postgres = True
|
|
is_sqlite = False
|
|
|
|
def execute(self, sql, params=None):
|
|
raise RuntimeError('synthetic database failure')
|
|
|
|
def rollback(self):
|
|
return None
|
|
|
|
db = self.claim_db(FailedConnection())
|
|
with self.assertRaisesRegex(RuntimeError, 'synthetic database failure'):
|
|
db.persist_dockerhub_discovery_page(
|
|
'dockerhub', 'exact-query', [{'repo_name': 'fixture/database-failure'}],
|
|
)
|
|
|
|
def test_docker_shadow_controls_are_bounded_policy_exact_and_hash_only(self):
|
|
class ShadowConnection:
|
|
is_postgres = True
|
|
is_sqlite = False
|
|
|
|
def __init__(self):
|
|
self.statements = []
|
|
|
|
def execute(self, sql, params=None):
|
|
params = tuple(params or ())
|
|
self.statements.append((sql, params))
|
|
normalized = ' '.join(str(sql).split())
|
|
if normalized.startswith('WITH eligible AS'):
|
|
return Cursor([{
|
|
'target_scan_id': value,
|
|
'normalized_target': (
|
|
f'owner/repo-{value}@sha256:' + format(value, '064x')
|
|
),
|
|
} for value in range(1, 51)])
|
|
if normalized.startswith('SELECT c.routed_service, kc.provider_key_hash'):
|
|
return Cursor([{
|
|
'routed_service': 'openai',
|
|
'provider_key_hash': 'c' * 64,
|
|
}])
|
|
if normalized.startswith('SELECT detector_secret_hash FROM findings'):
|
|
return Cursor([{'detector_secret_hash': 'd' * 64}])
|
|
raise AssertionError(normalized)
|
|
|
|
def commit(self):
|
|
return None
|
|
|
|
def rollback(self):
|
|
raise AssertionError('valid shadow reads must not roll back')
|
|
|
|
connection = ShadowConnection()
|
|
db = ScannerDB.__new__(ScannerDB)
|
|
db.conn = connection
|
|
controls = db.docker_adaptive_shadow_controls('a' * 64, 50, 'b' * 64)
|
|
routed, detectors = db.docker_adaptive_shadow_control_identities(controls[0]['target_scan_id'])
|
|
|
|
self.assertEqual(len(controls), 50)
|
|
self.assertEqual(routed, frozenset({('openai', 'c' * 64)}))
|
|
self.assertEqual(detectors, frozenset({'d' * 64}))
|
|
control_sql, control_params = connection.statements[0]
|
|
normalized = ' '.join(control_sql.split())
|
|
self.assertIn('ROW_NUMBER() OVER ( PARTITION BY ts.normalized_target', normalized)
|
|
self.assertIn("r.state = 'acknowledged' AND b.state = 'acknowledged'", normalized)
|
|
self.assertIn("ts.scan_options_json::jsonb ->> 'scan_policy_sha256'", normalized)
|
|
self.assertIn("docker_scan_assignment,effective_mode}') = 'full'", normalized)
|
|
self.assertIn('ORDER BY md5(? ||', normalized)
|
|
self.assertTrue(normalized.endswith('LIMIT ?'))
|
|
self.assertEqual(control_params, ('a' * 64, 'b' * 64, 50))
|
|
|
|
routed_sql = ' '.join(connection.statements[1][0].split()).lower()
|
|
detector_sql = ' '.join(connection.statements[2][0].split()).lower()
|
|
self.assertIn('c.routed_service, kc.provider_key_hash', routed_sql)
|
|
self.assertNotIn('secret_text', routed_sql)
|
|
self.assertNotIn('credential_hash', routed_sql)
|
|
self.assertEqual(
|
|
detector_sql,
|
|
'select detector_secret_hash from findings where target_scan_id = ?',
|
|
)
|
|
|
|
def test_exhausted_attempts_update_is_bounded_to_locked_ids(self):
|
|
connection = ClaimConnection(exhausted=True)
|
|
self.claim_db(connection).claim_targets(
|
|
'source', 'github', 1, 'owner', 60,
|
|
max_attempts=3, return_rows=True, claim_batch='batch',
|
|
)
|
|
exhausted_select = next(
|
|
sql for sql, _ in connection.statements if 'SELECT q.id FROM candidates c' in sql
|
|
)
|
|
exhausted_update = next(
|
|
sql for sql, _ in connection.statements
|
|
if "UPDATE target_queue SET status = 'failed'" in sql
|
|
)
|
|
self.assertIn('FOR UPDATE OF q SKIP LOCKED', exhausted_select)
|
|
self.assertIn('WHERE id IN (?)', exhausted_update)
|
|
self.assertNotIn('WHERE source = ?', exhausted_update)
|
|
|
|
def test_known_target_lookup_hard_caps_each_parameter_batch(self):
|
|
class KnownConnection:
|
|
is_postgres = True
|
|
|
|
def __init__(self):
|
|
self.calls = []
|
|
|
|
def execute(self, sql, params=None):
|
|
params = tuple(params or ())
|
|
self.calls.append((sql, params))
|
|
return Cursor([{'normalized_target': value} for value in params[2:]])
|
|
|
|
def commit(self):
|
|
return None
|
|
|
|
def rollback(self):
|
|
return None
|
|
|
|
connection = KnownConnection()
|
|
db = self.claim_db(connection)
|
|
targets = [f'https://github.com/example/repo-{index}' for index in range(150)]
|
|
known = db.known_target_normalizations_for('source', 'github', targets, batch_size=500)
|
|
self.assertEqual(len(known), 150)
|
|
self.assertEqual(len(connection.calls), 3)
|
|
self.assertLessEqual(max(len(params) for _, params in connection.calls), 66)
|
|
|
|
def test_runtime_schema_contains_exact_composite_and_partial_indexes(self):
|
|
with tempfile.TemporaryDirectory() as temp_dir, mock.patch.dict(os.environ, {
|
|
'SCANNER_DB_URL': '', 'DATABASE_URL': '',
|
|
}):
|
|
db = ScannerDB(db_path=os.path.join(temp_dir, 'scanner.db'))
|
|
try:
|
|
queue_indexes = db.conn.table_indexes('target_queue')
|
|
scan_indexes = db.conn.table_indexes('target_scans')
|
|
finally:
|
|
db.close()
|
|
self.assertEqual(
|
|
queue_indexes['idx_target_queue_source_platform_normalized']['columns'],
|
|
['source', 'platform', 'normalized_target'],
|
|
)
|
|
self.assertEqual(
|
|
scan_indexes['idx_target_scans_source_ended_id']['columns'],
|
|
['source', 'ended_at', 'id'],
|
|
)
|
|
self.assertEqual(
|
|
scan_indexes['idx_target_scans_queue_id']['columns'],
|
|
['queue_id'],
|
|
)
|
|
self.assertEqual(
|
|
queue_indexes['idx_target_queue_observe_source_status']['columns'],
|
|
['source', 'status'],
|
|
)
|
|
active_lease_index = queue_indexes['idx_target_queue_active_lease_owner_token']
|
|
self.assertEqual(active_lease_index['columns'], ['lease_owner', 'lease_token'])
|
|
self.assertEqual(
|
|
scanner_db._normalized_predicate(active_lease_index['predicate']),
|
|
"status='in_progress'",
|
|
)
|
|
cooldown_index = scan_indexes['idx_target_scans_cooldown_recent']
|
|
self.assertEqual(
|
|
cooldown_index['columns'],
|
|
['source', 'ended_at', 'id'],
|
|
)
|
|
self.assertEqual(
|
|
scanner_db._normalized_predicate(cooldown_index['predicate']),
|
|
scanner_db._normalized_predicate(scanner_db.CI_COOLDOWN_INDEX_PREDICATE),
|
|
)
|
|
for status in ('pending', 'deferred', 'in_progress'):
|
|
name = f'idx_target_queue_claim_{status}'
|
|
predicate = scanner_db._normalized_predicate(queue_indexes[name]['predicate'])
|
|
self.assertIn(scanner_db._normalized_predicate(f"status = '{status}'"), predicate)
|
|
self.assertIn('resolver_stateisnullorresolver_state=', predicate)
|
|
self.assertEqual(
|
|
queue_indexes['idx_target_queue_exhausted_attempts']['columns'],
|
|
['source', 'platform', 'status', 'attempts', 'id', 'lease_expires_at'],
|
|
)
|
|
|
|
def test_package_lookup_indexes_and_short_term_fallback_preserve_newest_order(self):
|
|
with tempfile.TemporaryDirectory() as temp_dir, mock.patch.dict(os.environ, {
|
|
'SCANNER_DB_URL': '', 'DATABASE_URL': '',
|
|
}):
|
|
db = ScannerDB(db_path=os.path.join(temp_dir, 'scanner.db'))
|
|
try:
|
|
indexes = db.conn.table_indexes('package_repo_candidates')
|
|
self.assertNotIn('idx_package_repo_candidates_name_trgm', indexes)
|
|
self.assertEqual(
|
|
indexes['idx_package_repo_candidates_query_seen']['columns'],
|
|
['query', 'last_seen_at', 'id'],
|
|
)
|
|
self.assertEqual(
|
|
indexes['idx_package_repo_candidates_recent_lookup']['columns'],
|
|
['last_seen_at', 'id', 'package_source', 'query', 'package_name'],
|
|
)
|
|
for name in (
|
|
'idx_package_repo_candidates_query_seen',
|
|
'idx_package_repo_candidates_recent_lookup',
|
|
):
|
|
self.assertEqual(
|
|
scanner_db._normalized_predicate(indexes[name]['predicate']),
|
|
scanner_db._normalized_predicate(scanner_db.PACKAGE_REPO_NONEMPTY_PREDICATE),
|
|
)
|
|
rows = (
|
|
('pypi', 'contains-ai-excluded', '1', 'other', 'https://example.invalid/excluded', 'github', '2026-07-20T05:00:00'),
|
|
('npm', 'contains-ai-newest', '1', 'other', 'https://example.invalid/newest', 'github', '2026-07-20T03:00:00'),
|
|
('npm', 'exact-query', '1', 'ai', 'https://example.invalid/exact', 'github', '2026-07-20T02:00:00'),
|
|
('npm', 'unrelated', '1', 'other', 'https://example.invalid/unrelated', 'github', '2026-07-20T04:00:00'),
|
|
)
|
|
for source, name, version, query, repo, provider, seen in rows:
|
|
db.conn.execute(
|
|
'''INSERT INTO package_repo_candidates(
|
|
package_source, package_name, package_version, query, repo_url,
|
|
provider, evidence_json, confidence, first_seen_at, last_seen_at
|
|
) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?)''',
|
|
(source, name, version, query, repo, provider, '[]', 'high', seen, seen),
|
|
)
|
|
db.conn.commit()
|
|
candidates = db.get_package_repo_candidates('ai', ['npm'], limit=2)
|
|
finally:
|
|
db.close()
|
|
self.assertEqual(
|
|
[candidate['name'] for candidate in candidates],
|
|
['contains-ai-newest', 'exact-query'],
|
|
)
|
|
|
|
def test_postgres_package_lookup_uses_ordered_metadata_scan_without_like(self):
|
|
class PackageConnection:
|
|
is_postgres = True
|
|
|
|
def __init__(self):
|
|
self.calls = []
|
|
|
|
def execute(self, sql, params=None):
|
|
params = tuple(params or ())
|
|
normalized = ' '.join(sql.split())
|
|
self.calls.append((normalized, params))
|
|
if normalized.startswith('SET LOCAL'):
|
|
return Cursor()
|
|
common = {
|
|
'package_source': 'npm', 'version': '1', 'provider': 'github',
|
|
'evidence_json': '[]', 'confidence': 'high',
|
|
}
|
|
if 'WHERE id IN' in normalized:
|
|
details = {
|
|
3: dict(common, id=3, name='workflow-substring-second', repo_url='https://example.invalid/3', last_seen_at='2026-07-20T03:00:00'),
|
|
4: dict(common, id=4, name='workflow-substring-third', repo_url='https://example.invalid/4', last_seen_at='2026-07-20T02:00:00'),
|
|
}
|
|
return Cursor([details[row_id] for row_id in params])
|
|
if 'AND query = ?' in normalized:
|
|
return Cursor([
|
|
dict(common, id=1, name='exact-newest', repo_url='https://example.invalid/1', last_seen_at='2026-07-20T04:00:00'),
|
|
dict(common, id=2, name='exact-older', repo_url='https://example.invalid/2', last_seen_at='2026-07-20T01:00:00'),
|
|
])
|
|
if 'SELECT id, package_source, package_name, query, last_seen_at' in normalized:
|
|
return Cursor([
|
|
{'id': 5, 'package_source': 'pypi', 'package_name': 'workflow-excluded', 'query': 'other', 'last_seen_at': '2026-07-20T05:00:00'},
|
|
{'id': 6, 'package_source': 'npm', 'package_name': 'unrelated', 'query': 'other', 'last_seen_at': '2026-07-20T04:30:00'},
|
|
{'id': 1, 'package_source': 'npm', 'package_name': 'exact-newest', 'query': 'workflow', 'last_seen_at': '2026-07-20T04:00:00'},
|
|
{'id': 3, 'package_source': 'npm', 'package_name': 'workflow-substring-second', 'query': 'other', 'last_seen_at': '2026-07-20T03:00:00'},
|
|
{'id': 4, 'package_source': 'npm', 'package_name': 'workflow-substring-third', 'query': 'other', 'last_seen_at': '2026-07-20T02:00:00'},
|
|
])
|
|
raise AssertionError(normalized)
|
|
|
|
def commit(self):
|
|
return None
|
|
|
|
def rollback(self):
|
|
return None
|
|
|
|
connection = PackageConnection()
|
|
db = self.claim_db(connection)
|
|
candidates = db.get_package_repo_candidates('workflow', ['npm'], limit=2)
|
|
self.assertEqual(
|
|
[candidate['name'] for candidate in candidates],
|
|
['exact-newest', 'workflow-substring-second'],
|
|
)
|
|
queries = [(sql, params) for sql, params in connection.calls if 'FROM package_repo_candidates' in sql]
|
|
self.assertEqual(len(queries), 3)
|
|
self.assertFalse(any('package_name LIKE' in sql for sql, _ in queries))
|
|
exact_sql, exact_params = queries[0]
|
|
metadata_sql, metadata_params = queries[1]
|
|
self.assertIn('package_source IN (?)', exact_sql)
|
|
self.assertEqual(exact_params, ('npm', 'workflow', 2))
|
|
self.assertIn('ORDER BY last_seen_at DESC, id DESC LIMIT ?', metadata_sql)
|
|
self.assertEqual(metadata_params, (250,))
|
|
|
|
def test_package_fallback_reports_when_bounded_scan_cannot_prove_final_top(self):
|
|
with tempfile.TemporaryDirectory() as temp_dir, mock.patch.dict(os.environ, {
|
|
'SCANNER_DB_URL': '', 'DATABASE_URL': '',
|
|
}):
|
|
db = ScannerDB(db_path=os.path.join(temp_dir, 'scanner.db'))
|
|
try:
|
|
for index in range(3):
|
|
seen = f'2026-07-20T0{index}:00:00'
|
|
query = 'ai' if index == 0 else 'other'
|
|
name = 'exact-outside-bounded-scan' if index == 0 else f'unrelated-{index}'
|
|
db.conn.execute(
|
|
'''INSERT INTO package_repo_candidates(
|
|
package_source, package_name, package_version, query, repo_url,
|
|
provider, evidence_json, confidence, first_seen_at, last_seen_at
|
|
) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?)''',
|
|
('npm', name, '1', query,
|
|
f'https://example.invalid/{index}', 'github', '[]', 'high', seen, seen),
|
|
)
|
|
db.conn.commit()
|
|
with mock.patch.object(
|
|
scanner_db, 'PACKAGE_CANDIDATE_SCAN_MAX_ROWS', 2,
|
|
), mock.patch.object(scanner_db.logger, 'warning') as warning:
|
|
candidates = db.get_package_repo_candidates('ai', ['npm'], limit=2)
|
|
finally:
|
|
db.close()
|
|
self.assertEqual(
|
|
[candidate['name'] for candidate in candidates],
|
|
['exact-outside-bounded-scan'],
|
|
)
|
|
self.assertTrue(any(
|
|
'bounded' in str(call.args[0]).lower() for call in warning.call_args_list
|
|
))
|
|
|
|
def test_migration_retires_dead_trigram_index_but_retains_extension_idempotently(self):
|
|
class MigrationConnection:
|
|
is_postgres = True
|
|
|
|
def __init__(self):
|
|
self.statements = []
|
|
|
|
def execute(self, sql, params=None):
|
|
normalized = ' '.join(sql.split())
|
|
self.statements.append((normalized, tuple(params or ())))
|
|
if 'pg_available_extensions' in normalized:
|
|
return Cursor([{'available': True}])
|
|
return Cursor()
|
|
|
|
connection = MigrationConnection()
|
|
scanner_db._migration_retain_pg_trgm_without_name_index(connection)
|
|
scanner_db._migration_retain_pg_trgm_without_name_index(connection)
|
|
sql = [statement for statement, _ in connection.statements]
|
|
self.assertEqual(sql.count('CREATE EXTENSION IF NOT EXISTS pg_trgm'), 2)
|
|
self.assertEqual(
|
|
sql.count('DROP INDEX IF EXISTS idx_package_repo_candidates_name_trgm'), 2,
|
|
)
|
|
self.assertFalse(any(
|
|
statement.startswith('CREATE INDEX idx_package_repo_candidates_name_trgm')
|
|
for statement in sql
|
|
))
|
|
|
|
def test_package_candidate_cache_keeps_committed_chunks_after_later_timeout(self):
|
|
self.assertEqual(scanner_db.PACKAGE_CANDIDATE_WRITE_CHUNK_SIZE, 25)
|
|
|
|
class ChunkConnection:
|
|
is_postgres = True
|
|
is_sqlite = False
|
|
|
|
def __init__(self):
|
|
self.insert_calls = []
|
|
self.pending = []
|
|
self.committed = []
|
|
self.commits = 0
|
|
self.rollbacks = 0
|
|
|
|
def execute(self, sql, params=None):
|
|
params = tuple(params or ())
|
|
if sql.lstrip().startswith('INSERT INTO package_repo_candidates'):
|
|
self.insert_calls.append((sql, params))
|
|
if len(self.insert_calls) == 3:
|
|
raise TimeoutError('statement timeout')
|
|
self.pending.append(params[1])
|
|
return Cursor(rowcount=1)
|
|
if sql == 'SELECT 1':
|
|
return Cursor([{'ok': 1}])
|
|
raise AssertionError('unexpected SQL')
|
|
|
|
def commit(self):
|
|
self.committed.extend(self.pending)
|
|
self.pending.clear()
|
|
self.commits += 1
|
|
|
|
def rollback(self):
|
|
self.pending.clear()
|
|
self.rollbacks += 1
|
|
|
|
candidates = [
|
|
{
|
|
'package_source': 'npm', 'name': f'candidate-{index}',
|
|
'version': '1', 'repo_url': f'https://example.invalid/{index}',
|
|
'provider': 'github', 'evidence': [], 'confidence': 'high',
|
|
}
|
|
for index in range(5)
|
|
]
|
|
connection = ChunkConnection()
|
|
db = self.claim_db(connection)
|
|
callback = console_runner.make_repo_candidate_callback(
|
|
db, run_id=7, cycle_id=9, query='fixture',
|
|
)
|
|
with mock.patch.object(
|
|
scanner_db, 'PACKAGE_CANDIDATE_WRITE_CHUNK_SIZE', 2,
|
|
), mock.patch.object(scanner_db.logger, 'warning') as warning:
|
|
report = callback(candidates)
|
|
|
|
self.assertEqual(connection.committed, ['candidate-0', 'candidate-1'])
|
|
self.assertEqual(connection.pending, [])
|
|
self.assertEqual(len(connection.insert_calls), 3)
|
|
self.assertEqual(connection.commits, 1)
|
|
self.assertEqual(connection.rollbacks, 1)
|
|
self.assertIn('ON CONFLICT(package_source, package_name, package_version, repo_url)', connection.insert_calls[0][0])
|
|
self.assertEqual(report, {
|
|
'offered_rows': 5,
|
|
'committed_rows': 2,
|
|
'skipped_rows': 3,
|
|
'failed': True,
|
|
'failed_chunk_rows': 2,
|
|
'connection_usable': True,
|
|
})
|
|
self.assertEqual(warning.call_count, 1)
|
|
self.assertNotIn('candidate-', str(warning.call_args))
|
|
self.assertEqual(connection.execute('SELECT 1').fetchone()['ok'], 1)
|
|
|
|
def test_lease_renewal_and_recheck_use_exact_owner_token_shapes(self):
|
|
class LeaseConnection:
|
|
is_postgres = True
|
|
|
|
def __init__(self):
|
|
self.calls = []
|
|
|
|
def execute(self, sql, params=None):
|
|
self.calls.append((' '.join(sql.split()), tuple(params or ())))
|
|
if sql.lstrip().startswith('SELECT lease_token'):
|
|
return Cursor([{'lease_token': 'token'}])
|
|
return Cursor(rowcount=1)
|
|
|
|
def commit(self):
|
|
return None
|
|
|
|
def rollback(self):
|
|
return None
|
|
|
|
connection = LeaseConnection()
|
|
db = self.claim_db(connection)
|
|
self.assertEqual(db.renew_target_leases('owner', 1800, ['token']), 1)
|
|
self.assertEqual(db.active_target_lease_tokens('owner', ['token']), {'token'})
|
|
self.assertEqual(len(connection.calls), 2)
|
|
for sql, params in connection.calls:
|
|
self.assertIn("status = 'in_progress'", sql)
|
|
self.assertIn('lease_owner = ?', sql)
|
|
self.assertIn('lease_token IN (?)', sql)
|
|
self.assertEqual(params[-2:], ('owner', 'token'))
|
|
|
|
def test_queue_observability_counts_are_bounded_and_use_stale_cache_on_timeout(self):
|
|
self.assertEqual(scanner_db.TARGET_QUEUE_OBSERVABILITY_STATUS_LIMIT, 100)
|
|
|
|
class CountConnection:
|
|
is_postgres = True
|
|
|
|
def __init__(self):
|
|
self.statements = []
|
|
self.execute_calls = 0
|
|
self.commits = 0
|
|
self.rollbacks = 0
|
|
self.fail = False
|
|
|
|
def execute(self, sql, params=None):
|
|
self.execute_calls += 1
|
|
if self.fail and not sql.startswith('SET LOCAL statement_timeout'):
|
|
raise TimeoutError('statement timeout')
|
|
params = tuple(params or ())
|
|
self.statements.append((sql, params))
|
|
if sql.startswith('SET LOCAL statement_timeout'):
|
|
return Cursor()
|
|
counts = {
|
|
'pending': scanner_db.TARGET_QUEUE_OBSERVABILITY_STATUS_LIMIT,
|
|
'done': 2,
|
|
}
|
|
return Cursor([{'count': counts.get(params[1], 0)}])
|
|
|
|
def commit(self):
|
|
self.commits += 1
|
|
|
|
def rollback(self):
|
|
self.rollbacks += 1
|
|
|
|
connection = CountConnection()
|
|
db = self.claim_db(connection)
|
|
clock = {'now': 100.0}
|
|
with mock.patch.object(scanner_db.time, 'monotonic', side_effect=lambda: clock['now']), \
|
|
mock.patch.object(scanner_db.logger, 'warning') as warning:
|
|
snapshot = db.target_queue_counts('package_git')
|
|
self.assertEqual(snapshot['counts']['pending'], scanner_db.TARGET_QUEUE_OBSERVABILITY_STATUS_LIMIT)
|
|
self.assertEqual(snapshot['counts']['done'], 2)
|
|
self.assertTrue(snapshot['degraded'])
|
|
self.assertFalse(snapshot['stale'])
|
|
self.assertEqual(snapshot['truncated_statuses'], ['pending'])
|
|
self.assertIn(
|
|
str(scanner_db.TARGET_QUEUE_OBSERVABILITY_TIMEOUT_MS),
|
|
connection.statements[0][0],
|
|
)
|
|
count_statements = connection.statements[1:]
|
|
self.assertEqual(len(count_statements), len(scanner_db.TARGET_QUEUE_OBSERVABILITY_STATUSES))
|
|
for sql, params in count_statements:
|
|
self.assertIn('LIMIT ?', sql)
|
|
self.assertEqual(params[-1], scanner_db.TARGET_QUEUE_OBSERVABILITY_STATUS_LIMIT)
|
|
self.assertEqual(connection.commits, 1)
|
|
|
|
connection.fail = True
|
|
clock['now'] = 101.0
|
|
stale = db.target_queue_counts('package_git')
|
|
self.assertEqual(stale['counts'], snapshot['counts'])
|
|
self.assertTrue(stale['degraded'])
|
|
self.assertTrue(stale['stale'])
|
|
self.assertEqual(stale['reason'], 'bounded_count_query_failed')
|
|
self.assertEqual(stale['retry_after_sec'], 300)
|
|
self.assertEqual(connection.rollbacks, 1)
|
|
self.assertEqual(warning.call_count, 1)
|
|
|
|
stale['counts']['pending'] = -1
|
|
stale['truncated_statuses'].append('failed')
|
|
connection.fail = False
|
|
calls_before_other_key = connection.execute_calls
|
|
other = db.target_queue_counts('github')
|
|
self.assertFalse(other['stale'])
|
|
self.assertGreater(connection.execute_calls, calls_before_other_key)
|
|
|
|
calls_before_backoff = connection.execute_calls
|
|
clock['now'] = 102.0
|
|
backed_off = db.target_queue_counts('package_git')
|
|
self.assertEqual(connection.execute_calls, calls_before_backoff)
|
|
self.assertEqual(warning.call_count, 1)
|
|
self.assertEqual(backed_off['reason'], 'bounded_count_retry_backoff')
|
|
self.assertEqual(
|
|
backed_off['counts']['pending'],
|
|
scanner_db.TARGET_QUEUE_OBSERVABILITY_STATUS_LIMIT,
|
|
)
|
|
self.assertEqual(backed_off['truncated_statuses'], ['pending'])
|
|
self.assertTrue(backed_off['stale'])
|
|
self.assertTrue(backed_off['degraded'])
|
|
|
|
clock['now'] = 402.0
|
|
refreshed = db.target_queue_counts('package_git')
|
|
self.assertFalse(refreshed['stale'])
|
|
self.assertEqual(refreshed['retry_after_sec'], 0)
|
|
self.assertNotIn('package_git', db._target_queue_counts_retry_after)
|
|
self.assertEqual(warning.call_count, 1)
|
|
|
|
def test_recent_operation_health_uses_status_scoped_bounded_queries(self):
|
|
class OperationConnection:
|
|
is_postgres = True
|
|
is_sqlite = False
|
|
|
|
def __init__(self):
|
|
self.statements = []
|
|
self.commits = 0
|
|
|
|
def execute(self, sql, params=None):
|
|
params = tuple(params or ())
|
|
self.statements.append((sql, params))
|
|
if sql.startswith('SET TRANSACTION ISOLATION LEVEL'):
|
|
return Cursor()
|
|
return Cursor([{
|
|
'operation_id': f'operation-{params[0]}',
|
|
'updated_at': f'2026-09-20T00:00:0{len(self.statements)}+00:00',
|
|
}])
|
|
|
|
def commit(self):
|
|
self.commits += 1
|
|
|
|
def rollback(self):
|
|
raise AssertionError('bounded operation query should not roll back')
|
|
|
|
connection = OperationConnection()
|
|
db = self.claim_db(connection)
|
|
db._runtime_operation_state = lambda row: dict(row)
|
|
|
|
result = db.recent_runtime_operations(limit=3)
|
|
|
|
self.assertEqual(len(connection.statements), 8)
|
|
self.assertEqual(
|
|
connection.statements[0][0],
|
|
'SET TRANSACTION ISOLATION LEVEL REPEATABLE READ READ ONLY',
|
|
)
|
|
for sql, params in connection.statements[1:]:
|
|
normalized = ' '.join(sql.split())
|
|
self.assertIn('WHERE status = ?', normalized)
|
|
self.assertIn('ORDER BY updated_at DESC, operation_id DESC LIMIT ?', normalized)
|
|
self.assertEqual(params[-1], 3)
|
|
self.assertEqual(connection.commits, 1)
|
|
self.assertEqual(len(result), 3)
|
|
|
|
connection.statements.clear()
|
|
before_at = '2026-09-20T00:00:05+00:00'
|
|
before_id = '00000000-0000-4000-8000-000000000050'
|
|
db.recent_runtime_operations(
|
|
limit=3, before_updated_at=before_at,
|
|
before_operation_id=before_id,
|
|
)
|
|
for sql, params in connection.statements[1:]:
|
|
normalized = ' '.join(sql.split())
|
|
self.assertIn(
|
|
'updated_at < ? OR ( updated_at = ? AND operation_id < ? )',
|
|
normalized,
|
|
)
|
|
self.assertEqual(params[1:4], (before_at, before_at, before_id))
|
|
self.assertEqual(params[-1], 3)
|
|
|
|
def test_queue_observability_backoff_without_cache_returns_bounded_empty_snapshot(self):
|
|
class FailedConnection:
|
|
is_postgres = True
|
|
|
|
def __init__(self):
|
|
self.execute_calls = 0
|
|
self.rollbacks = 0
|
|
|
|
def execute(self, sql, params=None):
|
|
self.execute_calls += 1
|
|
if sql.startswith('SET LOCAL statement_timeout'):
|
|
return Cursor()
|
|
raise TimeoutError('statement timeout')
|
|
|
|
def rollback(self):
|
|
self.rollbacks += 1
|
|
|
|
connection = FailedConnection()
|
|
db = self.claim_db(connection)
|
|
clock = {'now': 10.0}
|
|
with mock.patch.object(scanner_db.time, 'monotonic', side_effect=lambda: clock['now']), \
|
|
mock.patch.object(scanner_db.logger, 'warning') as warning:
|
|
first = db.target_queue_counts('source-a')
|
|
calls_after_failure = connection.execute_calls
|
|
clock['now'] = 11.0
|
|
second = db.target_queue_counts('source-a')
|
|
self.assertEqual(connection.rollbacks, 1)
|
|
self.assertEqual(connection.execute_calls, calls_after_failure)
|
|
self.assertEqual(warning.call_count, 1)
|
|
self.assertEqual(first['reason'], 'bounded_count_query_failed')
|
|
self.assertEqual(second['reason'], 'bounded_count_retry_backoff')
|
|
self.assertEqual(second['counts'], {})
|
|
self.assertEqual(second['sample_limit_per_status'], 100)
|
|
self.assertEqual(second['timeout_ms'], scanner_db.TARGET_QUEUE_OBSERVABILITY_TIMEOUT_MS)
|
|
self.assertEqual(second['truncated_statuses'], [])
|
|
self.assertTrue(second['stale'])
|
|
self.assertTrue(second['degraded'])
|
|
|
|
def test_admin_queue_health_uses_fixed_source_scoped_bounded_queries(self):
|
|
class CountConnection:
|
|
is_postgres = False
|
|
|
|
def __init__(self):
|
|
self.statements = []
|
|
|
|
def execute(self, sql, params=None):
|
|
params = tuple(params or ())
|
|
self.statements.append((sql, params))
|
|
return Cursor([{'count': 1 if params[1] == 'pending' else 0}])
|
|
|
|
def commit(self):
|
|
return None
|
|
|
|
def rollback(self):
|
|
return None
|
|
|
|
connection = CountConnection()
|
|
db = self.claim_db(connection)
|
|
snapshot = db.admin_target_queue_health(
|
|
('gitlab', 'dockerhub', 'huggingface'),
|
|
)
|
|
|
|
self.assertEqual(snapshot['counts'], {'pending': 3})
|
|
self.assertEqual(snapshot['sample_limit_per_status'], 300)
|
|
self.assertEqual(
|
|
len(connection.statements),
|
|
3 * len(scanner_db.TARGET_QUEUE_OBSERVABILITY_STATUSES),
|
|
)
|
|
for sql, params in connection.statements:
|
|
self.assertIn('WHERE source = ? AND status = ?', sql)
|
|
self.assertIn('LIMIT ?', sql)
|
|
self.assertIn(params[0], ('gitlab', 'dockerhub', 'huggingface'))
|
|
self.assertEqual(params[-1], scanner_db.TARGET_QUEUE_OBSERVABILITY_STATUS_LIMIT)
|
|
|
|
def test_windows_console_fallback_does_not_mutate_unicode_log_text(self):
|
|
class Cp1251Stream:
|
|
encoding = 'cp1251'
|
|
|
|
original = 'keychecks launch \U0001f680'
|
|
rendered = supervisor.console_safe_text(original, Cp1251Stream())
|
|
rendered.encode('cp1251')
|
|
self.assertEqual(original, 'keychecks launch \U0001f680')
|
|
self.assertIn('\\U0001f680', rendered)
|
|
|
|
|
|
if __name__ == '__main__':
|
|
unittest.main()
|