import os from sqlalchemy import create_engine, text from sqlalchemy.orm import sessionmaker, declarative_base from sqlalchemy.pool import NullPool # Try to use CockroachDB dialect if available try: import cockroachdb.sqlalchemy.dialect COCKROACHDB_AVAILABLE = True except ImportError: COCKROACHDB_AVAILABLE = False # Get database URL from environment variable # Default to SQLite for local development if not set ORIGINAL_DATABASE_URL = os.getenv( "DATABASE_URL", "sqlite:///./postgen.db" ) DATABASE_URL = ORIGINAL_DATABASE_URL # For CockroachDB, we need to handle SSL and connection pooling if ORIGINAL_DATABASE_URL.startswith("postgresql://") or ORIGINAL_DATABASE_URL.startswith("postgres://") or ORIGINAL_DATABASE_URL.startswith("cockroachdb://"): # Check if this is a CockroachDB connection (use original URL before modifications) is_cockroach = "cockroachlabs" in ORIGINAL_DATABASE_URL.lower() # CockroachDB connection - use NullPool to avoid connection issues # CockroachDB requires SSL, so we ensure sslmode is set # Use 'require' mode which uses SSL but doesn't require certificate file # For production with certificate, use 'verify-full' and provide sslrootcert cert_path = os.path.expanduser("~/.postgresql/root.crt") if "sslmode" not in DATABASE_URL: separator = "&" if "?" in DATABASE_URL else "?" # Use 'require' instead of 'verify-full' to work without certificate file # Still secure (uses SSL) but doesn't verify the certificate DATABASE_URL = f"{DATABASE_URL}{separator}sslmode=require" elif "sslmode=verify-full" in DATABASE_URL and not os.path.exists(cert_path): # If verify-full is set but cert file doesn't exist, change to require DATABASE_URL = DATABASE_URL.replace("sslmode=verify-full", "sslmode=require") print("⚠ Certificate file not found, using sslmode=require instead of verify-full") # Use CockroachDB dialect if available and this is a CockroachDB connection if is_cockroach and COCKROACHDB_AVAILABLE: # Replace postgresql:// with cockroachdb:// to use CockroachDB dialect DATABASE_URL = DATABASE_URL.replace("postgresql://", "cockroachdb://", 1) DATABASE_URL = DATABASE_URL.replace("postgres://", "cockroachdb://", 1) # Configure engine engine = create_engine( DATABASE_URL, poolclass=NullPool, # CockroachDB works better with NullPool echo=False, # Set to True for SQL query debugging connect_args={} # No special connect args needed ) else: # SQLite for local development engine = create_engine( DATABASE_URL, connect_args={"check_same_thread": False} if "sqlite" in DATABASE_URL else {} ) SessionLocal = sessionmaker(autocommit=False, autoflush=False, bind=engine) Base = declarative_base() def get_db(): """Dependency to get database session""" db = SessionLocal() try: yield db finally: db.close() def get_direct_psycopg2_connection(): """Get a direct psycopg2 connection bypassing SQLAlchemy version parsing""" try: import psycopg2 from urllib.parse import urlparse, parse_qs # Only for PostgreSQL/CockroachDB connections if not (ORIGINAL_DATABASE_URL.startswith("postgresql://") or ORIGINAL_DATABASE_URL.startswith("postgres://") or ORIGINAL_DATABASE_URL.startswith("cockroachdb://")): return None # Parse connection string url_for_parsing = ORIGINAL_DATABASE_URL.replace("cockroachdb://", "postgresql://") parsed = urlparse(url_for_parsing) dbname = parsed.path[1:] if parsed.path else "defaultdb" user = parsed.username password = parsed.password host = parsed.hostname port = parsed.port or 26257 # Get sslmode from query params params = parse_qs(parsed.query) sslmode_list = params.get('sslmode', ['require']) sslmode = sslmode_list[0] if sslmode_list else 'require' # Check if certificate file exists for verify-full mode cert_path = os.path.expanduser("~/.postgresql/root.crt") if sslmode == 'verify-full' and not os.path.exists(cert_path): # Change to require mode if certificate doesn't exist sslmode = 'require' # Create connection conn = psycopg2.connect( dbname=dbname, user=user, password=password, host=host, port=port, sslmode=sslmode ) return conn except Exception as e: print(f"Failed to create direct psycopg2 connection: {e}") return None def ensure_default_user(): """Ensure a default user (id=1) exists in the database""" try: conn = get_direct_psycopg2_connection() if not conn: return 1 # Return default ID if connection fails try: cursor = conn.cursor() # Check if user with id=1 exists cursor.execute("SELECT id FROM users WHERE id = 1") if cursor.fetchone(): cursor.close() conn.close() return 1 # Create default user if it doesn't exist # Try with ON CONFLICT first (PostgreSQL/CockroachDB) try: cursor.execute(""" INSERT INTO users (id, email, name, created_at) VALUES (1, 'default@postgen.app', 'Default User', NOW()) ON CONFLICT (id) DO NOTHING """) except Exception: # If ON CONFLICT fails, try without it (might be unique constraint on email) try: cursor.execute(""" INSERT INTO users (id, email, name, created_at) VALUES (1, 'default@postgen.app', 'Default User', NOW()) """) except Exception as insert_error: # User might already exist (race condition), check again cursor.execute("SELECT id FROM users WHERE id = 1 OR email = 'default@postgen.app' LIMIT 1") row = cursor.fetchone() if row: cursor.close() conn.close() return row[0] # If still fails, re-raise raise insert_error conn.commit() cursor.close() conn.close() return 1 except Exception as e: # If everything fails, try to get any existing user or return default try: cursor.execute("SELECT id FROM users LIMIT 1") row = cursor.fetchone() cursor.close() conn.close() if row: return row[0] except: if conn: conn.close() print(f"Warning: Could not ensure default user: {e}") return 1 # Return default ID as fallback except Exception as e: print(f"Error ensuring default user: {e}") return 1 # Return default ID as fallback def init_db(): """Initialize database tables""" try: from app.models import User, Integration, Asset, Post, Campaign # Try to create tables # For CockroachDB, version parsing may fail but connection still works try: Base.metadata.create_all(bind=engine) print("✓ Database tables created successfully") return True except Exception as create_error: error_str = str(create_error) # Check if it's a version parsing error (non-fatal for CockroachDB) if "Could not determine version" in error_str: # Version parsing failed, but CockroachDB connection works # Use psycopg2 directly to bypass SQLAlchemy's version parsing try: import psycopg2 from urllib.parse import urlparse, parse_qs # Parse connection string to get connection parameters # Handle both cockroachdb:// and postgresql:// schemes # Use original URL before any modifications url_for_parsing = ORIGINAL_DATABASE_URL.replace("cockroachdb://", "postgresql://") parsed = urlparse(url_for_parsing) dbname = parsed.path[1:] if parsed.path else "defaultdb" user = parsed.username password = parsed.password host = parsed.hostname port = parsed.port or 26257 # Get sslmode from query params (use require as default for CockroachDB) params = parse_qs(parsed.query) sslmode_list = params.get('sslmode', ['require']) sslmode = sslmode_list[0] if sslmode_list else 'require' # Connect directly with psycopg2 (bypasses SQLAlchemy version parsing) conn = psycopg2.connect( dbname=dbname, user=user, password=password, host=host, port=port, sslmode=sslmode ) cursor = conn.cursor() # Create tables using IF NOT EXISTS tables_sql = [ """CREATE TABLE IF NOT EXISTS users ( id SERIAL PRIMARY KEY, email VARCHAR UNIQUE, name VARCHAR, created_at TIMESTAMP DEFAULT NOW() )""", """CREATE TABLE IF NOT EXISTS integrations ( id SERIAL PRIMARY KEY, user_id INTEGER REFERENCES users(id), provider VARCHAR, access_token TEXT, refresh_token TEXT, expires_at TIMESTAMP, account_info JSONB, connected BOOLEAN DEFAULT FALSE, created_at TIMESTAMP DEFAULT NOW(), updated_at TIMESTAMP DEFAULT NOW() )""", """CREATE TABLE IF NOT EXISTS assets ( id SERIAL PRIMARY KEY, user_id INTEGER REFERENCES users(id), name VARCHAR, file_path VARCHAR, file_type VARCHAR, product_category VARCHAR, sub_category VARCHAR, size INTEGER, extra_metadata JSONB, extracted_content JSONB, analysis_status VARCHAR DEFAULT 'pending', analyzed_at TIMESTAMP, created_at TIMESTAMP DEFAULT NOW() )""", """CREATE TABLE IF NOT EXISTS posts ( id SERIAL PRIMARY KEY, user_id INTEGER REFERENCES users(id), title VARCHAR, content TEXT, post_type VARCHAR, product_category VARCHAR, scheduled_date TIMESTAMP, status VARCHAR, linkedin_post_id VARCHAR, canva_design_id VARCHAR, assets JSONB, extra_metadata JSONB, created_at TIMESTAMP DEFAULT NOW(), updated_at TIMESTAMP DEFAULT NOW() )""", """CREATE TABLE IF NOT EXISTS campaigns ( id SERIAL PRIMARY KEY, user_id INTEGER REFERENCES users(id), name VARCHAR, date_range_start TIMESTAMP, date_range_end TIMESTAMP, products JSONB, post_types JSONB, posts_per_week INTEGER, status VARCHAR, created_at TIMESTAMP DEFAULT NOW(), updated_at TIMESTAMP DEFAULT NOW() )""" ] for sql in tables_sql: cursor.execute(sql) conn.commit() # Add new columns to assets table if they don't exist (migration) # CockroachDB doesn't support ALTER TABLE in DO blocks, so we check first try: cursor = conn.cursor() # Check if columns exist and add them if they don't cursor.execute(""" SELECT column_name FROM information_schema.columns WHERE table_name='assets' AND column_name='extracted_content' """) if not cursor.fetchone(): cursor.execute("ALTER TABLE assets ADD COLUMN extracted_content JSONB") print("✓ Added extracted_content column") cursor.execute(""" SELECT column_name FROM information_schema.columns WHERE table_name='assets' AND column_name='analysis_status' """) if not cursor.fetchone(): cursor.execute("ALTER TABLE assets ADD COLUMN analysis_status VARCHAR DEFAULT 'pending'") print("✓ Added analysis_status column") cursor.execute(""" SELECT column_name FROM information_schema.columns WHERE table_name='assets' AND column_name='analyzed_at' """) if not cursor.fetchone(): cursor.execute("ALTER TABLE assets ADD COLUMN analyzed_at TIMESTAMP") print("✓ Added analyzed_at column") conn.commit() cursor.close() print("✓ Database migration completed (added new asset columns)") except Exception as migration_error: # Migration might fail if columns already exist, that's okay print(f"Migration note: {migration_error}") conn.close() print("✓ CockroachDB tables created successfully (using direct psycopg2 connection)") return True except Exception as raw_error: print(f"⚠ Table creation failed: {raw_error}") print("✓ Database connection works - tables will be created on first use") return True # Connection works, tables will be created later else: # Real error, not just version parsing raise create_error except Exception as e: error_str = str(e) if "Could not determine version" in error_str: print("⚠ CockroachDB version parsing issue (non-fatal)") print("✓ Database connection works - tables will be created on first use") return True # Connection works, return True else: print(f"Database connection failed: {e}") return False