938 lines
32 KiB
Python
938 lines
32 KiB
Python
"""Database client for storing and retrieving LLM messages and user authentication."""
|
|
|
|
import contextlib
|
|
import hashlib
|
|
import json
|
|
from datetime import datetime, timedelta
|
|
from typing import Dict, List, Optional
|
|
|
|
import bcrypt
|
|
import psycopg2
|
|
from psycopg2.extras import RealDictCursor
|
|
from psycopg2.pool import ThreadedConnectionPool
|
|
|
|
|
|
class DatabaseClient:
|
|
"""Handles database operations for message storage and retrieval."""
|
|
|
|
def __init__(self, database_url: str, minconn: int = 2, maxconn: int = 10):
|
|
"""Initialize the database client.
|
|
|
|
Args:
|
|
database_url: PostgreSQL connection string
|
|
minconn: Minimum connections in the pool
|
|
maxconn: Maximum connections in the pool
|
|
"""
|
|
self.database_url = database_url
|
|
self._pool: ThreadedConnectionPool | None = None
|
|
self._minconn = minconn
|
|
self._maxconn = maxconn
|
|
# Legacy single connection kept for backwards compatibility
|
|
self.conn = None
|
|
|
|
def _ensure_pool(self):
|
|
"""Create the connection pool if it doesn't exist yet."""
|
|
if self._pool is None or self._pool.closed:
|
|
self._pool = ThreadedConnectionPool(
|
|
self._minconn, self._maxconn, self.database_url
|
|
)
|
|
|
|
@contextlib.contextmanager
|
|
def get_conn(self):
|
|
"""Check out a connection from the pool. Returns it on exit."""
|
|
self._ensure_pool()
|
|
conn = self._pool.getconn()
|
|
try:
|
|
yield conn
|
|
finally:
|
|
self._pool.putconn(conn)
|
|
|
|
def connect(self):
|
|
"""Establish database connection (legacy single-connection path)."""
|
|
if not self.conn or self.conn.closed:
|
|
self.conn = psycopg2.connect(self.database_url)
|
|
|
|
def close(self):
|
|
"""Close database connection and pool."""
|
|
if self.conn and not self.conn.closed:
|
|
self.conn.close()
|
|
if self._pool and not self._pool.closed:
|
|
self._pool.closeall()
|
|
|
|
def initialize_schema(self):
|
|
"""Create database tables if they don't exist."""
|
|
self.connect()
|
|
|
|
with self.conn.cursor() as cursor:
|
|
# Create messages table
|
|
cursor.execute("""
|
|
CREATE TABLE IF NOT EXISTS llm_messages (
|
|
id SERIAL PRIMARY KEY,
|
|
timestamp TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
|
|
company_name VARCHAR(255),
|
|
analysis_type VARCHAR(50),
|
|
model VARCHAR(100),
|
|
prompt TEXT NOT NULL,
|
|
prompt_hash VARCHAR(64),
|
|
response TEXT,
|
|
metadata JSONB,
|
|
token_usage JSONB,
|
|
created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
|
|
is_cached BOOLEAN DEFAULT FALSE
|
|
)
|
|
""")
|
|
|
|
# Create index on timestamp for analytics queries
|
|
cursor.execute("""
|
|
CREATE INDEX IF NOT EXISTS idx_messages_timestamp
|
|
ON llm_messages(timestamp)
|
|
""")
|
|
|
|
# Create index on company_name for filtering
|
|
cursor.execute("""
|
|
CREATE INDEX IF NOT EXISTS idx_messages_company
|
|
ON llm_messages(company_name)
|
|
""")
|
|
|
|
# Add prompt_hash and is_cached columns if they don't exist (for existing tables)
|
|
# This must run BEFORE creating the index on prompt_hash
|
|
cursor.execute("""
|
|
DO $$
|
|
BEGIN
|
|
IF NOT EXISTS (
|
|
SELECT 1 FROM information_schema.columns
|
|
WHERE table_name = 'llm_messages' AND column_name = 'prompt_hash'
|
|
) THEN
|
|
ALTER TABLE llm_messages ADD COLUMN prompt_hash VARCHAR(64);
|
|
END IF;
|
|
IF NOT EXISTS (
|
|
SELECT 1 FROM information_schema.columns
|
|
WHERE table_name = 'llm_messages' AND column_name = 'is_cached'
|
|
) THEN
|
|
ALTER TABLE llm_messages ADD COLUMN is_cached BOOLEAN DEFAULT FALSE;
|
|
END IF;
|
|
END $$;
|
|
""")
|
|
|
|
# Create index on prompt_hash for cache lookups
|
|
cursor.execute("""
|
|
CREATE INDEX IF NOT EXISTS idx_messages_prompt_hash
|
|
ON llm_messages(prompt_hash)
|
|
""")
|
|
|
|
# Create users table for authentication
|
|
cursor.execute("""
|
|
CREATE TABLE IF NOT EXISTS users (
|
|
id SERIAL PRIMARY KEY,
|
|
email VARCHAR(255) UNIQUE NOT NULL,
|
|
password_hash VARCHAR(255) NOT NULL,
|
|
role VARCHAR(20) DEFAULT 'user' CHECK (role IN ('admin', 'user')),
|
|
created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
|
|
updated_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
|
|
)
|
|
""")
|
|
|
|
# Create index on email for fast lookups
|
|
cursor.execute("""
|
|
CREATE INDEX IF NOT EXISTS idx_users_email
|
|
ON users(email)
|
|
""")
|
|
|
|
# Create patents cache table
|
|
cursor.execute("""
|
|
CREATE TABLE IF NOT EXISTS patents (
|
|
patent_id VARCHAR(64) PRIMARY KEY,
|
|
company_name VARCHAR(255),
|
|
pdf_link TEXT,
|
|
raw_sections JSONB,
|
|
minimized_content TEXT,
|
|
created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
|
|
)
|
|
""")
|
|
|
|
cursor.execute("""
|
|
CREATE INDEX IF NOT EXISTS idx_patents_company
|
|
ON patents(company_name)
|
|
""")
|
|
|
|
# Create SERP query cache table
|
|
cursor.execute("""
|
|
CREATE TABLE IF NOT EXISTS serp_queries (
|
|
id SERIAL PRIMARY KEY,
|
|
company_name VARCHAR(255),
|
|
query_hash VARCHAR(64) UNIQUE,
|
|
result_patent_ids TEXT[],
|
|
expires_at TIMESTAMP NOT NULL,
|
|
created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
|
|
)
|
|
""")
|
|
|
|
cursor.execute("""
|
|
CREATE INDEX IF NOT EXISTS idx_serp_queries_hash
|
|
ON serp_queries(query_hash)
|
|
""")
|
|
|
|
# Create jobs table for persisting async batch job state
|
|
cursor.execute("""
|
|
CREATE TABLE IF NOT EXISTS jobs (
|
|
job_id VARCHAR(128) PRIMARY KEY,
|
|
status VARCHAR(20) NOT NULL DEFAULT 'pending',
|
|
progress INTEGER NOT NULL DEFAULT 0,
|
|
total_companies INTEGER NOT NULL DEFAULT 0,
|
|
completed_companies INTEGER NOT NULL DEFAULT 0,
|
|
result_json JSONB,
|
|
error TEXT,
|
|
created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
|
|
updated_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
|
|
)
|
|
""")
|
|
|
|
cursor.execute("""
|
|
CREATE INDEX IF NOT EXISTS idx_jobs_status
|
|
ON jobs(status)
|
|
""")
|
|
|
|
# Create tracked companies table for scheduled analysis
|
|
cursor.execute("""
|
|
CREATE TABLE IF NOT EXISTS tracked_companies (
|
|
id SERIAL PRIMARY KEY,
|
|
company_name VARCHAR(255) UNIQUE NOT NULL,
|
|
last_patent_count INTEGER DEFAULT 0,
|
|
last_analysis_at TIMESTAMP,
|
|
created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
|
|
)
|
|
""")
|
|
|
|
# Create alerts table for significant changes
|
|
cursor.execute("""
|
|
CREATE TABLE IF NOT EXISTS alerts (
|
|
id SERIAL PRIMARY KEY,
|
|
company_name VARCHAR(255) NOT NULL,
|
|
alert_type VARCHAR(50) NOT NULL,
|
|
message TEXT NOT NULL,
|
|
old_value NUMERIC,
|
|
new_value NUMERIC,
|
|
created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
|
|
)
|
|
""")
|
|
|
|
cursor.execute("""
|
|
CREATE INDEX IF NOT EXISTS idx_alerts_company
|
|
ON alerts(company_name)
|
|
""")
|
|
|
|
self.conn.commit()
|
|
|
|
@staticmethod
|
|
def hash_prompt(prompt: str) -> str:
|
|
"""Generate a hash of the prompt for cache lookups.
|
|
|
|
Args:
|
|
prompt: The prompt text to hash
|
|
|
|
Returns:
|
|
SHA-256 hash of the prompt
|
|
"""
|
|
return hashlib.sha256(prompt.encode()).hexdigest()
|
|
|
|
def get_cached_response(
|
|
self,
|
|
prompt: str,
|
|
company_name: Optional[str] = None,
|
|
analysis_type: Optional[str] = None,
|
|
) -> Optional[Dict]:
|
|
"""Look up a cached response for a given prompt.
|
|
|
|
Args:
|
|
prompt: The prompt to look up
|
|
company_name: Optional company name filter
|
|
analysis_type: Optional analysis type filter
|
|
|
|
Returns:
|
|
Cached message dict if found, None otherwise
|
|
"""
|
|
prompt_hash = self.hash_prompt(prompt)
|
|
|
|
query = """
|
|
SELECT * FROM llm_messages
|
|
WHERE prompt_hash = %s
|
|
AND response IS NOT NULL
|
|
AND response NOT LIKE '[DATABASE MODE]%%'
|
|
AND response NOT LIKE '[TEST MODE]%%'
|
|
AND response NOT LIKE '[NO API]%%'
|
|
"""
|
|
params = [prompt_hash]
|
|
|
|
if company_name:
|
|
query += " AND company_name = %s"
|
|
params.append(company_name)
|
|
|
|
if analysis_type:
|
|
query += " AND analysis_type = %s"
|
|
params.append(analysis_type)
|
|
|
|
query += " ORDER BY timestamp DESC LIMIT 1"
|
|
|
|
with self.get_conn() as conn:
|
|
with conn.cursor(cursor_factory=RealDictCursor) as cursor:
|
|
cursor.execute(query, params)
|
|
result = cursor.fetchone()
|
|
return dict(result) if result else None
|
|
|
|
def store_message(
|
|
self,
|
|
prompt: str,
|
|
response: str,
|
|
company_name: Optional[str] = None,
|
|
analysis_type: Optional[str] = None,
|
|
model: Optional[str] = None,
|
|
metadata: Optional[Dict] = None,
|
|
token_usage: Optional[Dict] = None,
|
|
is_cached: bool = False,
|
|
) -> int:
|
|
"""Store an LLM message exchange in the database.
|
|
|
|
Args:
|
|
prompt: The prompt sent to the LLM
|
|
response: The response from the LLM
|
|
company_name: Name of company being analyzed
|
|
analysis_type: Type of analysis (e.g., 'single_patent', 'portfolio')
|
|
model: Model identifier used
|
|
metadata: Additional metadata as dict
|
|
token_usage: Token usage information
|
|
is_cached: Whether this response was served from cache
|
|
|
|
Returns:
|
|
The ID of the inserted record
|
|
"""
|
|
prompt_hash = self.hash_prompt(prompt)
|
|
|
|
with self.get_conn() as conn:
|
|
with conn.cursor() as cursor:
|
|
cursor.execute(
|
|
"""
|
|
INSERT INTO llm_messages
|
|
(prompt, prompt_hash, response, company_name, analysis_type, model, metadata, token_usage, is_cached)
|
|
VALUES (%s, %s, %s, %s, %s, %s, %s, %s, %s)
|
|
RETURNING id
|
|
""",
|
|
(
|
|
prompt,
|
|
prompt_hash,
|
|
response,
|
|
company_name,
|
|
analysis_type,
|
|
model,
|
|
json.dumps(metadata) if metadata else None,
|
|
json.dumps(token_usage) if token_usage else None,
|
|
is_cached,
|
|
),
|
|
)
|
|
|
|
message_id = cursor.fetchone()[0]
|
|
conn.commit()
|
|
|
|
return message_id
|
|
|
|
def get_messages(
|
|
self,
|
|
company_name: Optional[str] = None,
|
|
analysis_type: Optional[str] = None,
|
|
limit: int = 100,
|
|
offset: int = 0,
|
|
) -> List[Dict]:
|
|
"""Retrieve messages from the database.
|
|
|
|
Args:
|
|
company_name: Filter by company name
|
|
analysis_type: Filter by analysis type
|
|
limit: Maximum number of records to return
|
|
offset: Number of records to skip
|
|
|
|
Returns:
|
|
List of message dictionaries
|
|
"""
|
|
query = "SELECT * FROM llm_messages WHERE 1=1"
|
|
params = []
|
|
|
|
if company_name:
|
|
query += " AND company_name = %s"
|
|
params.append(company_name)
|
|
|
|
if analysis_type:
|
|
query += " AND analysis_type = %s"
|
|
params.append(analysis_type)
|
|
|
|
query += " ORDER BY timestamp DESC LIMIT %s OFFSET %s"
|
|
params.extend([limit, offset])
|
|
|
|
with self.get_conn() as conn:
|
|
with conn.cursor(cursor_factory=RealDictCursor) as cursor:
|
|
cursor.execute(query, params)
|
|
return [dict(row) for row in cursor.fetchall()]
|
|
|
|
def get_analytics(self, days: int = 30) -> Dict:
|
|
"""Get analytics on message usage.
|
|
|
|
Args:
|
|
days: Number of days to look back
|
|
|
|
Returns:
|
|
Dictionary with analytics data
|
|
"""
|
|
with self.get_conn() as conn:
|
|
with conn.cursor(cursor_factory=RealDictCursor) as cursor:
|
|
# Total messages
|
|
cursor.execute(
|
|
"""
|
|
SELECT COUNT(*) as total_messages
|
|
FROM llm_messages
|
|
WHERE timestamp >= NOW() - INTERVAL '%s days'
|
|
""",
|
|
(days,),
|
|
)
|
|
total = cursor.fetchone()["total_messages"]
|
|
|
|
# Messages by company
|
|
cursor.execute(
|
|
"""
|
|
SELECT company_name, COUNT(*) as count
|
|
FROM llm_messages
|
|
WHERE timestamp >= NOW() - INTERVAL '%s days'
|
|
GROUP BY company_name
|
|
ORDER BY count DESC
|
|
LIMIT 10
|
|
""",
|
|
(days,),
|
|
)
|
|
by_company = cursor.fetchall()
|
|
|
|
# Messages by type
|
|
cursor.execute(
|
|
"""
|
|
SELECT analysis_type, COUNT(*) as count
|
|
FROM llm_messages
|
|
WHERE timestamp >= NOW() - INTERVAL '%s days'
|
|
GROUP BY analysis_type
|
|
ORDER BY count DESC
|
|
""",
|
|
(days,),
|
|
)
|
|
by_type = cursor.fetchall()
|
|
|
|
return {
|
|
"total_messages": total,
|
|
"by_company": [dict(row) for row in by_company],
|
|
"by_type": [dict(row) for row in by_type],
|
|
"period_days": days,
|
|
}
|
|
|
|
# Patent Cache Methods
|
|
|
|
def get_cached_patent(self, patent_id: str) -> Optional[Dict]:
|
|
"""Look up a cached patent by ID.
|
|
|
|
Returns:
|
|
Dict with raw_sections and minimized_content, or None.
|
|
"""
|
|
with self.get_conn() as conn:
|
|
with conn.cursor(cursor_factory=RealDictCursor) as cursor:
|
|
cursor.execute(
|
|
"SELECT * FROM patents WHERE patent_id = %s",
|
|
(patent_id,),
|
|
)
|
|
row = cursor.fetchone()
|
|
return dict(row) if row else None
|
|
|
|
def store_patent(
|
|
self,
|
|
patent_id: str,
|
|
company_name: str,
|
|
pdf_link: str,
|
|
raw_sections: Dict,
|
|
minimized_content: str,
|
|
) -> None:
|
|
"""Store a processed patent in the cache."""
|
|
with self.get_conn() as conn:
|
|
with conn.cursor() as cursor:
|
|
cursor.execute(
|
|
"""
|
|
INSERT INTO patents (patent_id, company_name, pdf_link, raw_sections, minimized_content)
|
|
VALUES (%s, %s, %s, %s, %s)
|
|
ON CONFLICT (patent_id) DO UPDATE SET
|
|
raw_sections = EXCLUDED.raw_sections,
|
|
minimized_content = EXCLUDED.minimized_content
|
|
""",
|
|
(patent_id, company_name, pdf_link, json.dumps(raw_sections), minimized_content),
|
|
)
|
|
conn.commit()
|
|
|
|
def get_cached_serp_query(self, query_hash: str) -> Optional[List[str]]:
|
|
"""Look up cached SERP query results.
|
|
|
|
Returns:
|
|
List of patent IDs if cache hit and not expired, None otherwise.
|
|
"""
|
|
with self.get_conn() as conn:
|
|
with conn.cursor(cursor_factory=RealDictCursor) as cursor:
|
|
cursor.execute(
|
|
"""
|
|
SELECT result_patent_ids FROM serp_queries
|
|
WHERE query_hash = %s AND expires_at > NOW()
|
|
""",
|
|
(query_hash,),
|
|
)
|
|
row = cursor.fetchone()
|
|
return row["result_patent_ids"] if row else None
|
|
|
|
def store_serp_query(
|
|
self,
|
|
company_name: str,
|
|
query_hash: str,
|
|
patent_ids: List[str],
|
|
ttl_hours: int = 24,
|
|
) -> None:
|
|
"""Store SERP query results in the cache."""
|
|
expires_at = datetime.now() + timedelta(hours=ttl_hours)
|
|
with self.get_conn() as conn:
|
|
with conn.cursor() as cursor:
|
|
cursor.execute(
|
|
"""
|
|
INSERT INTO serp_queries (company_name, query_hash, result_patent_ids, expires_at)
|
|
VALUES (%s, %s, %s, %s)
|
|
ON CONFLICT (query_hash) DO UPDATE SET
|
|
result_patent_ids = EXCLUDED.result_patent_ids,
|
|
expires_at = EXCLUDED.expires_at
|
|
""",
|
|
(company_name, query_hash, patent_ids, expires_at),
|
|
)
|
|
conn.commit()
|
|
|
|
# Job Persistence Methods
|
|
|
|
def create_job(
|
|
self,
|
|
job_id: str,
|
|
total_companies: int,
|
|
) -> Dict:
|
|
"""Create a new job record.
|
|
|
|
Args:
|
|
job_id: Unique job identifier
|
|
total_companies: Number of companies in the batch
|
|
|
|
Returns:
|
|
Job dict
|
|
"""
|
|
with self.get_conn() as conn:
|
|
with conn.cursor(cursor_factory=RealDictCursor) as cursor:
|
|
cursor.execute(
|
|
"""
|
|
INSERT INTO jobs (job_id, status, progress, total_companies, completed_companies)
|
|
VALUES (%s, 'pending', 0, %s, 0)
|
|
RETURNING *
|
|
""",
|
|
(job_id, total_companies),
|
|
)
|
|
job = cursor.fetchone()
|
|
conn.commit()
|
|
return dict(job)
|
|
|
|
def update_job(
|
|
self,
|
|
job_id: str,
|
|
status: Optional[str] = None,
|
|
progress: Optional[int] = None,
|
|
completed_companies: Optional[int] = None,
|
|
result_json: Optional[str] = None,
|
|
error: Optional[str] = None,
|
|
) -> Optional[Dict]:
|
|
"""Update a job's state.
|
|
|
|
Only non-None fields are updated.
|
|
"""
|
|
updates = []
|
|
params = []
|
|
if status is not None:
|
|
updates.append("status = %s")
|
|
params.append(status)
|
|
if progress is not None:
|
|
updates.append("progress = %s")
|
|
params.append(progress)
|
|
if completed_companies is not None:
|
|
updates.append("completed_companies = %s")
|
|
params.append(completed_companies)
|
|
if result_json is not None:
|
|
updates.append("result_json = %s")
|
|
params.append(result_json)
|
|
if error is not None:
|
|
updates.append("error = %s")
|
|
params.append(error)
|
|
|
|
if not updates:
|
|
return self.get_job(job_id)
|
|
|
|
updates.append("updated_at = CURRENT_TIMESTAMP")
|
|
params.append(job_id)
|
|
|
|
with self.get_conn() as conn:
|
|
with conn.cursor(cursor_factory=RealDictCursor) as cursor:
|
|
cursor.execute(
|
|
f"UPDATE jobs SET {', '.join(updates)} WHERE job_id = %s RETURNING *",
|
|
params,
|
|
)
|
|
job = cursor.fetchone()
|
|
conn.commit()
|
|
return dict(job) if job else None
|
|
|
|
def get_job(self, job_id: str) -> Optional[Dict]:
|
|
"""Get a job by ID."""
|
|
with self.get_conn() as conn:
|
|
with conn.cursor(cursor_factory=RealDictCursor) as cursor:
|
|
cursor.execute("SELECT * FROM jobs WHERE job_id = %s", (job_id,))
|
|
job = cursor.fetchone()
|
|
return dict(job) if job else None
|
|
|
|
def list_jobs(
|
|
self,
|
|
status: Optional[str] = None,
|
|
limit: int = 10,
|
|
cursor: Optional[str] = None,
|
|
) -> List[Dict]:
|
|
"""List jobs with optional status filter and cursor-based pagination.
|
|
|
|
Args:
|
|
status: Optional status filter (pending, running, completed, failed).
|
|
limit: Maximum number of jobs to return.
|
|
cursor: Opaque cursor (``created_at|job_id``) from a previous
|
|
response. When provided, only jobs older than the cursor are
|
|
returned.
|
|
|
|
Returns:
|
|
List of job dicts ordered by created_at descending.
|
|
"""
|
|
conditions: list[str] = []
|
|
params: list = []
|
|
|
|
if status:
|
|
conditions.append("status = %s")
|
|
params.append(status)
|
|
|
|
if cursor:
|
|
try:
|
|
ts_str, cursor_job_id = cursor.rsplit("|", 1)
|
|
conditions.append("(created_at, job_id) < (%s, %s)")
|
|
params.extend([ts_str, cursor_job_id])
|
|
except ValueError:
|
|
pass # Ignore malformed cursors; return from start
|
|
|
|
query = "SELECT * FROM jobs"
|
|
if conditions:
|
|
query += " WHERE " + " AND ".join(conditions)
|
|
query += " ORDER BY created_at DESC, job_id DESC LIMIT %s"
|
|
params.append(limit)
|
|
|
|
with self.get_conn() as conn:
|
|
with conn.cursor(cursor_factory=RealDictCursor) as cur:
|
|
cur.execute(query, params)
|
|
return [dict(row) for row in cur.fetchall()]
|
|
|
|
def mark_stale_jobs_failed(self) -> int:
|
|
"""Mark any jobs in 'running' or 'pending' state as 'failed'.
|
|
|
|
Called at startup to clean up jobs that were interrupted by a restart.
|
|
|
|
Returns:
|
|
Number of jobs marked as failed.
|
|
"""
|
|
with self.get_conn() as conn:
|
|
with conn.cursor() as cursor:
|
|
cursor.execute(
|
|
"""
|
|
UPDATE jobs SET status = 'failed', error = 'Interrupted by server restart',
|
|
updated_at = CURRENT_TIMESTAMP
|
|
WHERE status IN ('running', 'pending')
|
|
"""
|
|
)
|
|
count = cursor.rowcount
|
|
conn.commit()
|
|
return count
|
|
|
|
# User Authentication Methods
|
|
|
|
@staticmethod
|
|
def hash_password(password: str) -> str:
|
|
"""Hash a password using bcrypt.
|
|
|
|
Args:
|
|
password: Plain text password
|
|
|
|
Returns:
|
|
Hashed password string
|
|
"""
|
|
return bcrypt.hashpw(password.encode(), bcrypt.gensalt()).decode()
|
|
|
|
@staticmethod
|
|
def verify_password(password: str, password_hash: str) -> bool:
|
|
"""Verify a password against its hash.
|
|
|
|
Args:
|
|
password: Plain text password
|
|
password_hash: Stored hash
|
|
|
|
Returns:
|
|
True if password matches
|
|
"""
|
|
return bcrypt.checkpw(password.encode(), password_hash.encode())
|
|
|
|
def create_user(
|
|
self,
|
|
email: str,
|
|
password: str,
|
|
role: str = "user",
|
|
) -> Optional[Dict]:
|
|
"""Create a new user.
|
|
|
|
Args:
|
|
email: User email
|
|
password: Plain text password
|
|
role: User role ('admin' or 'user')
|
|
|
|
Returns:
|
|
Created user dict or None if email exists
|
|
"""
|
|
password_hash = self.hash_password(password)
|
|
|
|
try:
|
|
with self.get_conn() as conn:
|
|
with conn.cursor(cursor_factory=RealDictCursor) as cursor:
|
|
cursor.execute(
|
|
"""
|
|
INSERT INTO users (email, password_hash, role)
|
|
VALUES (%s, %s, %s)
|
|
RETURNING id, email, role, created_at
|
|
""",
|
|
(email, password_hash, role),
|
|
)
|
|
user = cursor.fetchone()
|
|
conn.commit()
|
|
return dict(user) if user else None
|
|
except psycopg2.errors.UniqueViolation:
|
|
return None
|
|
|
|
def authenticate_user(self, email: str, password: str) -> Optional[Dict]:
|
|
"""Authenticate a user by email and password.
|
|
|
|
Args:
|
|
email: User email
|
|
password: Plain text password
|
|
|
|
Returns:
|
|
User dict if authenticated, None otherwise
|
|
"""
|
|
with self.get_conn() as conn:
|
|
with conn.cursor(cursor_factory=RealDictCursor) as cursor:
|
|
cursor.execute(
|
|
"SELECT * FROM users WHERE email = %s",
|
|
(email,),
|
|
)
|
|
user = cursor.fetchone()
|
|
|
|
if user and self.verify_password(password, user["password_hash"]):
|
|
return {
|
|
"id": user["id"],
|
|
"email": user["email"],
|
|
"role": user["role"],
|
|
"created_at": user["created_at"],
|
|
}
|
|
return None
|
|
|
|
def get_user_by_id(self, user_id: int) -> Optional[Dict]:
|
|
"""Get a user by ID.
|
|
|
|
Args:
|
|
user_id: User ID
|
|
|
|
Returns:
|
|
User dict or None
|
|
"""
|
|
with self.get_conn() as conn:
|
|
with conn.cursor(cursor_factory=RealDictCursor) as cursor:
|
|
cursor.execute(
|
|
"SELECT id, email, role, created_at FROM users WHERE id = %s",
|
|
(user_id,),
|
|
)
|
|
user = cursor.fetchone()
|
|
return dict(user) if user else None
|
|
|
|
def get_user_by_email(self, email: str) -> Optional[Dict]:
|
|
"""Get a user by email.
|
|
|
|
Args:
|
|
email: User email
|
|
|
|
Returns:
|
|
User dict or None
|
|
"""
|
|
with self.get_conn() as conn:
|
|
with conn.cursor(cursor_factory=RealDictCursor) as cursor:
|
|
cursor.execute(
|
|
"SELECT id, email, role, created_at FROM users WHERE email = %s",
|
|
(email,),
|
|
)
|
|
user = cursor.fetchone()
|
|
return dict(user) if user else None
|
|
|
|
def get_all_users(self, limit: int = 100, offset: int = 0) -> List[Dict]:
|
|
"""Get all users (admin only).
|
|
|
|
Args:
|
|
limit: Maximum number of users
|
|
offset: Offset for pagination
|
|
|
|
Returns:
|
|
List of user dicts
|
|
"""
|
|
with self.get_conn() as conn:
|
|
with conn.cursor(cursor_factory=RealDictCursor) as cursor:
|
|
cursor.execute(
|
|
"""
|
|
SELECT id, email, role, created_at
|
|
FROM users
|
|
ORDER BY created_at DESC
|
|
LIMIT %s OFFSET %s
|
|
""",
|
|
(limit, offset),
|
|
)
|
|
return [dict(row) for row in cursor.fetchall()]
|
|
|
|
def update_user_role(self, user_id: int, role: str) -> Optional[Dict]:
|
|
"""Update a user's role (admin only).
|
|
|
|
Args:
|
|
user_id: User ID
|
|
role: New role ('admin' or 'user')
|
|
|
|
Returns:
|
|
Updated user dict or None
|
|
"""
|
|
with self.get_conn() as conn:
|
|
with conn.cursor(cursor_factory=RealDictCursor) as cursor:
|
|
cursor.execute(
|
|
"""
|
|
UPDATE users
|
|
SET role = %s, updated_at = CURRENT_TIMESTAMP
|
|
WHERE id = %s
|
|
RETURNING id, email, role, created_at
|
|
""",
|
|
(role, user_id),
|
|
)
|
|
user = cursor.fetchone()
|
|
conn.commit()
|
|
return dict(user) if user else None
|
|
|
|
def delete_user(self, user_id: int) -> bool:
|
|
"""Delete a user (admin only).
|
|
|
|
Args:
|
|
user_id: User ID
|
|
|
|
Returns:
|
|
True if deleted
|
|
"""
|
|
with self.get_conn() as conn:
|
|
with conn.cursor() as cursor:
|
|
cursor.execute("DELETE FROM users WHERE id = %s", (user_id,))
|
|
deleted = cursor.rowcount > 0
|
|
conn.commit()
|
|
return deleted
|
|
|
|
def get_user_count(self) -> int:
|
|
"""Get total user count.
|
|
|
|
Returns:
|
|
Number of users
|
|
"""
|
|
with self.get_conn() as conn:
|
|
with conn.cursor() as cursor:
|
|
cursor.execute("SELECT COUNT(*) FROM users")
|
|
return cursor.fetchone()[0]
|
|
|
|
# Tracked Companies Methods
|
|
|
|
def add_tracked_company(self, company_name: str) -> Optional[Dict]:
|
|
"""Add a company to the tracking list."""
|
|
with self.get_conn() as conn:
|
|
with conn.cursor(cursor_factory=RealDictCursor) as cursor:
|
|
try:
|
|
cursor.execute(
|
|
"INSERT INTO tracked_companies (company_name) VALUES (%s) RETURNING *",
|
|
(company_name,),
|
|
)
|
|
row = cursor.fetchone()
|
|
conn.commit()
|
|
return dict(row) if row else None
|
|
except Exception:
|
|
conn.rollback()
|
|
return None
|
|
|
|
def remove_tracked_company(self, company_name: str) -> bool:
|
|
"""Remove a company from the tracking list."""
|
|
with self.get_conn() as conn:
|
|
with conn.cursor() as cursor:
|
|
cursor.execute(
|
|
"DELETE FROM tracked_companies WHERE LOWER(company_name) = LOWER(%s)",
|
|
(company_name,),
|
|
)
|
|
conn.commit()
|
|
return cursor.rowcount > 0
|
|
|
|
def list_tracked_companies(self) -> List[Dict]:
|
|
"""List all tracked companies."""
|
|
with self.get_conn() as conn:
|
|
with conn.cursor(cursor_factory=RealDictCursor) as cursor:
|
|
cursor.execute("SELECT * FROM tracked_companies ORDER BY company_name")
|
|
return [dict(row) for row in cursor.fetchall()]
|
|
|
|
def update_tracked_company(
|
|
self, company_name: str, patent_count: int
|
|
) -> None:
|
|
"""Update the last analysis stats for a tracked company."""
|
|
with self.get_conn() as conn:
|
|
with conn.cursor() as cursor:
|
|
cursor.execute(
|
|
"""UPDATE tracked_companies
|
|
SET last_patent_count = %s, last_analysis_at = CURRENT_TIMESTAMP
|
|
WHERE LOWER(company_name) = LOWER(%s)""",
|
|
(patent_count, company_name),
|
|
)
|
|
conn.commit()
|
|
|
|
def store_alert(
|
|
self,
|
|
company_name: str,
|
|
alert_type: str,
|
|
message: str,
|
|
old_value: float | None = None,
|
|
new_value: float | None = None,
|
|
) -> None:
|
|
"""Record an alert for a significant change."""
|
|
with self.get_conn() as conn:
|
|
with conn.cursor() as cursor:
|
|
cursor.execute(
|
|
"""INSERT INTO alerts (company_name, alert_type, message, old_value, new_value)
|
|
VALUES (%s, %s, %s, %s, %s)""",
|
|
(company_name, alert_type, message, old_value, new_value),
|
|
)
|
|
conn.commit()
|
|
|
|
def list_alerts(self, limit: int = 50) -> List[Dict]:
|
|
"""List recent alerts."""
|
|
with self.get_conn() as conn:
|
|
with conn.cursor(cursor_factory=RealDictCursor) as cursor:
|
|
cursor.execute(
|
|
"SELECT * FROM alerts ORDER BY created_at DESC LIMIT %s",
|
|
(limit,),
|
|
)
|
|
return [dict(row) for row in cursor.fetchall()]
|