Spaces:

dn-scribe
/

attstore

Runtime error

feat: support tenant deletion

38536e2 10 months ago

5.15 kB

# Import job_manager for healthz endpoint
from job_manager import job_manager
import os
import logging
from flask import Flask, request, jsonify
from flask_sqlalchemy import SQLAlchemy
from flask_jwt_extended import JWTManager
from flask_migrate import Migrate
from sqlalchemy.orm import DeclarativeBase
from werkzeug.middleware.proxy_fix import ProxyFix
from datetime import timedelta


# Configure logging
logging.basicConfig(level=logging.DEBUG)

class Base(DeclarativeBase):
    pass

db = SQLAlchemy(model_class=Base)
jwt = JWTManager()

# Create the app
app = Flask(__name__)
os.makedirs(app.instance_path, exist_ok=True)
# Ensure templates auto-reload for development
app.config['TEMPLATES_AUTO_RELOAD'] = True
app.jinja_env.auto_reload = True
app.secret_key = os.environ.get("SESSION_SECRET", "dev-secret-key-change-in-production")
# Honor proxy-provided protocol/host/port so generated URLs match PUBLIC_BASE_URL
app.wsgi_app = ProxyFix(app.wsgi_app, x_proto=1, x_host=1, x_port=1)

# Configure JWT
app.config['JWT_SECRET_KEY'] = os.environ.get("JWT_SECRET_KEY", "jwt-secret-change-in-production")
app.config['JWT_ACCESS_TOKEN_EXPIRES'] = timedelta(hours=1)  # Tokens expire in 1 hour

# Configure the database
app.config["SQLALCHEMY_DATABASE_URI"] = os.environ.get("DATABASE_URL", "sqlite:///attestation_store.db")
app.config["SQLALCHEMY_ENGINE_OPTIONS"] = {
    "pool_recycle": 300,
    "pool_pre_ping": True,
}

uri = app.config["SQLALCHEMY_DATABASE_URI"]
if uri.startswith("postgresql"):
    uri = f"postgresql://***@{uri.split('@')[1]}"  # Remove user info for logging
logging.info(f"Using database at {uri}")

# Initialize extensions
db.init_app(app)
jwt.init_app(app)
migrate = Migrate(app, db)

# Import config and setup
from config import setup_config
setup_config(app)

# Import models so SQLAlchemy metadata is populated for migrations
import models  # noqa: E402,F401

def register_blueprints(app):
    from api.evidence import evidence_bp
    from api.login import login_bp
    from api.syslog import syslog_bp
    from web.routes import web_bp
    from api.webhooks import webhooks_bp
    from web.events import events_bp
    from api.badges import badges_bp
    from api.git_policies import git_policies_bp
    from api.saml_config import saml_config_bp
    from saml import saml_bp

    app.register_blueprint(evidence_bp, url_prefix='/')
    app.register_blueprint(login_bp, url_prefix='/v1')
    app.register_blueprint(syslog_bp, url_prefix='/')
    app.register_blueprint(webhooks_bp, url_prefix='/api')
    app.register_blueprint(web_bp)
    app.register_blueprint(events_bp)
    app.register_blueprint(badges_bp, url_prefix='/')
    app.register_blueprint(git_policies_bp, url_prefix='/')
    app.register_blueprint(saml_config_bp)
    app.register_blueprint(saml_bp)

register_blueprints(app)

# Error handlers
@app.errorhandler(400)
def bad_request(error):
    """Handle 400 Bad Request errors"""
    logging.warning(f"400 Bad Request: {request.method} {request.path}")
    accept = request.headers.get('Accept', '')
    if 'application/json' in accept or request.path.startswith('/api/'):
        return jsonify({'error': 'bad_request', 'message': str(error)}), 400
    else:
        return error, 400

@app.errorhandler(403)
def forbidden(error):
    """Handle 403 Forbidden errors"""
    logging.warning(f"403 Forbidden: {request.method} {request.path}")
    accept = request.headers.get('Accept', '')
    if 'application/json' in accept or request.path.startswith('/api/'):
        return jsonify({'error': 'forbidden', 'message': str(error)}), 403
    else:
        return error, 403

@app.errorhandler(404)
def not_found(error):
    """Handle 404 errors"""
    logging.info(f"404 Not Found: {request.method} {request.path}")
    accept = request.headers.get('Accept', '')
    if 'application/json' in accept or request.path.startswith('/api/'):
        return jsonify({'error': 'not_found', 'message': str(error)}), 404
    else:
        return error, 404

@app.errorhandler(500)
def internal_error(error):
    """Handle 500 Internal Server errors"""
    logging.error(f"500 Internal Server Error: {request.method} {request.path} - {error}")
    accept = request.headers.get('Accept', '')
    if 'application/json' in accept or request.path.startswith('/api/'):
        return jsonify({'error': 'internal_server_error', 'message': 'An internal error occurred'}), 500
    else:
        return error, 500

# Public health endpoint (no auth)
@app.get('/healthz')
def healthz():
    """Lightweight health check endpoint used by LB/containers. Also reports background job status."""
    try:
        # Optional: ping DB to reflect readiness
        from sqlalchemy import text
        with db.engine.connect() as conn:
            conn.execute(text('SELECT 1'))
        db_status = 'ok'
        status_code = 200
    except Exception:
        db_status = 'unavailable'
        status_code = 503

    job_count = job_manager.job_count()
    return {"status": "ok", "db": db_status, "jobs_running_or_queued": job_count}, status_code

# Expose app and db for flask CLI/migrate
if __name__ == "__main__":
    app.run()

	# Import job_manager for healthz endpoint
	from job_manager import job_manager
	import os
	import logging
	from flask import Flask, request, jsonify
	from flask_sqlalchemy import SQLAlchemy
	from flask_jwt_extended import JWTManager
	from flask_migrate import Migrate
	from sqlalchemy.orm import DeclarativeBase
	from werkzeug.middleware.proxy_fix import ProxyFix
	from datetime import timedelta


	# Configure logging
	logging.basicConfig(level=logging.DEBUG)

	class Base(DeclarativeBase):
	pass

	db = SQLAlchemy(model_class=Base)
	jwt = JWTManager()

	# Create the app
	app = Flask(__name__)
	os.makedirs(app.instance_path, exist_ok=True)
	# Ensure templates auto-reload for development
	app.config['TEMPLATES_AUTO_RELOAD'] = True
	app.jinja_env.auto_reload = True
	app.secret_key = os.environ.get("SESSION_SECRET", "dev-secret-key-change-in-production")
	# Honor proxy-provided protocol/host/port so generated URLs match PUBLIC_BASE_URL
	app.wsgi_app = ProxyFix(app.wsgi_app, x_proto=1, x_host=1, x_port=1)

	# Configure JWT
	app.config['JWT_SECRET_KEY'] = os.environ.get("JWT_SECRET_KEY", "jwt-secret-change-in-production")
	app.config['JWT_ACCESS_TOKEN_EXPIRES'] = timedelta(hours=1) # Tokens expire in 1 hour

	# Configure the database
	app.config["SQLALCHEMY_DATABASE_URI"] = os.environ.get("DATABASE_URL", "sqlite:///attestation_store.db")
	app.config["SQLALCHEMY_ENGINE_OPTIONS"] = {
	"pool_recycle": 300,
	"pool_pre_ping": True,
	}

	uri = app.config["SQLALCHEMY_DATABASE_URI"]
	if uri.startswith("postgresql"):
	uri = f"postgresql://***@{uri.split('@')[1]}" # Remove user info for logging
	logging.info(f"Using database at {uri}")

	# Initialize extensions
	db.init_app(app)
	jwt.init_app(app)
	migrate = Migrate(app, db)

	# Import config and setup
	from config import setup_config
	setup_config(app)

	# Import models so SQLAlchemy metadata is populated for migrations
	import models # noqa: E402,F401

	def register_blueprints(app):
	from api.evidence import evidence_bp
	from api.login import login_bp
	from api.syslog import syslog_bp
	from web.routes import web_bp
	from api.webhooks import webhooks_bp
	from web.events import events_bp
	from api.badges import badges_bp
	from api.git_policies import git_policies_bp
	from api.saml_config import saml_config_bp
	from saml import saml_bp

	app.register_blueprint(evidence_bp, url_prefix='/')
	app.register_blueprint(login_bp, url_prefix='/v1')
	app.register_blueprint(syslog_bp, url_prefix='/')
	app.register_blueprint(webhooks_bp, url_prefix='/api')
	app.register_blueprint(web_bp)
	app.register_blueprint(events_bp)
	app.register_blueprint(badges_bp, url_prefix='/')
	app.register_blueprint(git_policies_bp, url_prefix='/')
	app.register_blueprint(saml_config_bp)
	app.register_blueprint(saml_bp)

	register_blueprints(app)

	# Error handlers
	@app.errorhandler(400)
	def bad_request(error):
	"""Handle 400 Bad Request errors"""
	logging.warning(f"400 Bad Request: {request.method} {request.path}")
	accept = request.headers.get('Accept', '')
	if 'application/json' in accept or request.path.startswith('/api/'):
	return jsonify({'error': 'bad_request', 'message': str(error)}), 400
	else:
	return error, 400

	@app.errorhandler(403)
	def forbidden(error):
	"""Handle 403 Forbidden errors"""
	logging.warning(f"403 Forbidden: {request.method} {request.path}")
	accept = request.headers.get('Accept', '')
	if 'application/json' in accept or request.path.startswith('/api/'):
	return jsonify({'error': 'forbidden', 'message': str(error)}), 403
	else:
	return error, 403

	@app.errorhandler(404)
	def not_found(error):
	"""Handle 404 errors"""
	logging.info(f"404 Not Found: {request.method} {request.path}")
	accept = request.headers.get('Accept', '')
	if 'application/json' in accept or request.path.startswith('/api/'):
	return jsonify({'error': 'not_found', 'message': str(error)}), 404
	else:
	return error, 404

	@app.errorhandler(500)
	def internal_error(error):
	"""Handle 500 Internal Server errors"""
	logging.error(f"500 Internal Server Error: {request.method} {request.path} - {error}")
	accept = request.headers.get('Accept', '')
	if 'application/json' in accept or request.path.startswith('/api/'):
	return jsonify({'error': 'internal_server_error', 'message': 'An internal error occurred'}), 500
	else:
	return error, 500

	# Public health endpoint (no auth)
	@app.get('/healthz')
	def healthz():
	"""Lightweight health check endpoint used by LB/containers. Also reports background job status."""
	try:
	# Optional: ping DB to reflect readiness
	from sqlalchemy import text
	with db.engine.connect() as conn:
	conn.execute(text('SELECT 1'))
	db_status = 'ok'
	status_code = 200
	except Exception:
	db_status = 'unavailable'
	status_code = 503

	job_count = job_manager.job_count()
	return {"status": "ok", "db": db_status, "jobs_running_or_queued": job_count}, status_code

	# Expose app and db for flask CLI/migrate
	if __name__ == "__main__":
	app.run()