app.py

Download
python 240 lines 6.8 KB
  1"""
  2Flask application with PostgreSQL and Redis integration.
  3
  4This demonstrates:
  5- Database connections in a containerized environment
  6- Cache integration for improved performance
  7- Environment-based configuration
  8- Health checks for all dependencies
  9"""
 10
 11from flask import Flask, jsonify, request
 12import psycopg2
 13from psycopg2.extras import RealDictCursor
 14import redis
 15import os
 16import logging
 17import json
 18from datetime import datetime
 19
 20# Configure logging
 21logging.basicConfig(
 22    level=logging.INFO,
 23    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
 24)
 25logger = logging.getLogger(__name__)
 26
 27app = Flask(__name__)
 28
 29# ============================================================================
 30# Configuration
 31# ============================================================================
 32DATABASE_URL = os.getenv('DATABASE_URL', 'postgresql://user:password@localhost:5432/mydb')
 33REDIS_URL = os.getenv('REDIS_URL', 'redis://localhost:6379/0')
 34
 35
 36# ============================================================================
 37# Database Connection
 38# ============================================================================
 39def get_db_connection():
 40    """Create a connection to PostgreSQL database."""
 41    try:
 42        conn = psycopg2.connect(DATABASE_URL)
 43        return conn
 44    except Exception as e:
 45        logger.error(f"Database connection error: {e}")
 46        raise
 47
 48
 49def init_db():
 50    """Initialize database with required tables."""
 51    try:
 52        conn = get_db_connection()
 53        cur = conn.cursor()
 54
 55        # Create a simple visitors table
 56        cur.execute("""
 57            CREATE TABLE IF NOT EXISTS visitors (
 58                id SERIAL PRIMARY KEY,
 59                ip_address VARCHAR(45),
 60                user_agent TEXT,
 61                visited_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
 62            )
 63        """)
 64
 65        conn.commit()
 66        cur.close()
 67        conn.close()
 68        logger.info("Database initialized successfully")
 69    except Exception as e:
 70        logger.error(f"Database initialization error: {e}")
 71
 72
 73# ============================================================================
 74# Redis Connection
 75# ============================================================================
 76def get_redis_client():
 77    """Create a Redis client connection."""
 78    try:
 79        client = redis.from_url(REDIS_URL, decode_responses=True)
 80        return client
 81    except Exception as e:
 82        logger.error(f"Redis connection error: {e}")
 83        raise
 84
 85
 86# ============================================================================
 87# Routes
 88# ============================================================================
 89@app.route('/')
 90def index():
 91    """Main endpoint with visitor tracking."""
 92    try:
 93        # Get visitor info
 94        ip = request.remote_addr
 95        user_agent = request.headers.get('User-Agent', 'Unknown')
 96
 97        # Store in database
 98        conn = get_db_connection()
 99        cur = conn.cursor()
100        cur.execute(
101            "INSERT INTO visitors (ip_address, user_agent) VALUES (%s, %s) RETURNING id",
102            (ip, user_agent)
103        )
104        visitor_id = cur.fetchone()[0]
105        conn.commit()
106        cur.close()
107        conn.close()
108
109        # Increment visitor count in Redis
110        r = get_redis_client()
111        total_visits = r.incr('total_visits')
112
113        return jsonify({
114            'message': 'Welcome to the Flask + PostgreSQL + Redis demo!',
115            'visitor_id': visitor_id,
116            'total_visits': total_visits,
117            'timestamp': datetime.now().isoformat()
118        })
119
120    except Exception as e:
121        logger.error(f"Error in index route: {e}")
122        return jsonify({'error': str(e)}), 500
123
124
125@app.route('/stats')
126def stats():
127    """Get visitor statistics."""
128    try:
129        # Check cache first
130        r = get_redis_client()
131        cached_stats = r.get('stats_cache')
132
133        if cached_stats:
134            logger.info("Returning cached statistics")
135            return jsonify(json.loads(cached_stats))
136
137        # Query database
138        conn = get_db_connection()
139        cur = conn.cursor(cursor_factory=RealDictCursor)
140
141        cur.execute("SELECT COUNT(*) as total FROM visitors")
142        total = cur.fetchone()['total']
143
144        cur.execute("""
145            SELECT ip_address, COUNT(*) as visits
146            FROM visitors
147            GROUP BY ip_address
148            ORDER BY visits DESC
149            LIMIT 5
150        """)
151        top_visitors = cur.fetchall()
152
153        cur.close()
154        conn.close()
155
156        stats_data = {
157            'total_visitors': total,
158            'top_visitors': [dict(v) for v in top_visitors],
159            'cached': False,
160            'timestamp': datetime.now().isoformat()
161        }
162
163        # Cache for 60 seconds
164        r.setex('stats_cache', 60, json.dumps(stats_data))
165
166        return jsonify(stats_data)
167
168    except Exception as e:
169        logger.error(f"Error in stats route: {e}")
170        return jsonify({'error': str(e)}), 500
171
172
173@app.route('/health')
174def health():
175    """
176    Comprehensive health check endpoint.
177
178    Checks:
179    - Application is running
180    - Database connectivity
181    - Redis connectivity
182    """
183    health_status = {
184        'status': 'healthy',
185        'checks': {}
186    }
187
188    # Check database
189    try:
190        conn = get_db_connection()
191        cur = conn.cursor()
192        cur.execute('SELECT 1')
193        cur.close()
194        conn.close()
195        health_status['checks']['database'] = 'healthy'
196    except Exception as e:
197        health_status['status'] = 'unhealthy'
198        health_status['checks']['database'] = f'unhealthy: {str(e)}'
199
200    # Check Redis
201    try:
202        r = get_redis_client()
203        r.ping()
204        health_status['checks']['redis'] = 'healthy'
205    except Exception as e:
206        health_status['status'] = 'unhealthy'
207        health_status['checks']['redis'] = f'unhealthy: {str(e)}'
208
209    status_code = 200 if health_status['status'] == 'healthy' else 503
210    return jsonify(health_status), status_code
211
212
213@app.route('/cache/clear')
214def clear_cache():
215    """Clear Redis cache."""
216    try:
217        r = get_redis_client()
218        r.delete('stats_cache')
219        return jsonify({'message': 'Cache cleared successfully'})
220    except Exception as e:
221        logger.error(f"Error clearing cache: {e}")
222        return jsonify({'error': str(e)}), 500
223
224
225# ============================================================================
226# Application Startup
227# ============================================================================
228if __name__ == '__main__':
229    # Initialize database on startup
230    try:
231        init_db()
232    except Exception as e:
233        logger.error(f"Failed to initialize database: {e}")
234        logger.warning("Continuing anyway - database might not be ready yet")
235
236    port = int(os.getenv('PORT', 5000))
237    logger.info(f"Starting Flask application on port {port}")
238
239    app.run(host='0.0.0.0', port=port, debug=False)