|
from flask import Flask, request, jsonify, Response |
|
from flask_executor import Executor |
|
import uuid |
|
import time |
|
import logging |
|
from collections import defaultdict, deque |
|
from threading import Lock |
|
|
|
app = Flask(__name__) |
|
executor = Executor(app) |
|
|
|
|
|
transfers = defaultdict(deque) |
|
locks = defaultdict(Lock) |
|
metadata = {} |
|
CHUNK_TIMEOUT = 300 |
|
|
|
@app.route('/create_transfer', methods=['POST']) |
|
def create_transfer(): |
|
transfer_id = str(uuid.uuid4()) |
|
metadata[transfer_id] = { |
|
'filename': request.json.get('filename', 'file'), |
|
'created_at': time.time(), |
|
'completed': False |
|
} |
|
return jsonify({'transfer_id': transfer_id}) |
|
|
|
@app.route('/upload/<transfer_id>', methods=['POST']) |
|
def upload_chunk(transfer_id): |
|
if transfer_id not in metadata: |
|
return jsonify({'error': 'Invalid transfer ID'}), 404 |
|
|
|
with locks[transfer_id]: |
|
chunk = request.data |
|
if chunk: |
|
transfers[transfer_id].append(chunk) |
|
else: |
|
metadata[transfer_id]['completed'] = True |
|
|
|
return jsonify({'status': 'ok'}) |
|
|
|
@app.route('/stream/<transfer_id>') |
|
def stream_chunks(transfer_id): |
|
def generate(): |
|
last_activity = time.time() |
|
|
|
while True: |
|
with locks[transfer_id]: |
|
if transfers[transfer_id]: |
|
chunk = transfers[transfer_id].popleft() |
|
last_activity = time.time() |
|
yield chunk |
|
elif metadata.get(transfer_id, {}).get('completed', False): |
|
break |
|
elif time.time() - last_activity > CHUNK_TIMEOUT: |
|
break |
|
|
|
time.sleep(0.1) |
|
|
|
|
|
if transfer_id in metadata: |
|
del metadata[transfer_id] |
|
del transfers[transfer_id] |
|
del locks[transfer_id] |
|
|
|
return Response( |
|
generate(), |
|
mimetype='application/octet-stream', |
|
headers={ |
|
'Content-Disposition': f'attachment; filename="{metadata[transfer_id]["filename"]}"', |
|
'Transfer-Encoding': 'chunked' |
|
} |
|
) |
|
|
|
if __name__ == '__main__': |
|
app.run(debug=True) |