Spaces:
Running
Running
File size: 3,307 Bytes
bb6d7b4 89ae94f bb6d7b4 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 |
"""
Database initialization for Hugging Face Spaces environment.
This creates an in-memory SQLite database for demo purposes.
"""
import logging
import os
import sqlite3
from sqlalchemy import create_engine, event
from sqlalchemy.orm import sessionmaker
from sqlalchemy.pool import StaticPool
from src.models.base import Base
from src.models.user import User
try:
# Try to import from src.api.security first (for local development)
from src.api.security import get_password_hash
except ImportError:
# Fall back to simplified security module for HF (copied during deployment)
from security_hf import get_password_hash
# Configure logging
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)
# Database URL for SQLite in-memory
DATABASE_URL = "sqlite:///:memory:"
# Create engine with special configuration for in-memory SQLite
engine = create_engine(
DATABASE_URL,
connect_args={"check_same_thread": False},
poolclass=StaticPool,
echo=False
)
# Add pragma for foreign key support
@event.listens_for(engine, "connect")
def set_sqlite_pragma(dbapi_connection, connection_record):
cursor = dbapi_connection.cursor()
cursor.execute("PRAGMA foreign_keys=ON")
cursor.close()
# Create all tables
Base.metadata.create_all(engine)
# Create session factory
SessionLocal = sessionmaker(autocommit=False, autoflush=False, bind=engine)
def init_demo_data():
"""Initialize demo data for the in-memory database."""
session = SessionLocal()
try:
# Check if we already have users
user_count = session.query(User).count()
if user_count == 0:
# Create admin user
admin_user = User(
username="admin",
email="[email protected]",
full_name="Admin User",
hashed_password=get_password_hash("adminpassword"),
is_active=True,
is_superuser=True
)
session.add(admin_user)
# Create regular user
regular_user = User(
username="user",
email="[email protected]",
full_name="Regular User",
hashed_password=get_password_hash("userpassword"),
is_active=True,
is_superuser=False
)
session.add(regular_user)
# Create API user
api_user = User(
username="api_user",
email="[email protected]",
full_name="API User",
hashed_password=get_password_hash("apipassword"),
is_active=True,
is_superuser=False
)
session.add(api_user)
# Commit the session
session.commit()
logger.info("Demo users created successfully")
else:
logger.info("Demo data already exists")
# Here you would add other demo data like threats, indicators, etc.
except Exception as e:
session.rollback()
logger.error(f"Error initializing demo data: {e}")
finally:
session.close()
# Initialize demo data
init_demo_data()
logger.info("Hugging Face database initialized with demo data") |