# -*- coding:utf-8 -*-
# This demo demonstrates the concurrency problem of MongDB's GridFS interface.

import os, tempfile, multiprocessing, time

import pymongo, gridfs
import mongoengine

DB_HOST = '127.0.0.1'
DB_PORT = 27017

PROCESS_NUM = 2

class TestDoc(mongoengine.Document):
    field_a = mongoengine.IntField()
    field_b = mongoengine.StringField()
    field_c = mongoengine.DictField()
    field_fs = mongoengine.FileField()

def _conect_database():
    mongoengine.connect('TestDB', host=DB_HOST, port=DB_PORT)

def _create_database():
    connection = pymongo.Connection(DB_HOST, DB_PORT)
    db = connection['TestDB']
    connection.close()

    _conect_database()

def _drop_database():
    connection = pymongo.Connection(DB_HOST, DB_PORT)
    connection.drop_database('TestDB')

def _write_file_field(obj_id):
    try:
        with tempfile.TemporaryFile() as f:
            # Create temporary file
            f.write("Hello XingCloud!\n" * 10)
            f.flush()

            obj = TestDoc.objects.with_id(obj_id)
            file_name = 'PID_' + str(os.getpid())

            f.seek(0)
            obj.field_fs.replace(f, filename=file_name)
            obj.save()
    except TestDoc.DoesNotExist as e:
        raise e

if __name__ == '__main__':
    _create_database()

    # Crate document object
    print 'Create document object...'
    obj = TestDoc(
            field_a = 0,
            field_b = 'Hello XingCloud!',
            field_fs = '',  # Without initial value, data cannot be written to
                            # FieldFiled.
        )
    obj.save()

    # Write GridFS field with multil-processes
    for i in xrange(PROCESS_NUM):
        multiprocessing.Process(
                target = _write_file_field,
                args = (obj.id, ),
            ).start()

    time.sleep(1)
    print 'Writing finished...'

    obj.reload()

    print '>>> Doc:'
    print '    ID: %s' % obj.id
    print '    File: %s' % obj.field_fs.grid_id

    connection = pymongo.Connection(DB_HOST, DB_PORT)
    db = connection['TestDB']
    grid_fs = gridfs.GridFS(db)

    print '>>> GridFS:'
    print grid_fs.list()

    _drop_database()