function/app.py

import json
import os
import time
from flask import Flask, Response, jsonify, redirect, render_template, render_template_string, request, url_for
import jinja_partials
from jinja2_fragments import render_block
import requests
from extensions import db, htmx, init_app, login_manager
from services import create_http_function_view_model, create_http_functions_view_model
from flask_login import current_user, login_required
from werkzeug.security import check_password_hash, generate_password_hash
import os
from dotenv import load_dotenv
from routes.timer import timer
from routes.test import test
from routes.home import home
from routes.http import http
from routes.llm import llm
from routes.auth import auth
from routes.settings import settings
from routes.community import community
from routes.shared_env import shared_env
from constants import DEFAULT_FUNCTION_NAME, DEFAULT_SCRIPT, DEFAULT_ENVIRONMENT
from flask_apscheduler import APScheduler
import asyncio
import aiohttp
from concurrent.futures import ThreadPoolExecutor

# Load environment variables from .env file in non-production environments
if os.environ.get('FLASK_ENV') != 'production':
    load_dotenv()

app = Flask(__name__)

app.config.from_pyfile('config.py')
app.secret_key = os.environ.get('SECRET_KEY', '2a661781919643cb8a5a8bc57642d99f')
login_manager.init_app(app)
login_manager.login_view = "auth.login"
jinja_partials.register_extensions(app)

# Remove scheduler configuration and initialization
init_app(app)

app.register_blueprint(timer, url_prefix='/timer')
app.register_blueprint(test, url_prefix='/test')
app.register_blueprint(home, url_prefix='/home')
app.register_blueprint(http, url_prefix='/http')
app.register_blueprint(llm, url_prefix='/llm')
app.register_blueprint(auth, url_prefix='/auth')
app.register_blueprint(settings, url_prefix='/settings')
app.register_blueprint(community, url_prefix='/community')
app.register_blueprint(shared_env, url_prefix='/shared_env')

# Swith to inter app routing, which results in speed up from ~400ms to ~270ms
# https://stackoverflow.com/questions/76886643/linking-two-not-exposed-dokku-apps
NODE_API_URL = os.environ.get('NODE_API_URL', 'http://isolator.web:5000/execute')
DENO_API_URL = os.environ.get('DENO_API_URL', 'http://deno-isolator.web:5000/execute')
PYTHON_API_URL = os.environ.get('PYTHON_API_URL', 'http://python-isolator.web:5000/execute')

def map_isolator_response_to_flask_response(response):
    """
    Maps a Node.js response to a Flask response.

    :param nodejs_response: The response from Node.js, expected to be a dictionary
                            with keys 'body', 'headers', and 'status'.
    :return: Flask Response object
    """
    result = response.get('result', {})
    body = result.get('body', '')
    headers = result.get('headers', {})
    status = result.get('status', 200)


    # Convert body to JSON if it's a dictionary
    body = str(body)

    return Response(response=body, status=status, headers=headers)


@ app.route("/", methods=["GET"])
def landing_page():
    public_functions = db.get_public_http_functions()
    # Limit to top 6 for the landing page
    public_functions = public_functions[:6] if public_functions else []
    return render_template("landing_page.html", name='Try me', script=DEFAULT_SCRIPT, environment_info=DEFAULT_ENVIRONMENT, public_functions=public_functions)

@app.route("/documentation", methods=["GET"])
def documentation():
    return render_template("documentation.html")

@app.route('/execute', methods=['POST'])
async def execute_code():
    try:
        # Extract code and convert request to a format acceptable by Node.js app
        code = request.json.get('code')
        runtime = request.json.get('runtime', 'node')  # Default to node
        if runtime == 'deno':
            api_url = DENO_API_URL
        elif runtime == 'python':
            api_url = PYTHON_API_URL
        else:
            api_url = NODE_API_URL

        request_obj = {
            'method': request.method,
            'headers': dict(request.headers),
            'body': request.json,
            'url': request.url,
            'path': '/',
        }
        environment = request.json.get('environment_info')
        environment_json = json.loads(environment)

        # Call the selected isolator API asynchronously
        async with aiohttp.ClientSession() as session:
            async with session.post(api_url, json={'code': code, 'request': request_obj, 'environment': environment_json, 'name': "anonymous"}) as response:
                response_data = await response.json()

        # check if playground=true is in the query string
        if request.args.get('playground') == 'true':
            return response_data

        # Map the Node.js response to Flask response
        flask_response = map_isolator_response_to_flask_response(response_data)
        return flask_response

    except Exception as e:
        return jsonify({'error': str(e)}), 500
    
@app.route('/f/<int:user_id>/<path:function>', methods=['GET', 'POST', 'PUT', 'DELETE', 'PATCH', 'OPTIONS', 'HEAD'])
async def execute_http_function(user_id, function):
    try:
        start_time = time.time()
        # Split the function_path into the function name and the sub-path
        parts = function.split('/', 1)
        function_name = parts[0]
        sub_path = '/' + parts[1] if len(parts) > 1 else '/'

        http_function = db.get_http_function(user_id, function_name)
        if not http_function:
            return jsonify({'error': 'Function not found'}), 404
        
        code = http_function['script_content']
        environment_info = http_function['environment_info']
        runtime = http_function.get('runtime', 'node') # Default to node
        
        # Ensure environment is a dictionary
        if isinstance(environment_info, str) and environment_info:
            try:
                environment = json.loads(environment_info)
            except json.JSONDecodeError:
                environment = {}
        elif isinstance(environment_info, dict):
            environment = environment_info
        else:
            environment = {}
        is_public = http_function['is_public']
        log_request = http_function['log_request']
        log_response = http_function['log_response']
        version_number = http_function['version_number']

        # Check if the function is public, if not check if the user is authenticated and owns the function
        if not is_public:
            is_authorized = False
            
            # 1. Session Authentication
            if current_user.is_authenticated and int(current_user.id) == user_id:
                is_authorized = True
            
            # 2. API Key Authentication
            elif 'X-API-Key' in request.headers:
                api_key_value = request.headers.get('X-API-Key')
                api_key = db.get_api_key(api_key_value)
                
                if api_key and api_key['user_id'] == user_id:
                    # Check rate limit
                    if not db.check_and_increment_api_key_usage(api_key['id']):
                        return jsonify({'error': 'Rate limit exceeded'}), 429

                    # Check Scopes
                    scopes = api_key['scopes']
                    if isinstance(scopes, str):
                        scopes = json.loads(scopes)
                        
                    if "*" in scopes or f"function:{http_function['id']}" in scopes:
                        is_authorized = True
                        db.update_api_key_last_used(api_key['id'])

            if not is_authorized:
                if not current_user.is_authenticated:
                     return redirect(url_for('auth.login', next=request.url))
                return jsonify({'error': 'Function belongs to another user'}), 404

        request_data = {
            'method': request.method,
            'headers': dict(request.headers),
            'url': request.url,
            'path': sub_path,
        }
        
        # Add JSON data if it exists
        if request.is_json:
            request_data['json'] = request.get_json()

        # Add form data if it exists
        if request.form:
            request_data['form'] = {
                key.rstrip('[]'): request.form.getlist(key) if key.endswith('[]') or len(request.form.getlist(key)) > 1 else request.form.getlist(key)[0]
                for key in request.form.keys()
            }

        # Add query parameters if they exist
        if request.args:
            request_data['query'] = {
                key.rstrip('[]'): request.args.getlist(key) if key.endswith('[]') or len(request.args.getlist(key)) > 1 else request.args.getlist(key)[0]
                for key in request.args.keys()
            }

        # Add plain text data if it exists
        if request.data and not request.is_json:
            request_data['text'] = request.data.decode('utf-8')

        # Load and inject shared environments (namespaced)
        shared_envs = db.execute('''
            SELECT se.id, se.name, se.environment
            FROM http_function_shared_envs hfse
            JOIN shared_environments se ON hfse.shared_env_id = se.id
            WHERE hfse.http_function_id = %s
            ORDER BY se.name
        ''', [http_function['id']])
        
        # Inject shared environments as nested objects
        combined_environment = environment.copy()
        shared_env_map = {}  # Track shared env IDs for later extraction
        if shared_envs:
            for se in shared_envs:
                env_data = json.loads(se['environment']) if isinstance(se['environment'], str) else se['environment']
                combined_environment[se['name']] = env_data
                shared_env_map[se['name']] = se['id']

        # Call the Node.js API asynchronously
        if runtime == 'deno':
            api_url = DENO_API_URL
        elif runtime == 'python':
            api_url = PYTHON_API_URL
        else:
            api_url = NODE_API_URL
        async with aiohttp.ClientSession() as session:
            async with session.post(api_url, json={'code': code, 'request': request_data, 'environment': combined_environment, 'name': function_name}) as response:
                response_data = await response.json()

        # Extract and persist shared environment mutations
        returned_env = response_data['environment']
        function_specific_env = {}
        
        # Separate function-specific properties from shared environments
        for key, value in returned_env.items():
            if key in shared_env_map:
                # This is a shared environment - save it back
                db.execute(
                    'UPDATE shared_environments SET environment=%s, updated_at=NOW() WHERE id=%s',
                    [json.dumps(value), shared_env_map[key]],
                    commit=True
                )
            else:
                # This is function-specific - keep it
                function_specific_env[key] = value
        
        # Update function's own environment (without shared envs)
        db.update_http_function_environment_info_and_invoked_count(user_id, function_name, function_specific_env)

        # Calculate execution time in milliseconds
        execution_time = (time.time() - start_time) * 1000

        db.add_http_function_invocation(
            http_function['id'],
            response_data['status'],
            request_data if log_request else {},
            response_data['result'] if (log_response or response_data['status'] != 'SUCCESS') else {},
            response_data['logs'],
            version_number,
            execution_time)
        
        if response_data['status'] != 'SUCCESS':
            return render_template("function_error.html", function_name=function_name ,error=response_data['result'], logs=response_data['logs'])

        # Map the Node.js response to Flask response
        flask_response = map_isolator_response_to_flask_response(response_data)
        return flask_response

    except Exception as e:
        return jsonify({'error': str(e)}), 500


if __name__ == '__main__':
    # Bind to PORT if defined, otherwise default to 5000.
    port = int(os.environ.get('PORT', 5000))
    app.run(host='127.0.0.1', port=port)

@app.teardown_appcontext
def teardown_db(exception):
    db.close_conn()