First Commit
This commit is contained in:
@@ -0,0 +1,241 @@
|
||||
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()
|
||||
Reference in New Issue
Block a user