ProspectIQ / db.py
Ashwin
Merge origin/main branch with our custom role-based features
2f574bb
Raw
History Blame Contribute Delete
29.6 kB
import sqlite3
import os
from pathlib import Path
import psycopg2
from psycopg2 import pool
DB_PATH = Path(__file__).resolve().parent / "prospectiq.db"
SCHEMA_PATH = Path(__file__).resolve().parent / "schema.sql"
# Global PG Connection Pool
pg_pool = None
def init_pg_pool():
global pg_pool
if pg_pool is not None:
return
host = os.environ.get("SUPABASE_DB_HOST") or ""
db = os.environ.get("SUPABASE_DB_NAME", "postgres")
user = os.environ.get("SUPABASE_DB_USER", "postgres")
password = os.environ.get("SUPABASE_DB_PASSWORD") or ""
port = os.environ.get("SUPABASE_DB_PORT", "6543")
if host.strip() and password.strip():
try:
pg_pool = psycopg2.pool.ThreadedConnectionPool(
minconn=2,
maxconn=20,
host=host,
database=db,
user=user,
password=password,
port=port
)
print("Successfully initialized PostgreSQL connection pool.")
except Exception as e:
print(f"Error initializing PostgreSQL connection pool: {e}")
pg_pool = None
class PostgresRow:
def __init__(self, data, description):
self._data = data
self._keys = [desc[0] for desc in description]
self._key_map = {name: i for i, name in enumerate(self._keys)}
def __getitem__(self, key):
if isinstance(key, int):
return self._data[key]
elif isinstance(key, str):
idx = self._key_map.get(key)
if idx is None:
raise KeyError(key)
return self._data[idx]
else:
raise TypeError("Index must be int or str")
def keys(self):
return self._keys
def __len__(self):
return len(self._data)
def __iter__(self):
return iter(self._keys)
def get(self, key, default=None):
if key in self._key_map:
return self[key]
return default
class PostgresCursorWrapper:
def __init__(self, pg_cursor, connection):
self.cursor = pg_cursor
self.connection = connection
self._lastrowid = None
def __enter__(self):
return self
def __exit__(self, exc_type, exc_val, exc_tb):
self.close()
def execute(self, sql, params=None):
sql = sql.replace('?', '%s')
sql_upper = sql.upper().strip()
if "INSERT OR REPLACE INTO INTERACTIONS" in sql_upper:
sql = """
INSERT INTO interactions (
id, kind, title, status, input_json, output_json, result_text, logs_json,
error, run_dir, profile_path, command_json, returncode, pid,
created_at, started_at, finished_at, updated_at
) VALUES (%s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s)
ON CONFLICT (id) DO UPDATE SET
kind = EXCLUDED.kind,
title = EXCLUDED.title,
status = EXCLUDED.status,
input_json = EXCLUDED.input_json,
output_json = EXCLUDED.output_json,
result_text = EXCLUDED.result_text,
logs_json = EXCLUDED.logs_json,
error = EXCLUDED.error,
run_dir = EXCLUDED.run_dir,
profile_path = EXCLUDED.profile_path,
command_json = EXCLUDED.command_json,
returncode = EXCLUDED.returncode,
pid = EXCLUDED.pid,
created_at = EXCLUDED.created_at,
started_at = EXCLUDED.started_at,
finished_at = EXCLUDED.finished_at,
updated_at = EXCLUDED.updated_at
"""
elif "INSERT OR IGNORE INTO" in sql_upper:
sql = sql.replace("INSERT OR IGNORE INTO", "INSERT INTO") + " ON CONFLICT DO NOTHING"
is_insert = sql_upper.startswith("INSERT")
if is_insert and "RETURNING" not in sql_upper and "ROLE_PERMISSIONS" not in sql_upper:
sql = sql.rstrip().rstrip(';') + " RETURNING id"
# Invalidate stage or user caches automatically on database writes
is_write = sql_upper.startswith("INSERT") or sql_upper.startswith("UPDATE") or sql_upper.startswith("DELETE")
if is_write:
try:
if "STAGES" in sql_upper:
cache_delete("prospectiq:stages")
if "USERS" in sql_upper:
cache_delete("prospectiq:users")
except Exception:
pass
if params is not None:
self.cursor.execute(sql, params)
else:
self.cursor.execute(sql)
self._lastrowid = None
if is_insert and "RETURNING" in sql.upper():
try:
row = self.cursor.fetchone()
if row:
self._lastrowid = row[0]
except Exception:
pass
return self
def executemany(self, sql, seq_of_params):
sql = sql.replace('?', '%s')
self.cursor.executemany(sql, seq_of_params)
return self
def fetchone(self):
row = self.cursor.fetchone()
if row is None:
return None
return PostgresRow(row, self.cursor.description)
def fetchall(self):
rows = self.cursor.fetchall()
if not rows:
return []
desc = self.cursor.description
return [PostgresRow(r, desc) for r in rows]
@property
def rowcount(self):
return self.cursor.rowcount
@property
def lastrowid(self):
return self._lastrowid
def close(self):
self.cursor.close()
class PostgresConnectionWrapper:
def __init__(self, pg_conn):
self.conn = pg_conn
def cursor(self):
return PostgresCursorWrapper(self.conn.cursor(), self)
def execute(self, sql, params=None):
cur = self.cursor()
cur.execute(sql, params)
return cur
def executemany(self, sql, seq_of_params):
cur = self.cursor()
cur.executemany(sql, seq_of_params)
return cur
def commit(self):
self.conn.commit()
def rollback(self):
self.conn.rollback()
def close(self):
global pg_pool
if pg_pool is not None:
try:
pg_pool.putconn(self.conn)
except Exception as e:
print(f"Error returning connection to pool: {e}")
try:
self.conn.close()
except Exception:
pass
else:
self.conn.close()
def __enter__(self):
return self
def __exit__(self, exc_type, exc_val, exc_tb):
try:
if exc_type is not None:
self.rollback()
else:
self.commit()
finally:
self.close()
def is_supabase_active() -> bool:
global pg_pool
if pg_pool is None:
init_pg_pool()
return pg_pool is not None
def sync_postgres_sequence(table_name: str, column_name: str = "id"):
if not is_supabase_active():
return
safe_table = "".join(ch for ch in str(table_name) if ch.isalnum() or ch == "_")
safe_column = "".join(ch for ch in str(column_name) if ch.isalnum() or ch == "_")
if not safe_table or not safe_column:
return
conn = get_db_connection()
cursor = conn.cursor()
try:
cursor.execute(f"""
SELECT setval(
pg_get_serial_sequence('{safe_table}', '{safe_column}'),
COALESCE((SELECT MAX({safe_column}) FROM {safe_table}), 0) + 1,
false
)
""")
conn.commit()
except Exception as e:
conn.rollback()
print(f"Warning: could not sync PostgreSQL sequence for {safe_table}.{safe_column}: {e}")
finally:
conn.close()
def sync_postgres_sequences():
if not is_supabase_active():
return
for table_name in (
"tenants",
"users",
"stages",
"companies",
"contacts",
"activities",
"notes",
"stage_changes",
"ai_field_values",
):
sync_postgres_sequence(table_name)
def get_db_connection():
global pg_pool
if pg_pool is None:
init_pg_pool()
if pg_pool is not None:
try:
conn = pg_pool.getconn()
conn.autocommit = False
return PostgresConnectionWrapper(conn)
except Exception as e:
print(f"PostgreSQL connection failed, falling back to SQLite: {e}")
conn = sqlite3.connect(DB_PATH)
conn.row_factory = sqlite3.Row
return conn
def dict_from_row(row):
return dict(row) if row is not None else None
def init_db():
# Execute schema
conn = get_db_connection()
cursor = conn.cursor()
with open(SCHEMA_PATH, 'r') as f:
schema_sql = f.read()
cursor.executescript(schema_sql)
conn.commit()
# Seed stages if empty
cursor.execute("SELECT COUNT(*) FROM stages")
if cursor.fetchone()[0] == 0:
default_stages = [
("New", 1, "#f1f5f9", "#475569"),
("Shortlisted", 2, "#eff6ff", "#1d4ed8"),
("Research Done", 3, "#f5f3ff", "#6d28d9"),
("Contacted", 4, "#ecfdf5", "#047857"),
("Engaged", 5, "#f0fdf4", "#15803d"),
("Meeting Scheduled", 6, "#fffbeb", "#b45309"),
("Meeting Done", 7, "#fef3c7", "#b45309"),
("Proposal Sent", 8, "#faf5ff", "#7e22ce"),
("Won", 9, "#d1fae5", "#065f46"),
("Paused / Not a Fit", 10, "#fee2e2", "#991b1b")
]
cursor.executemany(
"INSERT INTO stages (name, order_index, color_bg, color_text) VALUES (?, ?, ?, ?)",
default_stages
)
conn.commit()
# Seed tenant if empty
cursor.execute("SELECT COUNT(*) FROM tenants")
if cursor.fetchone()[0] == 0:
cursor.execute(
"INSERT INTO tenants (id, name, brand_color) VALUES (1, 'MapleMPSS', '#3b82f6')"
)
conn.commit()
# Seed users if empty
cursor.execute("SELECT COUNT(*) FROM users")
if cursor.fetchone()[0] == 0:
default_users = [
(1, 1, "admin1@maplempss.com", "admin1", "Admin"),
(2, 1, "bdrep1@maplempss.com", "bd rep1", "BD Rep"),
(3, 1, "viewer1@maplempss.com", "viewer1", "Viewer"),
(4, 1, "bdrep2@maplempss.com", "bd rep2", "BD Rep"),
(5, 1, "bdrep3@maplempss.com", "bd rep3", "BD Rep")
]
cursor.executemany(
"INSERT INTO users (id, tenant_id, email, name, role) VALUES (?, ?, ?, ?, ?)",
default_users
)
conn.commit()
# Seed 10 dummy companies if empty
cursor.execute("SELECT COUNT(*) FROM companies")
if cursor.fetchone()[0] == 0:
# Columns:
# id, tenant_id, name, legal_name, website, phone, nic_code, industry,
# street, city, province, postal_code, country, size_staff, revenue_band, business_type,
# products_made, parts_sourced, current_supplier_notes, maplempss_fit_score, stage_id, assigned_rep_id,
# created_by, updated_by
dummy_companies = [
# 1. Tesla Inc. - Large OEM (California), Contacted stage, has 2 contacts, active call history, upcoming follow-up.
(1, 1, "Tesla Inc.", "Tesla Inc.", "https://www.tesla.com", "1-800-TESLA", "3711", "Automotive OEM",
"3500 Deer Creek Road", "Palo Alto", "California", "94304", "USA", 100000, "1B+", "OEM",
"Electric Vehicles, Energy Storage Systems", "Stamping parts, cast aluminum knuckles, fasteners", "Currently sources cast parts globally, looking for localized North American CNC shops.", 4.5, 4, 2, 1, 1),
# 2. Northstar Sourcing - Mid distributor (Ontario), Engaged stage, has overdue follow-up.
(2, 1, "Northstar Sourcing", "Northstar Sourcing Ltd.", "https://www.northstarsourcing.ca", "416-555-0199", "3327", "Aerospace Parts",
"100 Airport Road", "Toronto", "Ontario", "M5V 2N8", "Canada", 120, "10M-50M", "distributor",
"Landing gear brackets, actuator assemblies", "Bushings, sleeves, pivot pins", "Supplier in Mexico failing quality checks. Seeking backup CNC capacity in Ontario.", 4.0, 5, 2, 1, 1),
# 3. Ontario Automotive OEM - Large OEM (Ontario), New stage, single contact, no activities but has notes.
(3, 1, "Ontario Automotive OEM", "Ontario Automotive OEM Corp.", "https://www.ontarioautooem.ca", "519-555-0123", "3714", "Auto Parts Mfg",
"500 Tecumseh Rd E", "Windsor", "Ontario", "N8X 2S2", "Canada", 450, "100M-500M", "OEM",
"Suspension arms, links, engine mounts", "Steel forgings, CNC machined journals", "Currently single sourced. High risk due to logistics bottlenecks.", 3.5, 1, 2, 1, 1),
# 4. Apex Machining Ltd. - Small machine shop (BC), Research Done stage, 0 contacts (edge case: no contacts logged).
(4, 1, "Apex Machining Ltd.", None, "https://www.apexmachining.ca", None, "3327", "CNC Shop",
"12330 Vulcan Way", "Richmond", "British Columbia", "V6V 1J8", "Canada", 45, "1M-10M", "other",
"Custom hydraulic fittings, prototype valves", "Aluminum bars, brass castings", "High precision shop. Might be a sub-contracting partner, not a direct client.", 4.2, 3, 2, 1, 1),
# 5. Calgary Metal Press - Small shop, Paused/Not a Fit stage, low fit score.
(5, 1, "Calgary Metal Press", "Calgary Metal Pressing Co.", None, "403-555-0450", "3321", "Metal Stamping",
"7800 51 St SE", "Calgary", "Alberta", "T2C 4E3", "Canada", 15, "<1M", "OEM",
"Stamped brackets, flat washers", "Sheet metal coils", "Too small for our outsourcing program. Margins do not justify logistics. Paused.", 1.8, 10, 2, 1, 1),
# 6. Quebec Aerospace Group - Large Aerospace OEM, Won stage, multiple contacts and activities.
(6, 1, "Quebec Aerospace Group", "Quebec Aerospace Group Inc.", "https://www.quebecaerospace.com", "514-555-0900", "3364", "Aerospace OEM",
"450 Boulevard Saint-Martin O", "Laval", "Quebec", "H7M 3Y6", "Canada", 850, "500M-1B", "OEM",
"Turbine blades, housing rings, landing gears", "Inconel castings, titanium shafts", "Closed deal on turbine flange prototypes. Moving to production batches.", 4.8, 9, 2, 1, 1),
# 7. Precision Components India - Large global supplier, Shortlisted stage (Edge case: foreign country, fit assessment).
(7, 1, "Precision Components India", "Precision Components India Pvt. Ltd.", "https://www.precisioncompindia.com", "+91-120-555-0100", "3714", "Auto Parts Mfg",
"Sector 63, Block C", "Noida", "Uttar Pradesh", "201301", "India", 2100, "100M-500M", "OEM / distributor",
"Steering gears, universal joints, steering columns", "Steel bars, pinion blanks", "High volume Indian manufacturer. Assessing if logistics cost overrides sourcing discount.", 3.9, 2, 2, 1, 1),
# 8. Hamilton Gearworks - Mid OEM, Meeting Scheduled stage, upcoming meeting follow-up.
(8, 1, "Hamilton Gearworks", "Hamilton Industrial Gearworks Ltd.", "https://www.hamiltongear.ca", "905-555-0111", "3327", "Industrial Gears OEM",
"250 Gage Ave N", "Hamilton", "Ontario", "L8L 7B1", "Canada", 85, "10M-50M", "OEM",
"Heavy industrial gearboxes, custom splines", "Cast iron gear cases, forging blanks", "Scheduled discovery call with VP Operations to discuss gearbox housing outsourcing.", 4.1, 6, 2, 1, 1),
# 9. WestCoast Castings - Small Foundry, Meeting Done stage, pending next steps.
(9, 1, "WestCoast Castings", "WestCoast Foundries Ltd.", "https://www.westcoastcastings.com", "604-555-0188", "3315", "Foundry",
"1550 Kingsway Ave", "Port Coquitlam", "British Columbia", "V3C 6N6", "Canada", 30, "1M-10M", "OEM",
"Sand cast aluminum & bronze housings", "Ingots, sand molds", "Meeting completed last week. Waiting for their capacity sheet regarding 10-ton casting capacity.", 3.0, 7, 2, 1, 1),
# 10. Montreal Custom Fasteners - Micro distributor, Proposal Sent stage, follow-up due today.
(10, 1, "Montreal Custom Fasteners", None, None, "514-555-0133", "3327", "Fastener Distributor",
"890 Rue Saint-Jacques", "Montreal", "Quebec", "H3C 1H8", "Canada", 8, "<1M", "distributor",
"Custom sockets, micro fasteners", "Wire coils", "Proposal for standard fasteners supply sent. Decision-maker review today.", 3.6, 8, 2, 1, 1)
]
cursor.executemany(
"""INSERT INTO companies (
id, tenant_id, name, legal_name, website, phone, nic_code, industry,
street, city, province, postal_code, country, size_staff, revenue_band, business_type,
products_made, parts_sourced, current_supplier_notes, maplempss_fit_score, stage_id, assigned_rep_id,
created_by, updated_by
) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)""",
dummy_companies
)
conn.commit()
# Seed contacts
# Columns: company_id, first_name, last_name, title, email, phone, linkedin, is_primary
dummy_contacts = [
(1, "Elon", "Musk", "CEO", "elon@tesla.com", "1-800-TESLA", "https://linkedin.com/in/elonmusk", 0),
(1, "Priya", "Sharma", "VP Procurement", "priya.sharma@tesla.com", "650-555-0102", "https://linkedin.com/in/priyasharmaprocurement", 1),
(2, "David", "Miller", "Sourcing Manager", "david@northstarsourcing.ca", "416-555-0144", "https://linkedin.com/in/davidmillersourcing", 1),
(3, "Robert", "Chen", "Director of Supply Chain", "r.chen@ontarioautooem.ca", "519-555-0199", "https://linkedin.com/in/robertchensupply", 1),
# Company 4 has 0 contacts
(5, "Frank", "Glover", "Operations Manager", "frank@calgarymetalpress.com", "403-555-0451", None, 1),
(6, "Jean", "Dupont", "VP Procurement", "j.dupont@quebecaerospace.com", "514-555-0901", "https://linkedin.com/in/jeandupont", 1),
(6, "Marc", "Leclerc", "Buyer", "m.leclerc@quebecaerospace.com", "514-555-0902", None, 0),
(7, "Rajesh", "Kumar", "Director Sourcing", "r.kumar@precisioncompindia.com", None, None, 1),
(8, "Sarah", "Connor", "VP Operations", "sarah.connor@hamiltongear.ca", "905-555-0112", "https://linkedin.com/in/sarahconnor", 1),
(9, "Kevin", "Bacon", "Plant Manager", "kevin@westcoastcastings.com", "604-555-0189", None, 1),
(10, "Pierre", "Trudeau", "Owner", "pierre@montrealfasteners.com", "514-555-0134", None, 1)
]
cursor.executemany(
"""INSERT INTO contacts (
company_id, first_name, last_name, title, email, phone, linkedin, is_primary
) VALUES (?, ?, ?, ?, ?, ?, ?, ?)""",
dummy_contacts
)
conn.commit()
# Seed activities and notes
# Columns: company_id, contact_id, user_id, type, outcome, notes, next_action, next_action_date
dummy_activities = [
# Tesla - upcoming follow-up
(1, 2, 2, "Cold Call", "Answered", "Had a brief conversation with Priya. She requested our machinery inventory list and casting capacity details.", "Send CNC capacity list and certifications", "2026-06-25"),
(1, 2, 2, "Email", "Sent", "Follow-up email containing CNC capacity sheet and ISO registration certificate pdf.", None, None),
# Northstar - overdue follow-up (next action date: 2026-06-15)
(2, 3, 2, "Meeting", "Answered", "Discovery call completed. Discussed landing gear bushing volumes. They source 50,000 units annually. Mexico vendor experiencing major delays.", "Prepare landing gear bushing price quote", "2026-06-15"),
# Quebec Aerospace - Won
(6, 6, 2, "Meeting", "discovery", "Discovery call completed. Inconel flange specifications discussed.", "Prepare proposal", "2026-05-10"),
(6, 6, 2, "Email", "Sent", "Sent formal proposal quote for 100 prototype flanges.", None, None),
(6, 6, 2, "Meeting", "negotiation", "Flange proposal accepted. First prototype PO issued.", "Deliver flanges by end of July", "2026-07-31"),
# Hamilton Gearworks - upcoming follow-up (next action date: 2026-06-22)
(8, 9, 2, "LinkedIn Message", "Accepted", "Sarah accepted connection request. Offered to show gearbox housing capabilities.", "Scheduled discovery call on Teams for June 22nd", "2026-06-22"),
# Montreal Fasteners - due today (next action date: 2026-06-20)
(10, 11, 2, "Email", "Sent", "Sent quotation for custom brass socket order. Pricing valid for 30 days.", "Follow up on proposal decision", "2026-06-20")
]
cursor.executemany(
"""INSERT INTO activities (
company_id, contact_id, user_id, type, outcome, notes, next_action, next_action_date
) VALUES (?, ?, ?, ?, ?, ?, ?, ?)""",
dummy_activities
)
conn.commit()
# Seed notes
# Columns: company_id, contact_id, user_id, body, is_pinned
dummy_notes = [
(1, 2, 2, "Priya mentioned that delivery delays with current CNC supplier in Mexico are holding up their lines. Extremely high urgency.", 1),
(2, 3, 2, "Northstar Sourcing prefers DDP shipping terms to their warehouse in Toronto. Must build freight pricing into quote.", 0),
(3, 4, 2, "Windsor plant operates on a 3-shift rotation. Will require a supplier with strong backup capacity planning.", 0),
(4, None, 2, "Assessed website products: Apex manufactures high-end hydraulic parts. Might be better suited as a strategic partner than a direct buyer.", 0)
]
cursor.executemany(
"""INSERT INTO notes (
company_id, contact_id, user_id, body, is_pinned
) VALUES (?, ?, ?, ?, ?)""",
dummy_notes
)
conn.commit()
# Seed initial stage_changes
dummy_stage_changes = [
(1, None, 4, 2, "manual"),
(2, None, 5, 2, "manual"),
(3, None, 1, 2, "manual"),
(4, None, 3, 2, "manual"),
(5, None, 10, 2, "manual"),
(6, None, 9, 2, "manual"),
(7, None, 2, 2, "manual"),
(8, None, 6, 2, "manual"),
(9, None, 7, 2, "manual"),
(10, None, 8, 2, "manual")
]
cursor.executemany(
"""INSERT INTO stage_changes (
company_id, from_stage_id, to_stage_id, changed_by_user_id, source
) VALUES (?, ?, ?, ?, ?)""",
dummy_stage_changes
)
conn.commit()
conn.close()
from flask import request
def get_request_user():
from flask import session
try:
if 'user_id' in session:
return session['user_id'], session['user_role'], session['user_name']
except Exception:
pass
try:
user_id = request.headers.get("X-User-Id")
role = request.headers.get("X-User-Role")
name = request.headers.get("X-User-Name")
try:
user_id = int(user_id) if user_id else 2
except (ValueError, TypeError):
user_id = 2
role = role if role else "BD Rep"
name = name if name else "bd rep1"
return user_id, role, name
except Exception:
return 2, "BD Rep", "bd rep1"
def get_request_tenant_id():
from flask import session
try:
if 'tenant_id' in session:
return session['tenant_id']
except Exception:
pass
try:
user_id, _, _ = get_request_user()
conn = get_db_connection()
cursor = conn.cursor()
cursor.execute("SELECT tenant_id FROM users WHERE id = ?", (user_id,))
row = cursor.fetchone()
conn.close()
if row:
return row[0]
except Exception:
pass
return 1
def get_user_access_level(db_conn, user_id, user_role, company_id):
if user_role in ('Super Admin', 'Admin'):
return 'full'
cursor = db_conn.cursor()
cursor.execute("SELECT assigned_rep_id FROM companies WHERE id = ?", (company_id,))
row = cursor.fetchone()
if not row:
return 'none'
assigned_rep = row[0]
if user_role == 'BD Rep':
if assigned_rep == user_id:
return 'full'
else:
return 'basic'
if user_role == 'Viewer':
return 'viewer'
return 'basic'
if __name__ == "__main__":
init_db()
print("Database initialized and 10 dummy records seeded successfully.")
# =====================================================================
# Redis Caching Utility Functions (Stage 4)
# — With in-memory fallback when Redis is unavailable
# =====================================================================
import redis
import json
import time
import threading
from collections import OrderedDict
redis_client = None
_redis_last_attempt = 0
_REDIS_RETRY_INTERVAL = 60 # seconds between Redis reconnection attempts
_redis_fallback_warned = False
# In-memory fallback cache (thread-safe, TTL-aware, LRU-eviction)
_mem_cache = OrderedDict() # key -> (value, expiry_timestamp)
_mem_cache_lock = threading.Lock()
_MEM_CACHE_MAX_SIZE = 1000
def get_redis_client():
"""Try to return a working Redis client, or None if unavailable."""
global redis_client, _redis_last_attempt, _redis_fallback_warned
if redis_client is not None:
return redis_client
now = time.time()
if now - _redis_last_attempt < _REDIS_RETRY_INTERVAL:
return None
_redis_last_attempt = now
try:
redis_url = (
os.environ.get("REDIS_URL")
or os.environ.get("CELERY_BROKER_URL")
or os.environ.get("CELERY_RESULT_BACKEND")
or "redis://127.0.0.1:6379/0"
)
client = redis.Redis.from_url(redis_url, decode_responses=True, socket_connect_timeout=2)
client.ping()
redis_client = client
if _redis_fallback_warned:
print("Redis cache connection restored — switching from in-memory fallback to Redis.")
_redis_fallback_warned = False
return redis_client
except Exception as e:
if not _redis_fallback_warned:
print(f"Warning: Redis cache unavailable ({e}). Using in-memory fallback cache.")
_redis_fallback_warned = True
redis_client = None
return None
# --- In-memory cache helpers ---
def _mem_get(key):
with _mem_cache_lock:
entry = _mem_cache.get(key)
if entry is None:
return None
value, expiry = entry
if time.time() > expiry:
del _mem_cache[key]
return None
# Move to end for LRU freshness
_mem_cache.move_to_end(key)
return value
def _mem_set(key, value, timeout):
with _mem_cache_lock:
expiry = time.time() + timeout
if key in _mem_cache:
_mem_cache.move_to_end(key)
_mem_cache[key] = (value, expiry)
# Evict oldest entries if over capacity
while len(_mem_cache) > _MEM_CACHE_MAX_SIZE:
_mem_cache.popitem(last=False)
def _mem_delete(key):
with _mem_cache_lock:
_mem_cache.pop(key, None)
# --- Public cache API (same signatures as before) ---
def cache_get(key):
client = get_redis_client()
if client:
try:
val = client.get(key)
if val:
return json.loads(val)
except Exception as e:
print(f"Redis cache_get error: {e}")
# Fallback to in-memory
return _mem_get(key)
def cache_set(key, value, timeout=3600):
client = get_redis_client()
if client:
try:
client.setex(key, timeout, json.dumps(value))
return True
except Exception as e:
print(f"Redis cache_set error: {e}")
# Always store in memory as well (acts as fallback + L1 cache)
_mem_set(key, value, timeout)
return True
def cache_delete(key):
client = get_redis_client()
if client:
try:
client.delete(key)
except Exception as e:
print(f"Redis cache_delete error: {e}")
_mem_delete(key)
return True