Spaces:
Running
Running
| from sqlalchemy import create_engine, text | |
| from sqlalchemy.ext.declarative import declarative_base | |
| from sqlalchemy.orm import sessionmaker | |
| import os | |
| from app.config import settings | |
| DATABASE_URL = settings.database_url or os.getenv("DATABASE_URL", "sqlite:///./amazon_intel.db") | |
| # Neon / HF secrets often use postgresql:// — SQLAlchemy 2 may pick psycopg v3. | |
| # Pin psycopg2 so the Space starts with the package already in requirements.txt. | |
| if DATABASE_URL.startswith("postgres://"): | |
| DATABASE_URL = "postgresql+psycopg2://" + DATABASE_URL[len("postgres://"):] | |
| elif DATABASE_URL.startswith("postgresql://") and "+" not in DATABASE_URL.split("://", 1)[0]: | |
| DATABASE_URL = "postgresql+psycopg2://" + DATABASE_URL[len("postgresql://"):] | |
| if DATABASE_URL.startswith("sqlite"): | |
| engine = create_engine(DATABASE_URL, connect_args={"check_same_thread": False}) | |
| else: | |
| engine = create_engine(DATABASE_URL, pool_pre_ping=True, pool_recycle=300) | |
| SessionLocal = sessionmaker(autocommit=False, autoflush=False, bind=engine) | |
| Base = declarative_base() | |
| def migrate_schema(): | |
| """Add columns when schema evolves (SQLite + PostgreSQL).""" | |
| user_alters_sqlite = [ | |
| "ALTER TABLE users ADD COLUMN token_version INTEGER DEFAULT 0", | |
| "ALTER TABLE users ADD COLUMN auth_provider VARCHAR(20) DEFAULT 'email'", | |
| "ALTER TABLE users ADD COLUMN firebase_uid VARCHAR(128)", | |
| ] | |
| user_alters_pg = [ | |
| "ALTER TABLE users ADD COLUMN IF NOT EXISTS token_version INTEGER DEFAULT 0", | |
| "ALTER TABLE users ADD COLUMN IF NOT EXISTS auth_provider VARCHAR(20) DEFAULT 'email'", | |
| "ALTER TABLE users ADD COLUMN IF NOT EXISTS firebase_uid VARCHAR(128)", | |
| ] | |
| sqlite_product_alters = [ | |
| "ALTER TABLE products ADD COLUMN seller_count INTEGER", | |
| "ALTER TABLE products ADD COLUMN buy_box_winner VARCHAR(255)", | |
| "ALTER TABLE products ADD COLUMN buy_box_price NUMERIC(10,2)", | |
| "ALTER TABLE products ADD COLUMN buy_box_is_fba BOOLEAN DEFAULT 1", | |
| "ALTER TABLE products ADD COLUMN is_amazon_sold BOOLEAN DEFAULT 0", | |
| "ALTER TABLE products ADD COLUMN last_data_source VARCHAR(20) DEFAULT 'cached'", | |
| "ALTER TABLE products ADD COLUMN upc VARCHAR(20)", | |
| "ALTER TABLE products ADD COLUMN has_buy_box BOOLEAN DEFAULT 1", | |
| "ALTER TABLE products ADD COLUMN package_weight_lbs NUMERIC(8,3)", | |
| "ALTER TABLE products ADD COLUMN other_sellers_json TEXT", | |
| "ALTER TABLE products ADD COLUMN offers_source VARCHAR(30)", | |
| "ALTER TABLE products ADD COLUMN marketplace VARCHAR(8) DEFAULT 'us'", | |
| "ALTER TABLE products ADD COLUMN currency VARCHAR(8) DEFAULT 'USD'", | |
| "ALTER TABLE products ADD COLUMN bought_in_past_month INTEGER", | |
| "ALTER TABLE products ADD COLUMN listing_intel_json TEXT", | |
| ] | |
| pg_product_alters = [ | |
| "ALTER TABLE products ADD COLUMN IF NOT EXISTS upc VARCHAR(20)", | |
| "ALTER TABLE products ADD COLUMN IF NOT EXISTS has_buy_box BOOLEAN DEFAULT TRUE", | |
| "ALTER TABLE products ADD COLUMN IF NOT EXISTS package_weight_lbs NUMERIC(8,3)", | |
| "ALTER TABLE products ADD COLUMN IF NOT EXISTS other_sellers_json TEXT", | |
| "ALTER TABLE products ADD COLUMN IF NOT EXISTS offers_source VARCHAR(30)", | |
| "ALTER TABLE products ADD COLUMN IF NOT EXISTS marketplace VARCHAR(8) DEFAULT 'us'", | |
| "ALTER TABLE products ADD COLUMN IF NOT EXISTS currency VARCHAR(8) DEFAULT 'USD'", | |
| "ALTER TABLE products ADD COLUMN IF NOT EXISTS bought_in_past_month INTEGER", | |
| "ALTER TABLE products ADD COLUMN IF NOT EXISTS listing_intel_json TEXT", | |
| ] | |
| is_pg = DATABASE_URL.startswith("postgresql") | |
| with engine.connect() as conn: | |
| for sql in (user_alters_pg if is_pg else user_alters_sqlite): | |
| try: | |
| conn.execute(text(sql)) | |
| conn.commit() | |
| except Exception: | |
| pass | |
| # Google-only accounts may have no password | |
| if is_pg: | |
| try: | |
| conn.execute(text("ALTER TABLE users ALTER COLUMN password_hash DROP NOT NULL")) | |
| conn.commit() | |
| except Exception: | |
| pass | |
| if DATABASE_URL.startswith("sqlite"): | |
| for sql in sqlite_product_alters: | |
| try: | |
| conn.execute(text(sql)) | |
| conn.commit() | |
| except Exception: | |
| pass | |
| if is_pg: | |
| for sql in pg_product_alters: | |
| try: | |
| conn.execute(text(sql)) | |
| conn.commit() | |
| except Exception: | |
| pass | |
| # Prefer (asin, marketplace) uniqueness so US + AE can coexist | |
| for sql in ( | |
| "UPDATE products SET marketplace = 'us' WHERE marketplace IS NULL", | |
| "UPDATE products SET currency = 'USD' WHERE currency IS NULL AND (marketplace = 'us' OR marketplace IS NULL)", | |
| "ALTER TABLE products DROP CONSTRAINT IF EXISTS products_asin_key", | |
| "ALTER TABLE products DROP CONSTRAINT IF EXISTS ix_products_asin", | |
| "CREATE UNIQUE INDEX IF NOT EXISTS uq_products_asin_marketplace ON products (asin, marketplace)", | |
| ): | |
| try: | |
| conn.execute(text(sql)) | |
| conn.commit() | |
| except Exception: | |
| pass | |
| else: | |
| for sql in ( | |
| "UPDATE products SET marketplace = 'us' WHERE marketplace IS NULL", | |
| "UPDATE products SET currency = 'USD' WHERE currency IS NULL", | |
| ): | |
| try: | |
| conn.execute(text(sql)) | |
| conn.commit() | |
| except Exception: | |
| pass | |
| def get_db(): | |
| db = SessionLocal() | |
| try: | |
| yield db | |
| finally: | |
| db.close() | |