81 lines
5.0 KiB
Python
81 lines
5.0 KiB
Python
"""Verify retirement and full news-AI reset only inside the disposable CI database."""
|
|
import asyncio
|
|
from datetime import UTC, datetime
|
|
|
|
from sqlalchemy import DateTime, Integer, MetaData, String, insert, text
|
|
from sqlalchemy.dialects.postgresql import JSONB
|
|
|
|
from alphapulse.core.config import get_settings
|
|
from alphapulse.core.db import Base, engine
|
|
from alphapulse.crawler.remove_cls import AI_TABLES, OPTIONAL_AI_TABLES, cleanup
|
|
from alphapulse.models import NewsItem, Source
|
|
import alphapulse.models.risk # noqa: F401
|
|
import alphapulse.news_research.models # noqa: F401
|
|
|
|
|
|
async def main():
|
|
assert get_settings().database_url.startswith('postgresql+asyncpg://ci:ci-only@')
|
|
assert get_settings().database_url.endswith('/ci')
|
|
try:
|
|
async with engine.connect() as connection:
|
|
transaction = await connection.begin()
|
|
try:
|
|
assert not await connection.scalar(text("SELECT count(*) FROM sources WHERE code IN ('cls_depth','cls_telegraph')"))
|
|
await connection.execute(insert(Source).values(code='cls_depth', name='CI CLS', list_url='https://example.com', interval_seconds=120))
|
|
for identity, source in ((9000001, 'cls_depth'), (9000002, 'wscn_live')):
|
|
await connection.execute(insert(NewsItem).values(id=identity, source_code=source, external_id='cleanup-ci',
|
|
url='https://example.com', published_at=datetime.now(UTC), content_hash=str(identity)))
|
|
assert set((await cleanup(connection))['absent_legacy_tables']) == set(OPTIONAL_AI_TABLES)
|
|
for name in OPTIONAL_AI_TABLES:
|
|
await connection.run_sync(lambda sync, name=name: Base.metadata.tables[name].create(sync))
|
|
reflected = MetaData()
|
|
await connection.run_sync(lambda sync: reflected.reflect(sync, only=list((*AI_TABLES, *OPTIONAL_AI_TABLES))))
|
|
for table in reflected.sorted_tables:
|
|
if table.name not in (*AI_TABLES, *OPTIONAL_AI_TABLES):
|
|
continue
|
|
values = {}
|
|
for column in table.columns:
|
|
if not column.primary_key and (column.nullable or column.default is not None or column.server_default is not None):
|
|
continue
|
|
if isinstance(column.type, Integer):
|
|
values[column.name] = 9000001
|
|
elif isinstance(column.type, DateTime):
|
|
values[column.name] = datetime.now(UTC)
|
|
elif isinstance(column.type, String):
|
|
values[column.name] = 'ci'
|
|
elif isinstance(column.type, JSONB):
|
|
values[column.name] = {}
|
|
else:
|
|
raise AssertionError(column)
|
|
await connection.execute(table.insert().values(**values))
|
|
await connection.execute(text("INSERT INTO raw_bodies VALUES (:shared,:body),(:exclusive,:body)"),
|
|
{'shared': 'a'*64, 'exclusive': 'b'*64, 'body': b'ci'})
|
|
for identity, source, digest in ((9000001,'cls_depth','a'*64),(9000002,'wscn_live','a'*64),(9000003,'cls_depth','b'*64)):
|
|
await connection.execute(text('''INSERT INTO raw_responses
|
|
(id,source_code,url,method,status,headers,body_sha256) VALUES
|
|
(:id,:source,'https://example.com','GET',200,'{}',:digest)'''),
|
|
{'id':identity,'source':source,'digest':digest})
|
|
preview = await cleanup(connection)
|
|
assert not preview['executed'] and preview['counts']['sources'] == 1
|
|
savepoint = await connection.begin_nested()
|
|
await cleanup(connection, True)
|
|
await savepoint.rollback()
|
|
assert await connection.scalar(text("SELECT count(*) FROM news_items WHERE id=9000001")) == 1
|
|
result = await cleanup(connection, True)
|
|
assert result['executed']
|
|
for table in (*AI_TABLES, *OPTIONAL_AI_TABLES):
|
|
assert await connection.scalar(text('SELECT count(*) FROM '+table)) == 0, table
|
|
assert await connection.scalar(text("SELECT count(*) FROM news_items WHERE id=9000001")) == 0
|
|
assert await connection.scalar(text("SELECT count(*) FROM news_items WHERE id=9000002")) == 1
|
|
assert await connection.scalar(text("SELECT count(*) FROM raw_bodies WHERE sha256=:digest"), {'digest':'a'*64}) == 1
|
|
assert await connection.scalar(text("SELECT count(*) FROM raw_bodies WHERE sha256=:digest"), {'digest':'b'*64}) == 0
|
|
assert (await cleanup(connection, True))['counts']['sources'] == 0
|
|
print('CLS removal: preview, rollback, shared bodies, retained originals, 26 AI tables and repeat execution passed')
|
|
finally:
|
|
await transaction.rollback()
|
|
finally:
|
|
await engine.dispose()
|
|
|
|
|
|
asyncio.run(main())
|