trax/scripts/db_setup.sh

183 lines
6.4 KiB
Bash
Executable File

#!/bin/bash
# Database setup for Trax project
set -e
# Color codes
GREEN='\033[0;32m'
YELLOW='\033[1;33m'
RED='\033[0;31m'
BLUE='\033[0;34m'
CYAN='\033[0;36m'
NC='\033[0m'
echo -e "${BLUE}🗄️ Trax Database Setup${NC}"
echo "================================"
# Default database configuration
DB_NAME="${TRAX_DB_NAME:-trax}"
DB_USER="${TRAX_DB_USER:-postgres}"
DB_HOST="${TRAX_DB_HOST:-localhost}"
DB_PORT="${TRAX_DB_PORT:-5432}"
TEST_DB_NAME="${DB_NAME}_test"
# Parse arguments
ACTION=${1:-setup}
# Check PostgreSQL is installed
if ! command -v psql &> /dev/null; then
echo -e "${RED}❌ PostgreSQL not installed${NC}"
echo -e "${YELLOW}💡 Install with: brew install postgresql@15${NC}"
exit 1
fi
# Check PostgreSQL is running
if ! pg_isready -h $DB_HOST -p $DB_PORT &> /dev/null; then
echo -e "${RED}❌ PostgreSQL is not running${NC}"
echo -e "${YELLOW}💡 Start with: brew services start postgresql@15${NC}"
exit 1
fi
echo -e "${GREEN}✅ PostgreSQL is running${NC}"
case "$ACTION" in
setup)
echo -e "\n${CYAN}Creating databases...${NC}"
# Create main database
if psql -U $DB_USER -h $DB_HOST -lqt | cut -d \| -f 1 | grep -qw $DB_NAME; then
echo -e " ${YELLOW}⚠️ Database '$DB_NAME' already exists${NC}"
else
createdb -U $DB_USER -h $DB_HOST $DB_NAME
echo -e " ${GREEN}✅ Created database '$DB_NAME'${NC}"
fi
# Create test database
if psql -U $DB_USER -h $DB_HOST -lqt | cut -d \| -f 1 | grep -qw $TEST_DB_NAME; then
echo -e " ${YELLOW}⚠️ Database '$TEST_DB_NAME' already exists${NC}"
else
createdb -U $DB_USER -h $DB_HOST $TEST_DB_NAME
echo -e " ${GREEN}✅ Created database '$TEST_DB_NAME'${NC}"
fi
# Create .env.local with database URL if it doesn't exist
ENV_FILE="$(dirname "$0")/../.env.local"
if [ ! -f "$ENV_FILE" ]; then
echo -e "\n${CYAN}Creating .env.local...${NC}"
cat > "$ENV_FILE" << EOF
# Trax Local Environment Variables
DATABASE_URL=postgresql://$DB_USER@$DB_HOST:$DB_PORT/$DB_NAME
TEST_DATABASE_URL=postgresql://$DB_USER@$DB_HOST:$DB_PORT/$TEST_DB_NAME
# Whisper Model Configuration
WHISPER_MODEL=distil-large-v3
WHISPER_DEVICE=cpu # Change to 'mps' for M3 Mac optimization
WHISPER_COMPUTE_TYPE=int8_float32
# Batch Processing
BATCH_SIZE=10
MAX_WORKERS=4
# Cache Configuration
CACHE_TTL_SECONDS=3600
EOF
echo -e " ${GREEN}✅ Created .env.local with database configuration${NC}"
else
echo -e " ${YELLOW}⚠️ .env.local already exists${NC}"
fi
# Initialize Alembic if not already done
if [ ! -d "alembic" ]; then
echo -e "\n${CYAN}Initializing Alembic...${NC}"
cd "$(dirname "$0")/.."
source .venv/bin/activate 2>/dev/null || true
alembic init alembic
echo -e " ${GREEN}✅ Alembic initialized${NC}"
# Update alembic.ini with database URL
sed -i.bak "s|sqlalchemy.url = .*|sqlalchemy.url = postgresql://$DB_USER@$DB_HOST:$DB_PORT/$DB_NAME|" alembic.ini
rm alembic.ini.bak
echo -e " ${GREEN}✅ Updated alembic.ini with database URL${NC}"
else
echo -e " ${YELLOW}⚠️ Alembic already initialized${NC}"
fi
echo -e "\n${GREEN}✅ Database setup complete!${NC}"
echo ""
echo "Database URLs:"
echo " Main: postgresql://$DB_USER@$DB_HOST:$DB_PORT/$DB_NAME"
echo " Test: postgresql://$DB_USER@$DB_HOST:$DB_PORT/$TEST_DB_NAME"
;;
drop)
echo -e "\n${RED}⚠️ Dropping databases...${NC}"
read -p "Are you sure? This will delete all data! (y/N): " confirm
if [ "$confirm" = "y" ] || [ "$confirm" = "Y" ]; then
dropdb -U $DB_USER -h $DB_HOST --if-exists $DB_NAME
echo -e " ${GREEN}✅ Dropped database '$DB_NAME'${NC}"
dropdb -U $DB_USER -h $DB_HOST --if-exists $TEST_DB_NAME
echo -e " ${GREEN}✅ Dropped database '$TEST_DB_NAME'${NC}"
else
echo " Cancelled"
fi
;;
reset)
echo -e "\n${YELLOW}Resetting databases...${NC}"
$0 drop
$0 setup
;;
status)
echo -e "\n${CYAN}Database Status:${NC}"
echo ""
# Check main database
if psql -U $DB_USER -h $DB_HOST -lqt | cut -d \| -f 1 | grep -qw $DB_NAME; then
echo -e " ${GREEN}✅ Main database '$DB_NAME' exists${NC}"
# Show table count
table_count=$(psql -U $DB_USER -h $DB_HOST -d $DB_NAME -t -c "SELECT COUNT(*) FROM information_schema.tables WHERE table_schema = 'public';" 2>/dev/null || echo "0")
echo -e " Tables: $table_count"
else
echo -e " ${RED}❌ Main database '$DB_NAME' does not exist${NC}"
fi
# Check test database
if psql -U $DB_USER -h $DB_HOST -lqt | cut -d \| -f 1 | grep -qw $TEST_DB_NAME; then
echo -e " ${GREEN}✅ Test database '$TEST_DB_NAME' exists${NC}"
else
echo -e " ${RED}❌ Test database '$TEST_DB_NAME' does not exist${NC}"
fi
# Check Alembic
if [ -d "alembic" ]; then
echo -e " ${GREEN}✅ Alembic is initialized${NC}"
# Check current revision
cd "$(dirname "$0")/.."
source .venv/bin/activate 2>/dev/null || true
current_rev=$(alembic current 2>/dev/null | grep -o '[a-f0-9]\{12\}' | head -1 || echo "none")
echo -e " Current revision: $current_rev"
else
echo -e " ${YELLOW}⚠️ Alembic not initialized${NC}"
fi
;;
*)
echo "Usage: $0 [setup|drop|reset|status]"
echo ""
echo "Commands:"
echo " setup - Create databases and initialize Alembic"
echo " drop - Drop all databases (requires confirmation)"
echo " reset - Drop and recreate databases"
echo " status - Show database status"
echo ""
echo "Environment variables:"
echo " TRAX_DB_NAME - Database name (default: trax)"
echo " TRAX_DB_USER - Database user (default: postgres)"
echo " TRAX_DB_HOST - Database host (default: localhost)"
echo " TRAX_DB_PORT - Database port (default: 5432)"
exit 1
;;
esac