242 lines
8.1 KiB
Python
242 lines
8.1 KiB
Python
import sqlite3
|
|
from pathlib import Path
|
|
from datetime import datetime, timezone
|
|
from .models import Paper
|
|
from .util import paper_key
|
|
|
|
SCHEMA = '''
|
|
CREATE TABLE IF NOT EXISTS papers (
|
|
paper_key TEXT PRIMARY KEY,
|
|
doi TEXT,
|
|
title TEXT NOT NULL,
|
|
authors TEXT,
|
|
abstract TEXT,
|
|
url TEXT,
|
|
venue TEXT,
|
|
publication_date TEXT,
|
|
year INTEGER,
|
|
citation_count INTEGER,
|
|
source TEXT,
|
|
source_id TEXT,
|
|
categories TEXT,
|
|
relevance INTEGER DEFAULT 0,
|
|
summary TEXT,
|
|
ai_reason TEXT,
|
|
first_seen TEXT NOT NULL,
|
|
first_seen_local_date TEXT,
|
|
last_seen TEXT NOT NULL
|
|
);
|
|
CREATE INDEX IF NOT EXISTS idx_papers_pubdate ON papers(publication_date);
|
|
CREATE INDEX IF NOT EXISTS idx_papers_relevance ON papers(relevance);
|
|
CREATE INDEX IF NOT EXISTS idx_papers_first_seen_local_date ON papers(first_seen_local_date);
|
|
'''
|
|
|
|
class PaperDB:
|
|
def __init__(self, path):
|
|
Path(path).parent.mkdir(parents=True, exist_ok=True)
|
|
self.conn = sqlite3.connect(path)
|
|
self.conn.row_factory = sqlite3.Row
|
|
self.conn.executescript(SCHEMA)
|
|
self._migrate()
|
|
|
|
def _migrate(self):
|
|
cols = {
|
|
r['name']
|
|
for r in self.conn.execute(
|
|
'PRAGMA table_info(papers)'
|
|
).fetchall()
|
|
}
|
|
|
|
if 'first_seen_local_date' not in cols:
|
|
self.conn.execute(
|
|
'ALTER TABLE papers ADD COLUMN first_seen_local_date TEXT'
|
|
)
|
|
|
|
if 'oa_status' not in cols:
|
|
self.conn.execute(
|
|
'ALTER TABLE papers ADD COLUMN oa_status TEXT'
|
|
)
|
|
|
|
if 'pdf_url' not in cols:
|
|
self.conn.execute(
|
|
'ALTER TABLE papers ADD COLUMN pdf_url TEXT'
|
|
)
|
|
|
|
if 'ai_analysis_level' not in cols:
|
|
self.conn.execute(
|
|
'ALTER TABLE papers ADD COLUMN ai_analysis_level TEXT'
|
|
)
|
|
|
|
if 'ai_status' not in cols:
|
|
self.conn.execute(
|
|
'ALTER TABLE papers ADD COLUMN ai_status TEXT'
|
|
)
|
|
|
|
self.conn.commit()
|
|
|
|
def exists(self, p: Paper) -> bool:
|
|
k = paper_key(p.doi, p.title)
|
|
row = self.conn.execute(
|
|
'SELECT 1 FROM papers WHERE paper_key=? LIMIT 1',
|
|
(k,)
|
|
).fetchone()
|
|
return row is not None
|
|
|
|
def upsert(self, p: Paper, local_date: str = ''):
|
|
k = paper_key(p.doi, p.title)
|
|
now = datetime.now(timezone.utc).isoformat()
|
|
old = self.conn.execute('SELECT * FROM papers WHERE paper_key=?', (k,)).fetchone()
|
|
is_new = old is None
|
|
authors, cats = '|||'.join(p.authors), '|||'.join(p.categories)
|
|
if is_new:
|
|
self.conn.execute('''INSERT INTO papers
|
|
(paper_key,doi,title,authors,abstract,url,venue,publication_date,year,citation_count,source,source_id,categories,relevance,summary,ai_reason,oa_status,pdf_url,ai_analysis_level,ai_status,first_seen,first_seen_local_date,last_seen)
|
|
VALUES (?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?)''',
|
|
(
|
|
k,
|
|
p.doi,
|
|
p.title,
|
|
authors,
|
|
p.abstract,
|
|
p.url,
|
|
p.venue,
|
|
p.publication_date,
|
|
p.year,
|
|
p.citation_count,
|
|
p.source,
|
|
p.source_id,
|
|
cats,
|
|
p.relevance,
|
|
p.summary,
|
|
p.ai_reason,
|
|
p.oa_status,
|
|
p.pdf_url,
|
|
p.ai_analysis_level,
|
|
p.ai_status,
|
|
now,
|
|
local_date,
|
|
now,
|
|
))
|
|
else:
|
|
def choose(new, oldv): return new if new not in ('', None, [], 0) else oldv
|
|
self.conn.execute('''UPDATE papers SET
|
|
doi=?,
|
|
title=?,
|
|
authors=?,
|
|
abstract=?,
|
|
url=?,
|
|
venue=?,
|
|
publication_date=?,
|
|
year=?,
|
|
citation_count=?,
|
|
source=?,
|
|
source_id=?,
|
|
categories=?,
|
|
relevance=?,
|
|
summary=?,
|
|
ai_reason=?,
|
|
oa_status=?,
|
|
pdf_url=?,
|
|
ai_analysis_level=?,
|
|
ai_status=?,
|
|
last_seen=?
|
|
WHERE paper_key=?''',
|
|
(
|
|
choose(p.doi, old['doi']),
|
|
choose(p.title, old['title']),
|
|
choose(authors, old['authors']),
|
|
choose(p.abstract, old['abstract']),
|
|
choose(p.url, old['url']),
|
|
choose(p.venue, old['venue']),
|
|
choose(p.publication_date, old['publication_date']),
|
|
choose(p.year, old['year']),
|
|
choose(p.citation_count, old['citation_count']),
|
|
choose(p.source, old['source']),
|
|
choose(p.source_id, old['source_id']),
|
|
choose(cats, old['categories']),
|
|
max(p.relevance, old['relevance'] or 0),
|
|
choose(p.summary, old['summary']),
|
|
choose(p.ai_reason, old['ai_reason']),
|
|
choose(p.oa_status, old['oa_status']),
|
|
choose(p.pdf_url, old['pdf_url']),
|
|
choose(p.ai_analysis_level, old['ai_analysis_level']),
|
|
choose(p.ai_status, old['ai_status']),
|
|
now,
|
|
k,
|
|
))
|
|
self.conn.commit()
|
|
return is_new
|
|
|
|
def list_first_seen_on(self, local_date: str) -> list[Paper]:
|
|
rows = self.conn.execute('SELECT * FROM papers WHERE first_seen_local_date=? ORDER BY relevance DESC, publication_date DESC', (local_date,)).fetchall()
|
|
out = []
|
|
for r in rows:
|
|
out.append(Paper(
|
|
title=r['title'],
|
|
authors=(r['authors'] or '').split('|||') if r['authors'] else [],
|
|
abstract=r['abstract'] or '',
|
|
doi=r['doi'] or '',
|
|
url=r['url'] or '',
|
|
venue=r['venue'] or '',
|
|
publication_date=r['publication_date'] or '',
|
|
year=r['year'],
|
|
citation_count=r['citation_count'],
|
|
source=r['source'] or '',
|
|
source_id=r['source_id'] or '',
|
|
categories=(r['categories'] or '').split('|||') if r['categories'] else [],
|
|
relevance=r['relevance'] or 0,
|
|
summary=r['summary'] or '',
|
|
ai_reason=r['ai_reason'] or '',
|
|
oa_status=r['oa_status'] or '',
|
|
pdf_url=r['pdf_url'] or '',
|
|
ai_analysis_level=r['ai_analysis_level'] or '',
|
|
ai_status=r['ai_status'] or '',
|
|
))
|
|
return out
|
|
|
|
def list_ai_failed(self, limit: int = 20) -> list[Paper]:
|
|
rows = self.conn.execute(
|
|
'''
|
|
SELECT *
|
|
FROM papers
|
|
WHERE ai_status='failed'
|
|
ORDER BY last_seen DESC
|
|
LIMIT ?
|
|
''',
|
|
(limit,)
|
|
).fetchall()
|
|
|
|
out = []
|
|
|
|
for r in rows:
|
|
out.append(
|
|
Paper(
|
|
title=r['title'],
|
|
authors=(r['authors'] or '').split('|||')
|
|
if r['authors'] else [],
|
|
abstract=r['abstract'] or '',
|
|
doi=r['doi'] or '',
|
|
url=r['url'] or '',
|
|
venue=r['venue'] or '',
|
|
publication_date=r['publication_date'] or '',
|
|
year=r['year'],
|
|
citation_count=r['citation_count'],
|
|
source=r['source'] or '',
|
|
source_id=r['source_id'] or '',
|
|
categories=(r['categories'] or '').split('|||')
|
|
if r['categories'] else [],
|
|
relevance=r['relevance'] or 0,
|
|
summary=r['summary'] or '',
|
|
ai_reason=r['ai_reason'] or '',
|
|
oa_status=r['oa_status'] or '',
|
|
pdf_url=r['pdf_url'] or '',
|
|
ai_analysis_level=r['ai_analysis_level'] or '',
|
|
ai_status=r['ai_status'] or '',
|
|
)
|
|
)
|
|
|
|
return out
|
|
|
|
def close(self):
|
|
self.conn.close()
|