<commit_before>import pytest
from trough.wsgi.segment_manager import server
import ujson
from trough.settings import settings
import doublethink
import rethinkdb as r
import requests # :-\ urllib3?

@pytest.fixture(scope="module")
def segment_manager_server():
    server.testing = True
    return server.test_client()

def test_simple_provision(segment_manager_server):
    result = segment_manager_server.get('/')
    assert result.status == '405 METHOD NOT ALLOWED'

    # hasn't been provisioned yet
    result = segment_manager_server.post('/', data='test_simple_provision_segment')
    assert result.status_code == 200
    assert result.mimetype == 'text/plain'
    assert b''.join(result.response).endswith(b':6222/?segment=test_simple_provision_segment')

    # now it has already been provisioned
    result = segment_manager_server.post('/', data='test_simple_provision_segment')
    assert result.status_code == 200
    assert result.mimetype == 'text/plain'
    assert b''.join(result.response).endswith(b':6222/?segment=test_simple_provision_segment')

def test_provision(segment_manager_server):
    result = segment_manager_server.get('/provision')
    assert result.status == '405 METHOD NOT ALLOWED'

    # hasn't been provisioned yet
    result = segment_manager_server.post(
            '/provision', content_type='application/json',
            data=ujson.dumps({'segment':'test_provision_segment'}))
    assert result.status_code == 200
    assert result.mimetype == 'application/json'
    result_bytes = b''.join(result.response)
    result_dict = ujson.loads(result_bytes) # ujson accepts bytes! 😻
    assert result_dict['write_url'].endswith(':6222/?segment=test_provision_segment')

    # now it has already been provisioned
    result = segment_manager_server.post(
            '/provision', content_type='application/json',
            data=ujson.dumps({'segment':'test_provision_segment'}))
    assert result.status_code == 200
    assert result.mimetype == 'application/json'
    result_bytes = b''.join(result.response)
    result_dict = ujson.loads(result_bytes)
    assert result_dict['write_url'].endswith(':6222/?segment=test_provision_segment')

def test_provision_with_schema(segment_manager_server):
    schema = '''CREATE TABLE test (id INTEGER PRIMARY KEY AUTOINCREMENT, test varchar(4));
INSERT INTO test (test) VALUES ("test");'''
    # create a schema by submitting sql
    result = segment_manager_server.put(
            '/schema/test1/sql', content_type='applicaton/sql', data=schema)
    assert result.status_code == 201

    # provision a segment with that schema
    result = segment_manager_server.post(
            '/provision', content_type='application/json',
            data=ujson.dumps({'segment':'test_provision_with_schema_1', 'schema':'test1'}))
    assert result.status_code == 200
    assert result.mimetype == 'application/json'
    result_bytes = b''.join(result.response)
    result_dict = ujson.loads(result_bytes) # ujson accepts bytes! 😻
    assert result_dict['write_url'].endswith(':6222/?segment=test_provision_with_schema_1')

    # get db read url from rethinkdb
    rethinker = doublethink.Rethinker(
            servers=settings['RETHINKDB_HOSTS'], db='trough_configuration')
    query = rethinker.table('services').get_all('test_provision_with_schema_1', index='segment').filter({'role': 'trough-read'}).filter(lambda svc: r.now().sub(svc['last_heartbeat']).lt(svc['ttl'])).order_by('load')[0]
    healthy_segment = query.run()
    read_url = healthy_segment.get('url')
    assert read_url.endswith(':6444/?segment=test_provision_with_schema_1')

    # run a query to check that the schema was used
    sql = 'SELECT * FROM test;'
    with requests.post(read_url, stream=True, data=sql) as response:
        assert response.status_code == 200
        result = ujson.loads(response.text)
        assert result == [{'test': 'test', 'id': 1}]

    # delete the schema from rethinkdb for the sake of other tests
    rethinker = doublethink.Rethinker(
            servers=settings['RETHINKDB_HOSTS'], db='trough_configuration')
    result = rethinker.table('schema').get('test1').delete().run()
    assert result == {'deleted': 1, 'inserted': 0, 'skipped': 0, 'errors': 0, 'unchanged': 0, 'replaced': 0}

def test_schemas(segment_manager_server):
    # initial list of schemas
    result = segment_manager_server.get('/schema')
    assert result.status_code == 200
    assert result.mimetype == 'application/json'
    result_bytes = b''.join(result.response)
    result_list = ujson.loads(result_bytes)
    assert set(result_list) == {'default'}

    # existent schema as json
    result = segment_manager_server.get('/schema/default')
    assert result.status_code == 200
    assert result.mimetype == 'application/json'
    result_bytes = b''.join(result.response)
    result_dict = ujson.loads(result_bytes)
    assert result_dict == {'id': 'default', 'sql': ''}

    # existent schema sql
    result = segment_manager_server.get('/schema/default/sql')
    assert result.status_code == 200
    assert result.mimetype == 'application/sql'
    result_bytes = b''.join(result.response)
    assert result_bytes == b''

    # schema doesn't exist yet
    result = segment_manager_server.get('/schema/schema1')
    assert result.status_code == 404

    # schema doesn't exist yet
    result = segment_manager_server.get('/schema/schema1/sql')
    assert result.status_code == 404

    # bad request: POST not accepted (must be PUT)
    result = segment_manager_server.post('/schema/schema1', data='{}')
    assert result.status_code == 405
    result = segment_manager_server.post('/schema/schema1/sql', data='')
    assert result.status_code == 405

    # bad request: invalid json
    result = segment_manager_server.put(
            '/schema/schema1', data=']]}what the not valid json' )
    assert result.status_code == 400
    assert b''.join(result.response) == b'input could not be parsed as json'

    # bad request: id in json does not match url
    result = segment_manager_server.put(
            '/schema/schema1', data=ujson.dumps({'id': 'schema2', 'sql': ''}))
    assert result.status_code == 400
    assert b''.join(result.response) == b"id in json 'schema2' does not match id in url 'schema1'"

    # bad request: missing sql
    result = segment_manager_server.put(
            '/schema/schema1', data=ujson.dumps({'id': 'schema1'}))
    assert result.status_code == 400
    assert b''.join(result.response) == b"input json has keys {'id'} (should be {'id', 'sql'})"

    # bad request: missing id
    result = segment_manager_server.put(
            '/schema/schema1', data=ujson.dumps({'sql': ''}))
    assert result.status_code == 400
    assert b''.join(result.response) == b"input json has keys {'sql'} (should be {'id', 'sql'})"

    # bad request: invalid sql
    result = segment_manager_server.put(
            '/schema/schema1', data=ujson.dumps({'id': 'schema1', 'sql': 'create create table table blah blooofdjaio'}))
    assert result.status_code == 400
    assert b''.join(result.response) == b'schema sql failed validation: near "create": syntax error'

    # create new schema by submitting sql
    result = segment_manager_server.put(
            '/schema/schema1/sql', content_type='applicaton/sql',
            data='create table foo (bar varchar(100));')
    assert result.status_code == 201

    # get the new schema as json
    result = segment_manager_server.get('/schema/schema1')
    assert result.status_code == 200
    assert result.mimetype == 'application/json'
    result_bytes = b''.join(result.response)
    result_dict = ujson.loads(result_bytes)
    assert result_dict == {'id': 'schema1', 'sql': 'create table foo (bar varchar(100));'}

    # get the new schema as sql
    result = segment_manager_server.get('/schema/schema1/sql')
    assert result.status_code == 200
    assert result.mimetype == 'application/sql'
    result_bytes = b''.join(result.response)
    assert result_bytes == b'create table foo (bar varchar(100));'

    # create new schema by submitting json
    result = segment_manager_server.put(
            '/schema/schema2', content_type='applicaton/sql',
            data=ujson.dumps({'id': 'schema2', 'sql': 'create table schema2_table (foo varchar(100));'}))
    assert result.status_code == 201

    # get the new schema as json
    result = segment_manager_server.get('/schema/schema2')
    assert result.status_code == 200
    assert result.mimetype == 'application/json'
    result_bytes = b''.join(result.response)
    result_dict = ujson.loads(result_bytes)
    assert result_dict == {'id': 'schema2', 'sql': 'create table schema2_table (foo varchar(100));'}

    # get the new schema as sql
    result = segment_manager_server.get('/schema/schema2/sql')
    assert result.status_code == 200
    assert result.mimetype == 'application/sql'
    result_bytes = b''.join(result.response)
    assert result_bytes == b'create table schema2_table (foo varchar(100));'

    # updated list of schemas
    result = segment_manager_server.get('/schema')
    assert result.status_code == 200
    assert result.mimetype == 'application/json'
    result_bytes = b''.join(result.response)
    result_list = ujson.loads(result_bytes)
    assert set(result_list) == {'default', 'schema1', 'schema2'}

    # overwrite schema1 with json api
    result = segment_manager_server.put(
            '/schema/schema1', content_type='applicaton/json',
            data=ujson.dumps({'id': 'schema1', 'sql': 'create table blah (toot varchar(100));'}))
    assert result.status_code == 204

    # get the modified schema as sql
    result = segment_manager_server.get('/schema/schema1/sql')
    assert result.status_code == 200
    assert result.mimetype == 'application/sql'
    result_bytes = b''.join(result.response)
    assert result_bytes == b'create table blah (toot varchar(100));'

    # overwrite schema1 with sql api
    result = segment_manager_server.put(
            '/schema/schema1/sql', content_type='applicaton/sql',
            data='create table haha (hehehe varchar(100));')
    assert result.status_code == 204

    # get the modified schema as json
    result = segment_manager_server.get('/schema/schema1')
    assert result.status_code == 200
    assert result.mimetype == 'application/json'
    result_bytes = b''.join(result.response)
    result_dict = ujson.loads(result_bytes)
    assert result_dict == {'id': 'schema1', 'sql': 'create table haha (hehehe varchar(100));'}

    # updated list of schemas
    result = segment_manager_server.get('/schema')
    assert result.status_code == 200
    assert result.mimetype == 'application/json'
    result_bytes = b''.join(result.response)
    result_list = ujson.loads(result_bytes)
    assert set(result_list) == {'default', 'schema1', 'schema2'}

    # XXX DELETE?

def test_promotion(segment_manager_server):
    result = segment_manager_server.get('/promote')
    assert result.status == '405 METHOD NOT ALLOWED'

    assert False
    # provision a test segment for write
    # result = segment_manager_server.post(
    #         '/provision', content_type='application/json',
    #         data=ujson.dumps({'segment':'test_provision_segment'}))
    # assert result.status_code == 200
    # assert result.mimetype == 'application/json'
    # result_bytes = b''.join(result.response)
    # result_dict = ujson.loads(result_bytes)
    # assert result_dict['write_url'].endswith(':6222/?segment=test_provision_segment')

    # # now it has already been provisioned
    # result = segment_manager_server.post('/promote', data='test_simple_provision_segment')
    # assert result.status_code == 200
    # assert result.mimetype == 'text/plain'
    # assert b''.join(result.response).endswith(b':6222/?segment=test_simple_provision_segment')

<commit_msg>WIP tests for segment promotion upstream<commit_after>import pytest
from trough.wsgi.segment_manager import server
import ujson
from trough.settings import settings
import doublethink
import rethinkdb as r
import requests # :-\ urllib3?

@pytest.fixture(scope="module")
def segment_manager_server():
    server.testing = True
    return server.test_client()

def test_simple_provision(segment_manager_server):
    result = segment_manager_server.get('/')
    assert result.status == '405 METHOD NOT ALLOWED'

    # hasn't been provisioned yet
    result = segment_manager_server.post('/', data='test_simple_provision_segment')
    assert result.status_code == 200
    assert result.mimetype == 'text/plain'
    assert b''.join(result.response).endswith(b':6222/?segment=test_simple_provision_segment')

    # now it has already been provisioned
    result = segment_manager_server.post('/', data='test_simple_provision_segment')
    assert result.status_code == 200
    assert result.mimetype == 'text/plain'
    assert b''.join(result.response).endswith(b':6222/?segment=test_simple_provision_segment')

def test_provision(segment_manager_server):
    result = segment_manager_server.get('/provision')
    assert result.status == '405 METHOD NOT ALLOWED'

    # hasn't been provisioned yet
    result = segment_manager_server.post(
            '/provision', content_type='application/json',
            data=ujson.dumps({'segment':'test_provision_segment'}))
    assert result.status_code == 200
    assert result.mimetype == 'application/json'
    result_bytes = b''.join(result.response)
    result_dict = ujson.loads(result_bytes) # ujson accepts bytes! 😻
    assert result_dict['write_url'].endswith(':6222/?segment=test_provision_segment')

    # now it has already been provisioned
    result = segment_manager_server.post(
            '/provision', content_type='application/json',
            data=ujson.dumps({'segment':'test_provision_segment'}))
    assert result.status_code == 200
    assert result.mimetype == 'application/json'
    result_bytes = b''.join(result.response)
    result_dict = ujson.loads(result_bytes)
    assert result_dict['write_url'].endswith(':6222/?segment=test_provision_segment')

def test_provision_with_schema(segment_manager_server):
    schema = '''CREATE TABLE test (id INTEGER PRIMARY KEY AUTOINCREMENT, test varchar(4));
INSERT INTO test (test) VALUES ("test");'''
    # create a schema by submitting sql
    result = segment_manager_server.put(
            '/schema/test1/sql', content_type='applicaton/sql', data=schema)
    assert result.status_code == 201

    # provision a segment with that schema
    result = segment_manager_server.post(
            '/provision', content_type='application/json',
            data=ujson.dumps({'segment':'test_provision_with_schema_1', 'schema':'test1'}))
    assert result.status_code == 200
    assert result.mimetype == 'application/json'
    result_bytes = b''.join(result.response)
    result_dict = ujson.loads(result_bytes) # ujson accepts bytes! 😻
    assert result_dict['write_url'].endswith(':6222/?segment=test_provision_with_schema_1')

    # get db read url from rethinkdb
    rethinker = doublethink.Rethinker(
            servers=settings['RETHINKDB_HOSTS'], db='trough_configuration')
    query = rethinker.table('services').get_all('test_provision_with_schema_1', index='segment').filter({'role': 'trough-read'}).filter(lambda svc: r.now().sub(svc['last_heartbeat']).lt(svc['ttl'])).order_by('load')[0]
    healthy_segment = query.run()
    read_url = healthy_segment.get('url')
    assert read_url.endswith(':6444/?segment=test_provision_with_schema_1')

    # run a query to check that the schema was used
    sql = 'SELECT * FROM test;'
    with requests.post(read_url, stream=True, data=sql) as response:
        assert response.status_code == 200
        result = ujson.loads(response.text)
        assert result == [{'test': 'test', 'id': 1}]

    # delete the schema from rethinkdb for the sake of other tests
    rethinker = doublethink.Rethinker(
            servers=settings['RETHINKDB_HOSTS'], db='trough_configuration')
    result = rethinker.table('schema').get('test1').delete().run()
    assert result == {'deleted': 1, 'inserted': 0, 'skipped': 0, 'errors': 0, 'unchanged': 0, 'replaced': 0}

def test_schemas(segment_manager_server):
    # initial list of schemas
    result = segment_manager_server.get('/schema')
    assert result.status_code == 200
    assert result.mimetype == 'application/json'
    result_bytes = b''.join(result.response)
    result_list = ujson.loads(result_bytes)
    assert set(result_list) == {'default'}

    # existent schema as json
    result = segment_manager_server.get('/schema/default')
    assert result.status_code == 200
    assert result.mimetype == 'application/json'
    result_bytes = b''.join(result.response)
    result_dict = ujson.loads(result_bytes)
    assert result_dict == {'id': 'default', 'sql': ''}

    # existent schema sql
    result = segment_manager_server.get('/schema/default/sql')
    assert result.status_code == 200
    assert result.mimetype == 'application/sql'
    result_bytes = b''.join(result.response)
    assert result_bytes == b''

    # schema doesn't exist yet
    result = segment_manager_server.get('/schema/schema1')
    assert result.status_code == 404

    # schema doesn't exist yet
    result = segment_manager_server.get('/schema/schema1/sql')
    assert result.status_code == 404

    # bad request: POST not accepted (must be PUT)
    result = segment_manager_server.post('/schema/schema1', data='{}')
    assert result.status_code == 405
    result = segment_manager_server.post('/schema/schema1/sql', data='')
    assert result.status_code == 405

    # bad request: invalid json
    result = segment_manager_server.put(
            '/schema/schema1', data=']]}what the not valid json' )
    assert result.status_code == 400
    assert b''.join(result.response) == b'input could not be parsed as json'

    # bad request: id in json does not match url
    result = segment_manager_server.put(
            '/schema/schema1', data=ujson.dumps({'id': 'schema2', 'sql': ''}))
    assert result.status_code == 400
    assert b''.join(result.response) == b"id in json 'schema2' does not match id in url 'schema1'"

    # bad request: missing sql
    result = segment_manager_server.put(
            '/schema/schema1', data=ujson.dumps({'id': 'schema1'}))
    assert result.status_code == 400
    assert b''.join(result.response) == b"input json has keys {'id'} (should be {'id', 'sql'})"

    # bad request: missing id
    result = segment_manager_server.put(
            '/schema/schema1', data=ujson.dumps({'sql': ''}))
    assert result.status_code == 400
    assert b''.join(result.response) == b"input json has keys {'sql'} (should be {'id', 'sql'})"

    # bad request: invalid sql
    result = segment_manager_server.put(
            '/schema/schema1', data=ujson.dumps({'id': 'schema1', 'sql': 'create create table table blah blooofdjaio'}))
    assert result.status_code == 400
    assert b''.join(result.response) == b'schema sql failed validation: near "create": syntax error'

    # create new schema by submitting sql
    result = segment_manager_server.put(
            '/schema/schema1/sql', content_type='applicaton/sql',
            data='create table foo (bar varchar(100));')
    assert result.status_code == 201

    # get the new schema as json
    result = segment_manager_server.get('/schema/schema1')
    assert result.status_code == 200
    assert result.mimetype == 'application/json'
    result_bytes = b''.join(result.response)
    result_dict = ujson.loads(result_bytes)
    assert result_dict == {'id': 'schema1', 'sql': 'create table foo (bar varchar(100));'}

    # get the new schema as sql
    result = segment_manager_server.get('/schema/schema1/sql')
    assert result.status_code == 200
    assert result.mimetype == 'application/sql'
    result_bytes = b''.join(result.response)
    assert result_bytes == b'create table foo (bar varchar(100));'

    # create new schema by submitting json
    result = segment_manager_server.put(
            '/schema/schema2', content_type='applicaton/sql',
            data=ujson.dumps({'id': 'schema2', 'sql': 'create table schema2_table (foo varchar(100));'}))
    assert result.status_code == 201

    # get the new schema as json
    result = segment_manager_server.get('/schema/schema2')
    assert result.status_code == 200
    assert result.mimetype == 'application/json'
    result_bytes = b''.join(result.response)
    result_dict = ujson.loads(result_bytes)
    assert result_dict == {'id': 'schema2', 'sql': 'create table schema2_table (foo varchar(100));'}

    # get the new schema as sql
    result = segment_manager_server.get('/schema/schema2/sql')
    assert result.status_code == 200
    assert result.mimetype == 'application/sql'
    result_bytes = b''.join(result.response)
    assert result_bytes == b'create table schema2_table (foo varchar(100));'

    # updated list of schemas
    result = segment_manager_server.get('/schema')
    assert result.status_code == 200
    assert result.mimetype == 'application/json'
    result_bytes = b''.join(result.response)
    result_list = ujson.loads(result_bytes)
    assert set(result_list) == {'default', 'schema1', 'schema2'}

    # overwrite schema1 with json api
    result = segment_manager_server.put(
            '/schema/schema1', content_type='applicaton/json',
            data=ujson.dumps({'id': 'schema1', 'sql': 'create table blah (toot varchar(100));'}))
    assert result.status_code == 204

    # get the modified schema as sql
    result = segment_manager_server.get('/schema/schema1/sql')
    assert result.status_code == 200
    assert result.mimetype == 'application/sql'
    result_bytes = b''.join(result.response)
    assert result_bytes == b'create table blah (toot varchar(100));'

    # overwrite schema1 with sql api
    result = segment_manager_server.put(
            '/schema/schema1/sql', content_type='applicaton/sql',
            data='create table haha (hehehe varchar(100));')
    assert result.status_code == 204

    # get the modified schema as json
    result = segment_manager_server.get('/schema/schema1')
    assert result.status_code == 200
    assert result.mimetype == 'application/json'
    result_bytes = b''.join(result.response)
    result_dict = ujson.loads(result_bytes)
    assert result_dict == {'id': 'schema1', 'sql': 'create table haha (hehehe varchar(100));'}

    # updated list of schemas
    result = segment_manager_server.get('/schema')
    assert result.status_code == 200
    assert result.mimetype == 'application/json'
    result_bytes = b''.join(result.response)
    result_list = ujson.loads(result_bytes)
    assert set(result_list) == {'default', 'schema1', 'schema2'}

    # XXX DELETE?

def test_promotion(segment_manager_server):
    from snakebite.client import Client

    snakebite_client = Client(settings['HDFS_HOST'], settings['HDFS_PORT'])

    result = segment_manager_server.get('/promote')
    assert result.status == '405 METHOD NOT ALLOWED'

    # provision a test segment for write
    result = segment_manager_server.post(
            '/provision', content_type='application/json',
            data=ujson.dumps({'segment':'test_promote_segment'}))
    assert result.status_code == 200
    assert result.mimetype == 'application/json'
    result_bytes = b''.join(result.response)
    result_dict = ujson.loads(result_bytes)
    assert result_dict['write_url'].endswith(':6222/?segment=test_promote_segment')
    write_url = result_dict['write_url']
    byte_count = result_dict['size']

    listing = snakebite_client.ls([settings['HDFS_PATH']])
    listing_before_promotion = {item['path']: item['length'] for item in listing}
    assert 'test_promote_segment.sqlite' not in listing_before_promotion

    # now write to the segment and promote it to HDFS
    result = segment_manager_server.post(
            '/promote', content_type='application/json',
            data=ujson.dumps({'segment': 'test_simple_promote_segment'}))
    assert result.status_code == 200
    assert result.mimetype == 'text/plain'
    assert b''.join(result.response).endswith(b':6222/?segment=test_simple_promote_segment')

    listing_after_promotion = {item['path']: item['length'] for item in snakebite_client.ls([settings['HDFS_PATH']])}
    assert 'test_promote_segment.sqlite' in listing_after_promotion
    assert byte_count == listing_after_promotion.get(item['filename'])
<|endoftext|>


<commit_before>from __future__ import print_function
from functools import reduce
import operator

class Prop:
    def __init__(self, *args):
        """Proposition constructor.
        Prop(True), Prop(False) constructs the True and False atoms.
        Prop(name) constructs an atomic proposition.
        Prop('&', p1, p2, ...) constructs the conjunction of propositions p1, p2, ...
        Prop('|', p1, p2, ...) constructs the disjunction of propositions p1, p2, ...
        Prop('~', p) constructs the negation of proposition p.

        Built-in rewritings:
            redundance-elimination-union: Prop1 | Prop2 -> (|Prop1) if Prop1 == Prop2
            redundance-elimination-intersection: Prop1 & Prop2 -> (&Prop1) if Prop1 == Prop2"""
        if len(args) == 0:
            raise NameError("Malformed proposition.")

        if len(args) == 1:
            self.name = args[0]
            self.atomic = True

        if len(args) > 1:
            if args[0] not in ['&', '|', '~'] or \
                args[0] == '~' and len(args)>2:
                raise NameError("Malformed proposition: "+str(args))
            self.oper = args[0]
            self.terms = frozenset(args[1:])
            self.atomic = False

    # My function to handle return the names of True and False propositions as the
    # lowercase strings 'true' and 'false'
    def get_name(self):
        if self.atomic:
            """
            if self.name == True:
                return 'true'
            elif self.name == False:
                return 'false'
            """
            return self.name
        else:
            return None

    def is_true(self):
        return self.atomic and self.name == True

    # My function to get all the atom names.
    def get_atoms(self):
        if self.atomic:
            return [self.get_name()]
        return reduce(operator.add, [term.get_atoms() for term in self.terms])

    def is_atomic(self):
        return self.atomic

    def is_positive(self):
        return self.atomic or (self.oper in ['&', '|'])

    def is_negative(self):
        return (not self.atomic) and (self.oper == '~')

    def is_literal(self):
        return self.atomic or \
                self.oper == '~' and [x.atomic for x in self.terms][0]

    def __eq__(self, other):
        """Syntactic equality testing. Fully commutative but not associative.
        If True, then the two propositions are equal. If False, unsure."""
        if self.atomic:
            if other.atomic:
                return self.name == other.name
            return False
        if other.atomic:
            return False
        if self.oper == self.oper:
            return self.terms == other.terms
        return False

    def __ne__(self, other):
        return not (self == other)

    def __hash__(self):
        if '_hash' in self.__dict__:
            return self._hash
        if self.atomic:
            self._hash = hash(self.name)
            return self._hash
        self._hash = hash(hash(self.oper)^hash(self.terms))
        return self._hash

    def __repr__(self):
        def no_parenthesing(s):
            return s.atomic or s.oper == '~'

        if self.atomic:
            return str(self.get_name())

        oper = self.oper
        if oper == '~':
            a = [x for x in self.terms][0]
            if no_parenthesing(a):
                return oper+repr(a)
            return oper+'('+repr(a)+')'

        if oper in ['&', '|']:
            subs = []
            for e in self.terms:
                if no_parenthesing(e):
                    subs.append(repr(e))
                else:
                    subs.append('('+repr(e)+')')
            assert(len(subs) > 0)
            if len(subs) == 1:
                return '('+oper+subs[0]+')'
            return (' '+oper+' ').join(subs)

        assert(False)

    def __and__(self, other):
        return Prop('&', self, other)

    def __rand__(self, other):
        return other.__and__(self)

    def __or__(self, other):
        return Prop('|', self, other)

    def __ror__(self, other):
        return other.__ror__(self)

    def __invert__(self):
        return Prop('~', self)

    def __gt__(self, other):
        """material implication."""
        return other | (~self)

    def __sub__(self, other):
        return self & ~other

    def __rsub__(self, other):
        return other.__sub__(self)

    def __pformat(self, indentation=0, space=' '):
        if self.is_literal():
            if self.is_positive():
                return str(self.name)
            return '~'+str(self.get_terms()[0].name)

        oper = self.oper
        if oper == '~':
            a = [x for x in self.terms][0]
            return oper+' '+a.__pformat(indentation+2)

        if oper in ['&', '|']:
            subs = []
            for e in self.terms:
                subs.append(e.__pformat(indentation+2))
            return oper+' '+('\n'+space*(indentation+2)).join(subs)

        assert(False)

    def __str__(self):
        return self.__pformat()

    def get_op(self):
        return self.__dict__.get('oper', None)

    def get_terms(self):
        return [x for x in self.__dict__.get('terms', [None])]

    def size(self):
        if '_size' in self.__dict__:
            return self._size
        if self.atomic:
            self._size = 1
            return 1
        self._size = 1 + sum([x.size() for x in self.terms])
        return self._size

    def depth(self):
        if '_depth' in self.__dict__:
            return self._depth
        if self.is_literal():
            self._depth = 1
            return 1
        self._depth = 1 + max([x.depth() for x in self.terms])
        return self._depth

    def var_stats(self):
        #if '_var_stats' in self.__dict__:
        #    return self._var_stats

        counts = dict()
        #self._var_stats = counts

        if self.is_literal():
            if self.is_positive():
                if self not in counts:
                    counts[self] = [0, 0]
                counts[self][0] += 1
            else:
                pos = self.get_terms()[0]
                if pos not in counts:
                    counts[pos] = [0, 0]
                counts[pos][1] += 1
            return counts

        if self.oper == '~':
            for sub in self.terms:
                nc = sub.var_stats()
                for e in nc:
                    if e not in counts:
                        counts[e] = [0, 0]
                    counts[e][0] += nc[e][1]
                    counts[e][1] += nc[e][0]
            return counts

        def update(a, counts):
            for e in a:
                if e not in counts:
                    counts[e] = a[e]
                else:
                    counts[e][0] += a[e][0]
                    counts[e][1] += a[e][1]

        if self.oper in ['&', '|']:
            for sub in self.terms:
                update(sub.var_stats(), counts)
            return counts

        raise NameError("ERROR")

<commit_msg>Small changes to proposition.py<commit_after>from __future__ import print_function
from functools import reduce
import operator

class Prop:
    def __init__(self, *args):
        """Proposition constructor.
        Prop(True), Prop(False) constructs the True and False atoms.
        Prop(name) constructs an atomic proposition.
        Prop('&', p1, p2, ...) constructs the conjunction of propositions p1, p2, ...
        Prop('|', p1, p2, ...) constructs the disjunction of propositions p1, p2, ...
        Prop('~', p) constructs the negation of proposition p.

        Built-in rewritings:
            redundance-elimination-union: Prop1 | Prop2 -> (|Prop1) if Prop1 == Prop2
            redundance-elimination-intersection: Prop1 & Prop2 -> (&Prop1) if Prop1 == Prop2"""
        if len(args) == 0:
            raise NameError("Malformed proposition.")

        if len(args) == 1:
            self.name = args[0]
            self.atomic = True

        if len(args) > 1:
            if args[0] not in ['&', '|', '~'] or \
                args[0] == '~' and len(args)>2:
                raise NameError("Malformed proposition: "+str(args))
            self.oper = args[0]
            self.terms = frozenset(args[1:])
            self.atomic = False

    def get_name(self):
        if self.atomic:
            return self.name
        else:
            return None

    def is_true(self):
        return self.atomic and self.name == True

    # My function to get all the atom names.
    def get_atoms(self):
        if self.atomic:
            return [self.get_name()]
        return reduce(operator.add, [term.get_atoms() for term in self.terms])

    def is_atomic(self):
        return self.atomic

    def is_positive(self):
        return self.atomic or (self.oper in ['&', '|'])

    def is_negative(self):
        return (not self.atomic) and (self.oper == '~')

    def is_literal(self):
        return self.atomic or \
                self.oper == '~' and [x.atomic for x in self.terms][0]

    def __eq__(self, other):
        """Syntactic equality testing. Fully commutative but not associative.
        If True, then the two propositions are equal. If False, unsure."""
        if self.atomic:
            if other.atomic:
                return self.name == other.name
            return False
        if other.atomic:
            return False
        if self.oper == self.oper:
            return self.terms == other.terms
        return False

    def __ne__(self, other):
        return not (self == other)

    def __hash__(self):
        if '_hash' in self.__dict__:
            return self._hash
        if self.atomic:
            self._hash = hash(self.name)
            return self._hash
        self._hash = hash(hash(self.oper)^hash(self.terms))
        return self._hash

    def __repr__(self):
        def no_parenthesing(s):
            return s.atomic or s.oper == '~'

        if self.atomic:
            return str(self.get_name())

        oper = self.oper
        if oper == '~':
            a = [x for x in self.terms][0]
            if no_parenthesing(a):
                return oper+repr(a)
            return oper+'('+repr(a)+')'

        if oper in ['&', '|']:
            subs = []
            for e in self.terms:
                if no_parenthesing(e):
                    subs.append(repr(e))
                else:
                    subs.append('('+repr(e)+')')
            assert(len(subs) > 0)
            if len(subs) == 1:
                return '('+oper+subs[0]+')'
            return (' '+oper+' ').join(subs)

        assert(False)

    def __and__(self, other):
        return Prop('&', self, other)

    def __rand__(self, other):
        return other.__and__(self)

    def __or__(self, other):
        return Prop('|', self, other)

    def __ror__(self, other):
        return other.__ror__(self)

    def __invert__(self):
        return Prop('~', self)

    def __gt__(self, other):
        """material implication."""
        return other | (~self)

    def __sub__(self, other):
        return self & ~other

    def __rsub__(self, other):
        return other.__sub__(self)

    def __pformat(self, indentation=0, space=' '):
        if self.is_literal():
            if self.is_positive():
                return str(self.name)
            return '~'+str(self.get_terms()[0].name)

        oper = self.oper
        if oper == '~':
            a = [x for x in self.terms][0]
            return oper+' '+a.__pformat(indentation+2)

        if oper in ['&', '|']:
            subs = []
            for e in self.terms:
                subs.append(e.__pformat(indentation+2))
            return oper+' '+('\n'+space*(indentation+2)).join(subs)

        assert(False)

    def __str__(self):
        return self.__pformat()

    def get_op(self):
        return self.__dict__.get('oper', None)

    def get_terms(self):
        return [x for x in self.__dict__.get('terms', [None])]

    def size(self):
        if '_size' in self.__dict__:
            return self._size
        if self.atomic:
            self._size = 1
            return 1
        self._size = 1 + sum([x.size() for x in self.terms])
        return self._size

    def depth(self):
        if '_depth' in self.__dict__:
            return self._depth
        if self.is_literal():
            self._depth = 1
            return 1
        self._depth = 1 + max([x.depth() for x in self.terms])
        return self._depth

    def var_stats(self):
        #if '_var_stats' in self.__dict__:
        #    return self._var_stats

        counts = dict()
        #self._var_stats = counts

        if self.is_literal():
            if self.is_positive():
                if self not in counts:
                    counts[self] = [0, 0]
                counts[self][0] += 1
            else:
                pos = self.get_terms()[0]
                if pos not in counts:
                    counts[pos] = [0, 0]
                counts[pos][1] += 1
            return counts

        if self.oper == '~':
            for sub in self.terms:
                nc = sub.var_stats()
                for e in nc:
                    if e not in counts:
                        counts[e] = [0, 0]
                    counts[e][0] += nc[e][1]
                    counts[e][1] += nc[e][0]
            return counts

        def update(a, counts):
            for e in a:
                if e not in counts:
                    counts[e] = a[e]
                else:
                    counts[e][0] += a[e][0]
                    counts[e][1] += a[e][1]

        if self.oper in ['&', '|']:
            for sub in self.terms:
                update(sub.var_stats(), counts)
            return counts

        raise NameError("ERROR")

<|endoftext|>


<commit_before>import os, sys
import time
import getpass
import logging

from process import UPID, Process, async

from mesos_pb2 import TASK_LOST
from messages_pb2 import (RegisterFrameworkMessage, ReregisterFrameworkMessage,
        DeactivateFrameworkMessage, UnregisterFrameworkMessage,
        ResourceRequestMessage, ReviveOffersMessage, LaunchTasksMessage, KillTaskMessage,
        StatusUpdate, StatusUpdateAcknowledgementMessage, FrameworkToExecutorMessage)

logger = logging.getLogger(__name__)

class Scheduler(object):
    def registered(self, driver, framework_id, masterInfo): pass
    def reregistered(self, driver, masterInfo): pass
    def disconnected(self, driver): pass
    def frameworkMessage(self, driver, slave_id, executor_id, message): pass
    def resourceOffers(self, driver, offers): pass
    def offerRescinded(self, driver, offer_id): pass
    def statusUpdate(self, driver, status): pass
    def executorLost(self, driver, executor_id, slave_id, status): pass
    def slaveLost(self, driver, slave_id): pass
    def error(self, driver, message): pass

class SchedulerDriver(object):
    def start(self): pass
    def join(self): pass
    def run(self): pass
    def abort(self): pass
    def stop(self, failover=False): pass
    def reviveOffers(self): pass
    def requestResources(self, requests): pass
    def declineOffer(self, offerId, filters=None): pass
    def launchTasks(self, offerId, tasks, filters=None): pass
    def killTask(self, taskId): pass
    def sendFrameworkMessage(self, executorId, slaveId, data): pass

class MesosSchedulerDriver(Process):
    def __init__(self, sched, framework, master_uri):
        Process.__init__(self, 'scheduler')
        self.sched = sched
        #self.executor_info = executor_info
        self.master_uri = master_uri
        self.framework = framework
        self.framework.failover_timeout = 100
        self.framework_id = framework.id

        self.master = None
        self.detector = None

        self.connected = False
        self.savedOffers = {}
        self.savedSlavePids = {}

    @async # called by detector
    def onNewMasterDetectedMessage(self, pid):
        self.master = UPID(pid)
        self.register()

    @async # called by detector
    def onNoMasterDetectedMessage(self):
        self.connected = False
        self.master = None

    def register(self):
        if self.connected or self.aborted:
            return

        if self.master:
            if not self.framework_id.value:
                msg = RegisterFrameworkMessage()
                msg.framework.MergeFrom(self.framework)
            else:
                msg = ReregisterFrameworkMessage()
                msg.framework.MergeFrom(self.framework)
                msg.failover = True
            self.send(self.master, msg)

        self.delay(2, self.register)

    def onFrameworkRegisteredMessage(self, framework_id, master_info):
        self.framework_id = framework_id
        self.framework.id.MergeFrom(framework_id)
        self.connected = True
        self.link(self.master, self.onDisconnected)
        self.sched.registered(self, framework_id, master_info)

    def onFrameworkReregisteredMessage(self, framework_id, master_info):
        assert self.framework_id == framework_id
        self.connected = True
        self.link(self.master, self.onDisconnected)
        self.sched.reregistered(self, master_info)

    def onDisconnected(self):
        self.connected = False
        logger.warning("disconnected from master")
        self.delay(5, self.register)

    def onResourceOffersMessage(self, offers, pids):
        for offer, pid in zip(offers, pids):
            self.savedOffers.setdefault(offer.id.value, {})[offer.slave_id.value] = UPID(pid)
        self.sched.resourceOffers(self, list(offers))

    def onRescindResourceOfferMessage(self, offer_id):
        self.savedOffers.pop(offer_id.value, None)
        self.sched.offerRescinded(self, offer_id)

    def onStatusUpdateMessage(self, update, pid=''):
        assert self.framework_id == update.framework_id
        self.sched.statusUpdate(self, update.status)
        if not self.aborted and pid:
            reply = StatusUpdateAcknowledgementMessage()
            reply.framework_id.MergeFrom(self.framework_id)
            reply.slave_id.MergeFrom(update.slave_id)
            reply.task_id.MergeFrom(update.status.task_id)
            reply.uuid = update.uuid
            self.send(UPID(pid), reply)

    def onLostSlaveMessage(self, slave_id):
        self.sched.slaveLost(self, slave_id)

    def onExecutorToFrameworkMessage(self, slave_id, framework_id, executor_id, data):
        self.sched.frameworkMessage(self, slave_id, executor_id, data)

    def onFrameworkErrorMessage(self, message, code=0):
        self.sched.error(self, code, message)

    def start(self):
        Process.start(self)
        uri = self.master_uri
        if uri.startswith('zk://') or uri.startswith('zoo://'):
            from .detector import MasterDetector
            self.detector = MasterDetector(uri[uri.index('://') + 3:], self)
            self.detector.start()
        else:
            if not ':' in uri:
                uri += ':5050'
            self.onNewMasterDetectedMessage('master@%s' % uri)

    def abort(self):
        if self.connected:
            msg = UnregisterFrameworkMessage()
            msg.framework_id.MergeFrom(self.framework_id)
            self.send(self.master, msg)
        Process.abort(self)

    def stop(self, failover=False):
        if self.connected and not failover:
            msg = DeactivateFrameworkMessage()
            msg.framework_id.MergeFrom(self.framework_id)
            self.send(self.master, msg)
        if self.detector:
            self.detector.stop()
        Process.stop(self)

    @async
    def requestResources(self, requests):
        msg = ResourceRequestMessage()
        msg.framework_id.MergeFrom(self.framework_id)
        for req in requests:
            msg.requests.add().MergeFrom(req)
        self.send(self.master, msg)

    @async
    def reviveOffers(self):
        msg = ReviveOffersMessage()
        msg.framework_id.MergeFrom(self.framework_id)
        self.send(self.master, msg)

    def launchTasks(self, offer_id, tasks, filters):
        if not self.connected or offer_id.value not in self.savedOffers:
            update = StatusUpdate()
            update.framework_id.MergeFrom(self.framework_id)
            update.status.task_id.MergeFrom(task.task_id)
            update.status.state = TASK_LOST
            update.status.message = 'Master disconnected' if not self.connected else "invalid offer_id"
            update.timestamp = time.time()
            update.uuid = ''
            return self.onStatusUpdateMessage(update)

        msg = LaunchTasksMessage()
        msg.framework_id.MergeFrom(self.framework_id)
        msg.offer_id.MergeFrom(offer_id)
        msg.filters.MergeFrom(filters)
        for task in tasks:
            msg.tasks.add().MergeFrom(task)
            pid = self.savedOffers.get(offer_id.value, {}).get(task.slave_id.value)
            if pid and task.slave_id.value not in self.savedSlavePids:
                self.savedSlavePids[task.slave_id.value] = pid
        self.savedOffers.pop(offer_id.value)
        self.send(self.master, msg)

    def declineOffer(self, offer_id, filters=None):
        if not self.connected:
            return
        msg = LaunchTasksMessage()
        msg.framework_id.MergeFrom(self.framework_id)
        msg.offer_id.MergeFrom(offer_id)
        if filters:
             msg.filters.MergeFrom(filters)
        self.send(self.master, msg)

    @async
    def killTask(self, task_id):
        if not self.connected:
            return
        msg = KillTaskMessage()
        msg.framework_id.MergeFrom(self.framework_id)
        msg.task_id.MergeFrom(task_id)
        self.send(self.master, msg)

    @async
    def sendFrameworkMessage(self, executor_id, slave_id, data):
        if not self.connected:
            return

        msg = FrameworkToExecutorMessage()
        msg.framework_id.MergeFrom(self.framework_id)
        msg.executor_id.MergeFrom(executor_id)
        msg.slave_id.MergeFrom(slave_id)
        msg.data = data

        slave = self.savedSlavePids.get(slave_id.value, self.master) # can not send to slave directly
        self.send(slave, msg)
<commit_msg>ignore invalid pid<commit_after>import os, sys
import time
import getpass
import logging

from process import UPID, Process, async

from mesos_pb2 import TASK_LOST
from messages_pb2 import (RegisterFrameworkMessage, ReregisterFrameworkMessage,
        DeactivateFrameworkMessage, UnregisterFrameworkMessage,
        ResourceRequestMessage, ReviveOffersMessage, LaunchTasksMessage, KillTaskMessage,
        StatusUpdate, StatusUpdateAcknowledgementMessage, FrameworkToExecutorMessage)

logger = logging.getLogger(__name__)

class Scheduler(object):
    def registered(self, driver, framework_id, masterInfo): pass
    def reregistered(self, driver, masterInfo): pass
    def disconnected(self, driver): pass
    def frameworkMessage(self, driver, slave_id, executor_id, message): pass
    def resourceOffers(self, driver, offers): pass
    def offerRescinded(self, driver, offer_id): pass
    def statusUpdate(self, driver, status): pass
    def executorLost(self, driver, executor_id, slave_id, status): pass
    def slaveLost(self, driver, slave_id): pass
    def error(self, driver, message): pass

class SchedulerDriver(object):
    def start(self): pass
    def join(self): pass
    def run(self): pass
    def abort(self): pass
    def stop(self, failover=False): pass
    def reviveOffers(self): pass
    def requestResources(self, requests): pass
    def declineOffer(self, offerId, filters=None): pass
    def launchTasks(self, offerId, tasks, filters=None): pass
    def killTask(self, taskId): pass
    def sendFrameworkMessage(self, executorId, slaveId, data): pass

class MesosSchedulerDriver(Process):
    def __init__(self, sched, framework, master_uri):
        Process.__init__(self, 'scheduler')
        self.sched = sched
        #self.executor_info = executor_info
        self.master_uri = master_uri
        self.framework = framework
        self.framework.failover_timeout = 100
        self.framework_id = framework.id

        self.master = None
        self.detector = None

        self.connected = False
        self.savedOffers = {}
        self.savedSlavePids = {}

    @async # called by detector
    def onNewMasterDetectedMessage(self, pid):
        self.master = UPID(pid)
        self.register()

    @async # called by detector
    def onNoMasterDetectedMessage(self):
        self.connected = False
        self.master = None

    def register(self):
        if self.connected or self.aborted:
            return

        if self.master:
            if not self.framework_id.value:
                msg = RegisterFrameworkMessage()
                msg.framework.MergeFrom(self.framework)
            else:
                msg = ReregisterFrameworkMessage()
                msg.framework.MergeFrom(self.framework)
                msg.failover = True
            self.send(self.master, msg)

        self.delay(2, self.register)

    def onFrameworkRegisteredMessage(self, framework_id, master_info):
        self.framework_id = framework_id
        self.framework.id.MergeFrom(framework_id)
        self.connected = True
        self.link(self.master, self.onDisconnected)
        self.sched.registered(self, framework_id, master_info)

    def onFrameworkReregisteredMessage(self, framework_id, master_info):
        assert self.framework_id == framework_id
        self.connected = True
        self.link(self.master, self.onDisconnected)
        self.sched.reregistered(self, master_info)

    def onDisconnected(self):
        self.connected = False
        logger.warning("disconnected from master")
        self.delay(5, self.register)

    def onResourceOffersMessage(self, offers, pids):
        for offer, pid in zip(offers, pids):
            self.savedOffers.setdefault(offer.id.value, {})[offer.slave_id.value] = UPID(pid)
        self.sched.resourceOffers(self, list(offers))

    def onRescindResourceOfferMessage(self, offer_id):
        self.savedOffers.pop(offer_id.value, None)
        self.sched.offerRescinded(self, offer_id)

    def onStatusUpdateMessage(self, update, pid=''):
        assert self.framework_id == update.framework_id
        self.sched.statusUpdate(self, update.status)
        if not self.aborted and pid and not pid.endswith('0.0.0.0'):
            reply = StatusUpdateAcknowledgementMessage()
            reply.framework_id.MergeFrom(self.framework_id)
            reply.slave_id.MergeFrom(update.slave_id)
            reply.task_id.MergeFrom(update.status.task_id)
            reply.uuid = update.uuid
            self.send(UPID(pid), reply)

    def onLostSlaveMessage(self, slave_id):
        self.sched.slaveLost(self, slave_id)

    def onExecutorToFrameworkMessage(self, slave_id, framework_id, executor_id, data):
        self.sched.frameworkMessage(self, slave_id, executor_id, data)

    def onFrameworkErrorMessage(self, message, code=0):
        self.sched.error(self, code, message)

    def start(self):
        Process.start(self)
        uri = self.master_uri
        if uri.startswith('zk://') or uri.startswith('zoo://'):
            from .detector import MasterDetector
            self.detector = MasterDetector(uri[uri.index('://') + 3:], self)
            self.detector.start()
        else:
            if not ':' in uri:
                uri += ':5050'
            self.onNewMasterDetectedMessage('master@%s' % uri)

    def abort(self):
        if self.connected:
            msg = UnregisterFrameworkMessage()
            msg.framework_id.MergeFrom(self.framework_id)
            self.send(self.master, msg)
        Process.abort(self)

    def stop(self, failover=False):
        if self.connected and not failover:
            msg = DeactivateFrameworkMessage()
            msg.framework_id.MergeFrom(self.framework_id)
            self.send(self.master, msg)
        if self.detector:
            self.detector.stop()
        Process.stop(self)

    @async
    def requestResources(self, requests):
        msg = ResourceRequestMessage()
        msg.framework_id.MergeFrom(self.framework_id)
        for req in requests:
            msg.requests.add().MergeFrom(req)
        self.send(self.master, msg)

    @async
    def reviveOffers(self):
        msg = ReviveOffersMessage()
        msg.framework_id.MergeFrom(self.framework_id)
        self.send(self.master, msg)

    def launchTasks(self, offer_id, tasks, filters):
        if not self.connected or offer_id.value not in self.savedOffers:
            update = StatusUpdate()
            update.framework_id.MergeFrom(self.framework_id)
            update.status.task_id.MergeFrom(task.task_id)
            update.status.state = TASK_LOST
            update.status.message = 'Master disconnected' if not self.connected else "invalid offer_id"
            update.timestamp = time.time()
            update.uuid = ''
            return self.onStatusUpdateMessage(update)

        msg = LaunchTasksMessage()
        msg.framework_id.MergeFrom(self.framework_id)
        msg.offer_id.MergeFrom(offer_id)
        msg.filters.MergeFrom(filters)
        for task in tasks:
            msg.tasks.add().MergeFrom(task)
            pid = self.savedOffers.get(offer_id.value, {}).get(task.slave_id.value)
            if pid and task.slave_id.value not in self.savedSlavePids:
                self.savedSlavePids[task.slave_id.value] = pid
        self.savedOffers.pop(offer_id.value)
        self.send(self.master, msg)

    def declineOffer(self, offer_id, filters=None):
        if not self.connected:
            return
        msg = LaunchTasksMessage()
        msg.framework_id.MergeFrom(self.framework_id)
        msg.offer_id.MergeFrom(offer_id)
        if filters:
             msg.filters.MergeFrom(filters)
        self.send(self.master, msg)

    @async
    def killTask(self, task_id):
        if not self.connected:
            return
        msg = KillTaskMessage()
        msg.framework_id.MergeFrom(self.framework_id)
        msg.task_id.MergeFrom(task_id)
        self.send(self.master, msg)

    @async
    def sendFrameworkMessage(self, executor_id, slave_id, data):
        if not self.connected:
            return

        msg = FrameworkToExecutorMessage()
        msg.framework_id.MergeFrom(self.framework_id)
        msg.executor_id.MergeFrom(executor_id)
        msg.slave_id.MergeFrom(slave_id)
        msg.data = data

        slave = self.savedSlavePids.get(slave_id.value, self.master) # can not send to slave directly
        self.send(slave, msg)
<|endoftext|>


<commit_before># -*- coding: utf-8 -*-
# Copyright (c) 2015, Universal Resource Trading Limited and contributors
# For license information, please see license.txt

import os
import re
import shutil
import subprocess
from datetime import date

import json
import frappe

uploads_path = os.path.join(os.sep, 'home', 'uploads')

re_digitsearch = re.compile('([0-9]+)')

def realtime_eval(rte_id, tag, event, msg):
    """
    Use publish_realtime to run javascript code using custom event handler.

    rte_id: persistent id for the current form
    tag: communications tag specific to one series of communications
    event: name of custom event to trigger on client.
    msg: JSON serializable object (e.g. dict) to push to client.
    """

    rte_msg = json.dumps(msg)
    js = "erpnext_ebay_realtime_event('{}', '{}', '{}', '{}');"
    js = js.format(rte_id, tag, event, rte_msg)
    frappe.publish_realtime(event='eval_js', message=js,
                            user=frappe.session.user)

def resize_image(filename, out=None, thumbnail=False):
    """
    Use convert or mogrify to resize images
    Uses one of the below commands:

    mogrify -auto-orient -resize LxL> filename
    convert -auto-orient -resize LxL> filename out

    Resizes, rotates according to EXIF information, and resizes to fit into
    an L x L size box (trailing '>' symbol means that images will never
    be enlarged). The image size is taken from the ebay_image_size or
    ebay_thumbnail_size setting in Ebay_Settings_Manager.

    If out is set, the resized version is output to a new file (convert).
    Otherwise the modification takes place on the original file (mogrify).
    """

    if thumbnail:
        image_size = frappe.db.get_value(
            'eBay Manager Settings', filters=None,
            fieldname='ebay_thumbnail_size')
    else:
        image_size = frappe.db.get_value(
            'eBay Manager Settings', filters=None,
            fieldname='ebay_image_size')
    image_size = int(image_size)

    if image_size < 1:
        frappe.throw('Invalid image size: ' + str(image_size))

    size_string = '{}x{}>'.format(image_size, image_size)
    if out is not None:
        subprocess.call(
            ['convert', '-auto-orient', '-resize', size_string,
             filename, out])
    else:
        subprocess.call(
            ['mogrify', '-auto-orient', '-resize', size_string, filename])

def ugs_save_file_on_filesystem_hook(*args, **kwargs):
    """Intercept all write_file events, and mogrify images

    Replaces the standard write_file event. Obtains the filename, content_type
    etc. Calls the normal backup 'save_file_on_filesystem' with all arguments.
    If we do not handle the attachment specially, this function is entirely
    transparent. However, we can handle specific file types as we see fit - in
    this case we mogrify JPG and JPEG images. save_file_on_filesystem does
    strange things, so we need to reconstruct the filename using analagous
    logic - this could break in future with Frappe changes."""

    ret = frappe.utils.file_manager.save_file_on_filesystem(*args, **kwargs)

    file_name = ret['file_name']
    #file_url = ret['file_url'] # Not a consistent file system identifier

    if ('is_private' in kwargs) and kwargs['is_private']:
        file_path = os.path.abspath(
            frappe.get_site_path('private', 'files', file_name))
    else:
        file_path = os.path.abspath(
            frappe.get_site_path('public', 'files', file_name))

    extension = os.path.splitext(file_name)[1].lower()

    if extension in ('.jpg', '.jpeg'):
        # Resize and autoorient this image
        resize_image(file_path)

    return ret

@frappe.whitelist()
def view_slideshow_py(slideshow):

    # Whitelisted function; check permissions
    if not frappe.has_permission('Item', 'read'):
        frappe.throw('Need read permissions on Item!',
                     frappe.PermissionError)

    images_path = os.path.abspath(frappe.get_site_path('public'))

    image_list = frappe.db.sql("""
        SELECT image
        FROM `tabWebsite Slideshow Item`
        WHERE parent=%s
        ORDER BY idx;
        """, (slideshow,), as_dict=False)
    image_list = [x[0] for x in image_list]

    return image_list

@frappe.whitelist()
def process_new_images(item_code, rte_id, tag):
    """Read images from 'uploads' folder, sort and rename them, resize and
    auto-orient them, copy them to the site public images folder and finally
    create a website slideshow.

    Server-side part of auto-slideshow, called from a button on item page.
    """

    # Whitelisted function; check permissions
    if not frappe.has_permission('Item', 'write'):
        frappe.throw('Need write permissions on Item!',
                     frappe.PermissionError)

    public_site_files_path = os.path.abspath(
        frappe.get_site_path('public', 'files'))

    ret_val = {'success': False}

    # Get user
    current_user = frappe.db.get_value(
        "User", frappe.session.user, ["username"])
    upload_images_directory = os.path.join(uploads_path, current_user)

    slideshow_code = 'SS-' + item_code

    # Check that no slideshow exists for this item, and that we don't have
    # an existing slideshow with a matching name ('SS-ITEM-?????')
    if frappe.db.get_value("Item", item_code, "slideshow", slideshow_code):
        frappe.msgprint("A website slideshow is already set for this item.")
        return ret_val

    if frappe.db.exists("Website Slideshow", slideshow_code):
        frappe.msgprint("A website slideshow with the name " + slideshow_code +
                        " already exists.")
        return ret_val

    # Images should already be uploaded onto local uploads directory
    # Sort these images into a 'natural' order
    file_list = list_files(upload_images_directory)
    n_files = len(file_list)
    file_dict = {}
    for file in file_list:
        file_dict[file] = tuple(int(x) if x.isdigit() else x
                                for x in re_digitsearch.split(file) if x)
    file_list.sort(key=lambda x: file_dict[x])

    if(n_files == 0):
        frappe.msgprint("There are no images to process. " +
                        "Please upload images first.")
        return ret_val

    # Update the number of images to process
    msg = {'command': 'set_image_number',
           'n_images': n_files}
    realtime_eval(rte_id, tag, 'update_slideshow', msg)

    new_file_list = []
    file_sizes = []

    # Rename the files to ITEM-XXXXX-Y and move all the files to
    # public_site_files_path
    w = len(str(n_files))
    for i, fname in enumerate(file_list, 1):
        new_fname = item_code + '-{num:0{width}}.jpg'.format(num=i, width=w)
        new_file_list.append(new_fname)

        upload_fpath = os.path.join(upload_images_directory, fname)
        site_fpath = os.path.join(public_site_files_path, new_fname)
        shutil.move(upload_fpath, site_fpath)

        # Now auto resize the image
        resize_image(site_fpath)

        # Url (relative to hostname) of file
        file_url = os.path.join('files', new_fname)

        # File size
        file_sizes.append(os.path.getsize(site_fpath))

        # Now update the slideshow
        msg = {'command': 'new_image',
               'img_id': i,
               'n_images': n_files,
               'file_url': file_url}
        realtime_eval(rte_id, tag, 'update_slideshow', msg)

    if create_slideshow(slideshow_code):
        create_slideshow_items(slideshow_code, new_file_list, file_sizes)
    else:
        frappe.msgprint("There was a problem creating the slideshow. " +
                        "You will need to do this manually")
        return ret_val

    # For now, assume the first image is the primary image
    # Note this is the idx which is one-indexed, not zero-indexed.
    idx_main_image = 1

    # Update the website slideshow
    frappe.db.set_value('Item', item_code, 'slideshow', slideshow_code)

    # Update the item image
    file_name = new_file_list[idx_main_image-1]
    image_url = os.path.join(
        'files', new_file_list[idx_main_image-1])
    frappe.db.set_value('Item', item_code, 'image', image_url)

    # Set up website image
    web_url, thumb_url = create_website_image(file_name, item_code)

    # Set the website image and thumbnail
    frappe.db.set_value('Item', item_code, 'website_image', web_url)
    frappe.db.set_value('Item', item_code, 'thumbnail', thumb_url)

    # Add a comment to the Item
    frappe.get_doc("Item", item_code).add_comment(
        "Attachment",
        "Auto Create Slideshow: Website slideshow {}".format(slideshow_code))

    # Allow the slideshow to close and update to show completion
    msg = {'command': 'done'}
    realtime_eval(rte_id, tag, 'update_slideshow', msg)

    ret_val['success'] = True
    return ret_val

def create_slideshow(slideshow_name):
    """Create the doc for the Website Slideshow.

    Returns the slideshow doc.
    """

    ss = frappe.get_doc({"doctype": "Website Slideshow",
                         "slideshow_name": slideshow_name})

    ss.insert(ignore_permissions=True)
    return ss

def create_slideshow_items(parent, file_list, file_sizes):
    """Creates the docs for the website slideshow items and the files.

    Creating the docs for the Website Slideshow Items is necessary for
    the website slideshow to work.
    Creating the docs for the Files is not required, but it is what Frappe
    does so we can replicate that. The files are also attached to the Website
    Slideshow rather than the individual Items.
    However: creating individual File docs means that deleting the Website
    Slideshow deletes the files.

    For now: this is deliberately 'broken'.
    """

    idx = 1
    for fname, fsize in zip(file_list, file_sizes):
        filename = '/files/' + fname

        s = frappe.get_doc({
            "doctype": "Website Slideshow Item",
            "parent": parent,
            "parentfield": "slideshow_items",
            "parenttype": "Website Slideshow",
            "image": filename,
            "idx": idx
            })
        idx += 1
        s.insert(ignore_permissions=True)

        # Disabled as we don't want to delete the files if the
        # Website Slideshow is deleted.
        #f = frappe.get_doc({
            #"doctype": "File",
            #"file_url": filename,
            #"file_name": fname,
            #"attached_to_doctype": "Website Slideshow",
            #"attached_to_name": parent,
            #"attached_to_field": None,
            #"folder": 'Home/Attachments',
            #"file_size": fsize,
            #"is_private": 0
        #})
        #f.insert(ignore_permissions=True)

    return

def create_website_image(fname, item):
    """Create a symbolic link and a thumbnail, and set up file docs,
    ready to use as a Website Image.
    """

    public_site_files_path = os.path.abspath(
        frappe.get_site_path('public', 'files'))

    # Create a symbolic link and a thumbnail for the website image
    path, ext = os.path.splitext(fname)
    web_fname = path + '_web' + ext
    thumb_fname = path + '_thumb' + ext

    # Full paths to original file, web image symlink and thumbnail
    file_fpath = os.path.join(public_site_files_path, fname)
    web_fpath = os.path.join(public_site_files_path, web_fname)
    thumb_fpath = os.path.join(public_site_files_path, thumb_fname)

    # URLs on website for web image symlink and thumbnail
    web_url = '/' + os.path.join('files', web_fname)
    thumb_url = '/' + os.path.join('files', thumb_fname)

    # Create the symbolic link and create the thumbnail
    try:
        os.symlink(file_fpath, web_fpath)
    except OSError:
        if os.path.islink(web_fpath):
            os.remove(web_fpath)
            files = frappe.get_all(
                'File', filters={'file_url': web_url})
            for file in files:
                frappe.delete_doc('File', file['name'],
                                  ignore_permissions=True)
            os.symlink(file_fpath, web_fpath)
        else:
            raise
    resize_image(file_fpath, out=thumb_fpath, thumbnail=True)

    # Document for web image
    f = frappe.get_doc({
        "doctype": "File",
        "file_url": web_url,
        "file_name": web_fname,
        "attached_to_doctype": "Item",
        "attached_to_name": item,
        "attached_to_field": None,
        "folder": 'Home/Attachments',
        "file_size": os.path.getsize(web_fpath),
        "is_private": 0
    })
    f.insert(ignore_permissions=True)

    return web_url, thumb_url

def list_files(path):
    # returns a list of names (with extension, without full path) of all files
    # in folder path

    # requires import os

    files = []
    for name in os.listdir(path):
        if os.path.isfile(os.path.join(path, name)):
            files.append(name)
    return files
<commit_msg>fix: V12 issue with write_file hook<commit_after># -*- coding: utf-8 -*-
# Copyright (c) 2015, Universal Resource Trading Limited and contributors
# For license information, please see license.txt

import os
import re
import shutil
import subprocess
from datetime import date

import json
import frappe

from frappe.core.doctype.file.file import File as FrappeFileDoc

uploads_path = os.path.join(os.sep, 'home', 'uploads')

re_digitsearch = re.compile('([0-9]+)')

def realtime_eval(rte_id, tag, event, msg):
    """
    Use publish_realtime to run javascript code using custom event handler.

    rte_id: persistent id for the current form
    tag: communications tag specific to one series of communications
    event: name of custom event to trigger on client.
    msg: JSON serializable object (e.g. dict) to push to client.
    """

    rte_msg = json.dumps(msg)
    js = "erpnext_ebay_realtime_event('{}', '{}', '{}', '{}');"
    js = js.format(rte_id, tag, event, rte_msg)
    frappe.publish_realtime(event='eval_js', message=js,
                            user=frappe.session.user)

def resize_image(filename, out=None, thumbnail=False):
    """
    Use convert or mogrify to resize images
    Uses one of the below commands:

    mogrify -auto-orient -resize LxL> filename
    convert -auto-orient -resize LxL> filename out

    Resizes, rotates according to EXIF information, and resizes to fit into
    an L x L size box (trailing '>' symbol means that images will never
    be enlarged). The image size is taken from the ebay_image_size or
    ebay_thumbnail_size setting in Ebay_Settings_Manager.

    If out is set, the resized version is output to a new file (convert).
    Otherwise the modification takes place on the original file (mogrify).
    """

    if thumbnail:
        image_size = frappe.db.get_value(
            'eBay Manager Settings', filters=None,
            fieldname='ebay_thumbnail_size')
    else:
        image_size = frappe.db.get_value(
            'eBay Manager Settings', filters=None,
            fieldname='ebay_image_size')
    image_size = int(image_size)

    if image_size < 1:
        frappe.throw('Invalid image size: ' + str(image_size))

    size_string = '{}x{}>'.format(image_size, image_size)
    if out is not None:
        subprocess.call(
            ['convert', '-auto-orient', '-resize', size_string,
             filename, out])
    else:
        subprocess.call(
            ['mogrify', '-auto-orient', '-resize', size_string, filename])

def ugs_save_file_on_filesystem_hook(*args, **kwargs):
    """Intercept all write_file events, and mogrify images

    Replaces the standard write_file event. Obtains the filename, content_type
    etc. Calls the normal backup 'save_file_on_filesystem' with all arguments.
    If we do not handle the attachment specially, this function is entirely
    transparent. However, we can handle specific file types as we see fit - in
    this case we mogrify JPG and JPEG images. save_file_on_filesystem does
    strange things, so we need to reconstruct the filename using analagous
    logic - this could break in future with Frappe changes.

    Also this may now (V12) be called from a File doctype, which means we need
    to check for this and deal with it appropriately.
    """

    if len(args) == 1 and isinstance(args[0], FrappeFileDoc):
        # We are being called from a file-type doc
        ret = args[0].save_file_on_filesystem()
    else:
        ret = frappe.utils.file_manager.save_file_on_filesystem(*args, **kwargs)

    file_name = ret['file_name']
    #file_url = ret['file_url'] # Not a consistent file system identifier

    if ('is_private' in kwargs) and kwargs['is_private']:
        file_path = os.path.abspath(
            frappe.get_site_path('private', 'files', file_name))
    else:
        file_path = os.path.abspath(
            frappe.get_site_path('public', 'files', file_name))

    extension = os.path.splitext(file_name)[1].lower()

    if extension in ('.jpg', '.jpeg'):
        # Resize and autoorient this image
        resize_image(file_path)

    return ret

@frappe.whitelist()
def view_slideshow_py(slideshow):

    # Whitelisted function; check permissions
    if not frappe.has_permission('Item', 'read'):
        frappe.throw('Need read permissions on Item!',
                     frappe.PermissionError)

    images_path = os.path.abspath(frappe.get_site_path('public'))

    image_list = frappe.db.sql("""
        SELECT image
        FROM `tabWebsite Slideshow Item`
        WHERE parent=%s
        ORDER BY idx;
        """, (slideshow,), as_dict=False)
    image_list = [x[0] for x in image_list]

    return image_list

@frappe.whitelist()
def process_new_images(item_code, rte_id, tag):
    """Read images from 'uploads' folder, sort and rename them, resize and
    auto-orient them, copy them to the site public images folder and finally
    create a website slideshow.

    Server-side part of auto-slideshow, called from a button on item page.
    """

    # Whitelisted function; check permissions
    if not frappe.has_permission('Item', 'write'):
        frappe.throw('Need write permissions on Item!',
                     frappe.PermissionError)

    public_site_files_path = os.path.abspath(
        frappe.get_site_path('public', 'files'))

    ret_val = {'success': False}

    # Get user
    current_user = frappe.db.get_value(
        "User", frappe.session.user, ["username"])
    upload_images_directory = os.path.join(uploads_path, current_user)

    slideshow_code = 'SS-' + item_code

    # Check that no slideshow exists for this item, and that we don't have
    # an existing slideshow with a matching name ('SS-ITEM-?????')
    if frappe.db.get_value("Item", item_code, "slideshow", slideshow_code):
        frappe.msgprint("A website slideshow is already set for this item.")
        return ret_val

    if frappe.db.exists("Website Slideshow", slideshow_code):
        frappe.msgprint("A website slideshow with the name " + slideshow_code +
                        " already exists.")
        return ret_val

    # Images should already be uploaded onto local uploads directory
    # Sort these images into a 'natural' order
    file_list = list_files(upload_images_directory)
    n_files = len(file_list)
    file_dict = {}
    for file in file_list:
        file_dict[file] = tuple(int(x) if x.isdigit() else x
                                for x in re_digitsearch.split(file) if x)
    file_list.sort(key=lambda x: file_dict[x])

    if(n_files == 0):
        frappe.msgprint("There are no images to process. " +
                        "Please upload images first.")
        return ret_val

    # Update the number of images to process
    msg = {'command': 'set_image_number',
           'n_images': n_files}
    realtime_eval(rte_id, tag, 'update_slideshow', msg)

    new_file_list = []
    file_sizes = []

    # Rename the files to ITEM-XXXXX-Y and move all the files to
    # public_site_files_path
    w = len(str(n_files))
    for i, fname in enumerate(file_list, 1):
        new_fname = item_code + '-{num:0{width}}.jpg'.format(num=i, width=w)
        new_file_list.append(new_fname)

        upload_fpath = os.path.join(upload_images_directory, fname)
        site_fpath = os.path.join(public_site_files_path, new_fname)
        shutil.move(upload_fpath, site_fpath)

        # Now auto resize the image
        resize_image(site_fpath)

        # Url (relative to hostname) of file
        file_url = os.path.join('files', new_fname)

        # File size
        file_sizes.append(os.path.getsize(site_fpath))

        # Now update the slideshow
        msg = {'command': 'new_image',
               'img_id': i,
               'n_images': n_files,
               'file_url': file_url}
        realtime_eval(rte_id, tag, 'update_slideshow', msg)

    if create_slideshow(slideshow_code):
        create_slideshow_items(slideshow_code, new_file_list, file_sizes)
    else:
        frappe.msgprint("There was a problem creating the slideshow. " +
                        "You will need to do this manually")
        return ret_val

    # For now, assume the first image is the primary image
    # Note this is the idx which is one-indexed, not zero-indexed.
    idx_main_image = 1

    # Update the website slideshow
    frappe.db.set_value('Item', item_code, 'slideshow', slideshow_code)

    # Update the item image
    file_name = new_file_list[idx_main_image-1]
    image_url = os.path.join(
        'files', new_file_list[idx_main_image-1])
    frappe.db.set_value('Item', item_code, 'image', image_url)

    # Set up website image
    web_url, thumb_url = create_website_image(file_name, item_code)

    # Set the website image and thumbnail
    frappe.db.set_value('Item', item_code, 'website_image', web_url)
    frappe.db.set_value('Item', item_code, 'thumbnail', thumb_url)

    # Add a comment to the Item
    frappe.get_doc("Item", item_code).add_comment(
        "Attachment",
        "Auto Create Slideshow: Website slideshow {}".format(slideshow_code))

    # Allow the slideshow to close and update to show completion
    msg = {'command': 'done'}
    realtime_eval(rte_id, tag, 'update_slideshow', msg)

    ret_val['success'] = True
    return ret_val

def create_slideshow(slideshow_name):
    """Create the doc for the Website Slideshow.

    Returns the slideshow doc.
    """

    ss = frappe.get_doc({"doctype": "Website Slideshow",
                         "slideshow_name": slideshow_name})

    ss.insert(ignore_permissions=True)
    return ss

def create_slideshow_items(parent, file_list, file_sizes):
    """Creates the docs for the website slideshow items and the files.

    Creating the docs for the Website Slideshow Items is necessary for
    the website slideshow to work.
    Creating the docs for the Files is not required, but it is what Frappe
    does so we can replicate that. The files are also attached to the Website
    Slideshow rather than the individual Items.
    However: creating individual File docs means that deleting the Website
    Slideshow deletes the files.

    For now: this is deliberately 'broken'.
    """

    idx = 1
    for fname, fsize in zip(file_list, file_sizes):
        filename = '/files/' + fname

        s = frappe.get_doc({
            "doctype": "Website Slideshow Item",
            "parent": parent,
            "parentfield": "slideshow_items",
            "parenttype": "Website Slideshow",
            "image": filename,
            "idx": idx
            })
        idx += 1
        s.insert(ignore_permissions=True)

        # Disabled as we don't want to delete the files if the
        # Website Slideshow is deleted.
        #f = frappe.get_doc({
            #"doctype": "File",
            #"file_url": filename,
            #"file_name": fname,
            #"attached_to_doctype": "Website Slideshow",
            #"attached_to_name": parent,
            #"attached_to_field": None,
            #"folder": 'Home/Attachments',
            #"file_size": fsize,
            #"is_private": 0
        #})
        #f.insert(ignore_permissions=True)

    return

def create_website_image(fname, item):
    """Create a symbolic link and a thumbnail, and set up file docs,
    ready to use as a Website Image.
    """

    public_site_files_path = os.path.abspath(
        frappe.get_site_path('public', 'files'))

    # Create a symbolic link and a thumbnail for the website image
    path, ext = os.path.splitext(fname)
    web_fname = path + '_web' + ext
    thumb_fname = path + '_thumb' + ext

    # Full paths to original file, web image symlink and thumbnail
    file_fpath = os.path.join(public_site_files_path, fname)
    web_fpath = os.path.join(public_site_files_path, web_fname)
    thumb_fpath = os.path.join(public_site_files_path, thumb_fname)

    # URLs on website for web image symlink and thumbnail
    web_url = '/' + os.path.join('files', web_fname)
    thumb_url = '/' + os.path.join('files', thumb_fname)

    # Create the symbolic link and create the thumbnail
    try:
        os.symlink(file_fpath, web_fpath)
    except OSError:
        if os.path.islink(web_fpath):
            os.remove(web_fpath)
            files = frappe.get_all(
                'File', filters={'file_url': web_url})
            for file in files:
                frappe.delete_doc('File', file['name'],
                                  ignore_permissions=True)
            os.symlink(file_fpath, web_fpath)
        else:
            raise
    resize_image(file_fpath, out=thumb_fpath, thumbnail=True)

    # Document for web image
    f = frappe.get_doc({
        "doctype": "File",
        "file_url": web_url,
        "file_name": web_fname,
        "attached_to_doctype": "Item",
        "attached_to_name": item,
        "attached_to_field": None,
        "folder": 'Home/Attachments',
        "file_size": os.path.getsize(web_fpath),
        "is_private": 0
    })
    f.insert(ignore_permissions=True)

    return web_url, thumb_url

def list_files(path):
    # returns a list of names (with extension, without full path) of all files
    # in folder path

    # requires import os

    files = []
    for name in os.listdir(path):
        if os.path.isfile(os.path.join(path, name)):
            files.append(name)
    return files
<|endoftext|>


<commit_before>import os, sys
import time
import getpass
import logging

from process import UPID, Process, async

from mesos_pb2 import TASK_LOST
from messages_pb2 import (RegisterFrameworkMessage, ReregisterFrameworkMessage,
        DeactivateFrameworkMessage, UnregisterFrameworkMessage,
        ResourceRequestMessage, ReviveOffersMessage, LaunchTasksMessage, KillTaskMessage,
        StatusUpdate, StatusUpdateAcknowledgementMessage, FrameworkToExecutorMessage)

logger = logging.getLogger(__name__)

class Scheduler(object):
    def registered(self, driver, framework_id, masterInfo): pass
    def reregistered(self, driver, masterInfo): pass
    def disconnected(self, driver): pass
    def frameworkMessage(self, driver, slave_id, executor_id, message): pass
    def resourceOffers(self, driver, offers): pass
    def offerRescinded(self, driver, offer_id): pass
    def statusUpdate(self, driver, status): pass
    def executorLost(self, driver, executor_id, slave_id, status): pass
    def slaveLost(self, driver, slave_id): pass
    def error(self, driver, message): pass

class SchedulerDriver(object):
    def start(self): pass
    def join(self): pass
    def run(self): pass
    def abort(self): pass
    def stop(self, failover=False): pass
    def reviveOffers(self): pass
    def requestResources(self, requests): pass
    def declineOffer(self, offerId, filters=None): pass
    def launchTasks(self, offerId, tasks, filters=None): pass
    def killTask(self, taskId): pass
    def sendFrameworkMessage(self, executorId, slaveId, data): pass

class MesosSchedulerDriver(Process):
    def __init__(self, sched, framework, master_uri):
        Process.__init__(self, 'scheduler')
        self.sched = sched
        #self.executor_info = executor_info
        self.master_uri = master_uri
        self.framework = framework
        self.framework.failover_timeout = 100
        self.framework_id = framework.id

        self.master = None
        self.detector = None

        self.connected = False
        self.savedOffers = {}
        self.savedSlavePids = {}

    @async # called by detector
    def onNewMasterDetectedMessage(self, pid):
        self.master = UPID(pid)
        self.register()

    @async # called by detector
    def onNoMasterDetectedMessage(self):
        self.connected = False
        self.master = None

    def register(self):
        if self.connected or self.aborted:
            return

        if self.master:
            if not self.framework_id.value:
                msg = RegisterFrameworkMessage()
                msg.framework.MergeFrom(self.framework)
            else:
                msg = ReregisterFrameworkMessage()
                msg.framework.MergeFrom(self.framework)
                msg.failover = True
            self.send(self.master, msg)

        self.delay(2, self.register)

    def onFrameworkRegisteredMessage(self, framework_id, master_info):
        self.framework_id = framework_id
        self.framework.id.MergeFrom(framework_id)
        self.connected = True
        self.link(self.master, self.onDisconnected)
        self.sched.registered(self, framework_id, master_info)

    def onFrameworkReregisteredMessage(self, framework_id, master_info):
        assert self.framework_id == framework_id
        self.connected = True
        self.link(self.master, self.onDisconnected)
        self.sched.reregistered(self, master_info)

    def onDisconnected(self):
        self.connected = False
        logger.warning("disconnected from master")
        self.delay(5, self.register)

    def onResourceOffersMessage(self, offers, pids):
        for offer, pid in zip(offers, pids):
            self.savedOffers.setdefault(offer.id.value, {})[offer.slave_id.value] = UPID(pid)
        self.sched.resourceOffers(self, list(offers))

    def onRescindResourceOfferMessage(self, offer_id):
        self.savedOffers.pop(offer_id.value, None)
        self.sched.offerRescinded(self, offer_id)

    def onStatusUpdateMessage(self, update, pid=''):
        assert self.framework_id == update.framework_id
        self.sched.statusUpdate(self, update.status)
        if not self.aborted and pid and not pid.endswith('0.0.0.0'):
            reply = StatusUpdateAcknowledgementMessage()
            reply.framework_id.MergeFrom(self.framework_id)
            reply.slave_id.MergeFrom(update.slave_id)
            reply.task_id.MergeFrom(update.status.task_id)
            reply.uuid = update.uuid
            self.send(UPID(pid), reply)

    def onLostSlaveMessage(self, slave_id):
        self.sched.slaveLost(self, slave_id)

    def onExecutorToFrameworkMessage(self, slave_id, framework_id, executor_id, data):
        self.sched.frameworkMessage(self, slave_id, executor_id, data)

    def onFrameworkErrorMessage(self, message, code=0):
        self.sched.error(self, code, message)

    def start(self):
        Process.start(self)
        uri = self.master_uri
        if uri.startswith('zk://') or uri.startswith('zoo://'):
            from .detector import MasterDetector
            self.detector = MasterDetector(uri[uri.index('://') + 3:], self)
            self.detector.start()
        else:
            if not ':' in uri:
                uri += ':5050'
            self.onNewMasterDetectedMessage('master@%s' % uri)

    def abort(self):
        if self.connected:
            msg = UnregisterFrameworkMessage()
            msg.framework_id.MergeFrom(self.framework_id)
            self.send(self.master, msg)
        Process.abort(self)

    def stop(self, failover=False):
        if self.connected and not failover:
            msg = DeactivateFrameworkMessage()
            msg.framework_id.MergeFrom(self.framework_id)
            self.send(self.master, msg)
        if self.detector:
            self.detector.stop()
        Process.stop(self)

    @async
    def requestResources(self, requests):
        msg = ResourceRequestMessage()
        msg.framework_id.MergeFrom(self.framework_id)
        for req in requests:
            msg.requests.add().MergeFrom(req)
        self.send(self.master, msg)

    @async
    def reviveOffers(self):
        msg = ReviveOffersMessage()
        msg.framework_id.MergeFrom(self.framework_id)
        self.send(self.master, msg)

    def launchTasks(self, offer_id, tasks, filters):
        if not self.connected or offer_id.value not in self.savedOffers:
            update = StatusUpdate()
            update.framework_id.MergeFrom(self.framework_id)
            update.status.task_id.MergeFrom(task.task_id)
            update.status.state = TASK_LOST
            update.status.message = 'Master disconnected' if not self.connected else "invalid offer_id"
            update.timestamp = time.time()
            update.uuid = ''
            return self.onStatusUpdateMessage(update)

        msg = LaunchTasksMessage()
        msg.framework_id.MergeFrom(self.framework_id)
        msg.offer_id.MergeFrom(offer_id)
        msg.filters.MergeFrom(filters)
        for task in tasks:
            msg.tasks.add().MergeFrom(task)
            pid = self.savedOffers.get(offer_id.value, {}).get(task.slave_id.value)
            if pid and task.slave_id.value not in self.savedSlavePids:
                self.savedSlavePids[task.slave_id.value] = pid
        self.savedOffers.pop(offer_id.value)
        self.send(self.master, msg)

    def declineOffer(self, offer_id, filters=None):
        if not self.connected:
            return
        msg = LaunchTasksMessage()
        msg.framework_id.MergeFrom(self.framework_id)
        msg.offer_id.MergeFrom(offer_id)
        if filters:
             msg.filters.MergeFrom(filters)
        self.send(self.master, msg)

    @async
    def killTask(self, task_id):
        if not self.connected:
            return
        msg = KillTaskMessage()
        msg.framework_id.MergeFrom(self.framework_id)
        msg.task_id.MergeFrom(task_id)
        self.send(self.master, msg)

    @async
    def sendFrameworkMessage(self, executor_id, slave_id, data):
        if not self.connected:
            return

        msg = FrameworkToExecutorMessage()
        msg.framework_id.MergeFrom(self.framework_id)
        msg.executor_id.MergeFrom(executor_id)
        msg.slave_id.MergeFrom(slave_id)
        msg.data = data

        slave = self.savedSlavePids.get(slave_id.value, self.master) # can not send to slave directly
        self.send(slave, msg)
<commit_msg>ignore error when reply ack for status update<commit_after>import os, sys
import time
import getpass
import logging

from process import UPID, Process, async

from mesos_pb2 import TASK_LOST
from messages_pb2 import (RegisterFrameworkMessage, ReregisterFrameworkMessage,
        DeactivateFrameworkMessage, UnregisterFrameworkMessage,
        ResourceRequestMessage, ReviveOffersMessage, LaunchTasksMessage, KillTaskMessage,
        StatusUpdate, StatusUpdateAcknowledgementMessage, FrameworkToExecutorMessage)

logger = logging.getLogger(__name__)

class Scheduler(object):
    def registered(self, driver, framework_id, masterInfo): pass
    def reregistered(self, driver, masterInfo): pass
    def disconnected(self, driver): pass
    def frameworkMessage(self, driver, slave_id, executor_id, message): pass
    def resourceOffers(self, driver, offers): pass
    def offerRescinded(self, driver, offer_id): pass
    def statusUpdate(self, driver, status): pass
    def executorLost(self, driver, executor_id, slave_id, status): pass
    def slaveLost(self, driver, slave_id): pass
    def error(self, driver, message): pass

class SchedulerDriver(object):
    def start(self): pass
    def join(self): pass
    def run(self): pass
    def abort(self): pass
    def stop(self, failover=False): pass
    def reviveOffers(self): pass
    def requestResources(self, requests): pass
    def declineOffer(self, offerId, filters=None): pass
    def launchTasks(self, offerId, tasks, filters=None): pass
    def killTask(self, taskId): pass
    def sendFrameworkMessage(self, executorId, slaveId, data): pass

class MesosSchedulerDriver(Process):
    def __init__(self, sched, framework, master_uri):
        Process.__init__(self, 'scheduler')
        self.sched = sched
        #self.executor_info = executor_info
        self.master_uri = master_uri
        self.framework = framework
        self.framework.failover_timeout = 100
        self.framework_id = framework.id

        self.master = None
        self.detector = None

        self.connected = False
        self.savedOffers = {}
        self.savedSlavePids = {}

    @async # called by detector
    def onNewMasterDetectedMessage(self, pid):
        self.master = UPID(pid)
        self.register()

    @async # called by detector
    def onNoMasterDetectedMessage(self):
        self.connected = False
        self.master = None

    def register(self):
        if self.connected or self.aborted:
            return

        if self.master:
            if not self.framework_id.value:
                msg = RegisterFrameworkMessage()
                msg.framework.MergeFrom(self.framework)
            else:
                msg = ReregisterFrameworkMessage()
                msg.framework.MergeFrom(self.framework)
                msg.failover = True
            self.send(self.master, msg)

        self.delay(2, self.register)

    def onFrameworkRegisteredMessage(self, framework_id, master_info):
        self.framework_id = framework_id
        self.framework.id.MergeFrom(framework_id)
        self.connected = True
        self.link(self.master, self.onDisconnected)
        self.sched.registered(self, framework_id, master_info)

    def onFrameworkReregisteredMessage(self, framework_id, master_info):
        assert self.framework_id == framework_id
        self.connected = True
        self.link(self.master, self.onDisconnected)
        self.sched.reregistered(self, master_info)

    def onDisconnected(self):
        self.connected = False
        logger.warning("disconnected from master")
        self.delay(5, self.register)

    def onResourceOffersMessage(self, offers, pids):
        for offer, pid in zip(offers, pids):
            self.savedOffers.setdefault(offer.id.value, {})[offer.slave_id.value] = UPID(pid)
        self.sched.resourceOffers(self, list(offers))

    def onRescindResourceOfferMessage(self, offer_id):
        self.savedOffers.pop(offer_id.value, None)
        self.sched.offerRescinded(self, offer_id)

    def onStatusUpdateMessage(self, update, pid=''):
        assert self.framework_id == update.framework_id

        if pid and not pid.endswith('0.0.0.0'):
            reply = StatusUpdateAcknowledgementMessage()
            reply.framework_id.MergeFrom(self.framework_id)
            reply.slave_id.MergeFrom(update.slave_id)
            reply.task_id.MergeFrom(update.status.task_id)
            reply.uuid = update.uuid
            try: self.send(UPID(pid), reply)
            except IOError: pass

        self.sched.statusUpdate(self, update.status)

    def onLostSlaveMessage(self, slave_id):
        self.sched.slaveLost(self, slave_id)

    def onExecutorToFrameworkMessage(self, slave_id, framework_id, executor_id, data):
        self.sched.frameworkMessage(self, slave_id, executor_id, data)

    def onFrameworkErrorMessage(self, message, code=0):
        self.sched.error(self, code, message)

    def start(self):
        Process.start(self)
        uri = self.master_uri
        if uri.startswith('zk://') or uri.startswith('zoo://'):
            from .detector import MasterDetector
            self.detector = MasterDetector(uri[uri.index('://') + 3:], self)
            self.detector.start()
        else:
            if not ':' in uri:
                uri += ':5050'
            self.onNewMasterDetectedMessage('master@%s' % uri)

    def abort(self):
        if self.connected:
            msg = UnregisterFrameworkMessage()
            msg.framework_id.MergeFrom(self.framework_id)
            self.send(self.master, msg)
        Process.abort(self)

    def stop(self, failover=False):
        if self.connected and not failover:
            msg = DeactivateFrameworkMessage()
            msg.framework_id.MergeFrom(self.framework_id)
            self.send(self.master, msg)
        if self.detector:
            self.detector.stop()
        Process.stop(self)

    @async
    def requestResources(self, requests):
        msg = ResourceRequestMessage()
        msg.framework_id.MergeFrom(self.framework_id)
        for req in requests:
            msg.requests.add().MergeFrom(req)
        self.send(self.master, msg)

    @async
    def reviveOffers(self):
        msg = ReviveOffersMessage()
        msg.framework_id.MergeFrom(self.framework_id)
        self.send(self.master, msg)

    def launchTasks(self, offer_id, tasks, filters):
        if not self.connected or offer_id.value not in self.savedOffers:
            update = StatusUpdate()
            update.framework_id.MergeFrom(self.framework_id)
            update.status.task_id.MergeFrom(task.task_id)
            update.status.state = TASK_LOST
            update.status.message = 'Master disconnected' if not self.connected else "invalid offer_id"
            update.timestamp = time.time()
            update.uuid = ''
            return self.onStatusUpdateMessage(update)

        msg = LaunchTasksMessage()
        msg.framework_id.MergeFrom(self.framework_id)
        msg.offer_id.MergeFrom(offer_id)
        msg.filters.MergeFrom(filters)
        for task in tasks:
            msg.tasks.add().MergeFrom(task)
            pid = self.savedOffers.get(offer_id.value, {}).get(task.slave_id.value)
            if pid and task.slave_id.value not in self.savedSlavePids:
                self.savedSlavePids[task.slave_id.value] = pid
        self.savedOffers.pop(offer_id.value)
        self.send(self.master, msg)

    def declineOffer(self, offer_id, filters=None):
        if not self.connected:
            return
        msg = LaunchTasksMessage()
        msg.framework_id.MergeFrom(self.framework_id)
        msg.offer_id.MergeFrom(offer_id)
        if filters:
             msg.filters.MergeFrom(filters)
        self.send(self.master, msg)

    @async
    def killTask(self, task_id):
        if not self.connected:
            return
        msg = KillTaskMessage()
        msg.framework_id.MergeFrom(self.framework_id)
        msg.task_id.MergeFrom(task_id)
        self.send(self.master, msg)

    @async
    def sendFrameworkMessage(self, executor_id, slave_id, data):
        if not self.connected:
            return

        msg = FrameworkToExecutorMessage()
        msg.framework_id.MergeFrom(self.framework_id)
        msg.executor_id.MergeFrom(executor_id)
        msg.slave_id.MergeFrom(slave_id)
        msg.data = data

        slave = self.savedSlavePids.get(slave_id.value, self.master) # can not send to slave directly
        self.send(slave, msg)
<|endoftext|>