<commit_before>import os

class PipelineStage(object):

    def __init__(self, stage_config, config):
        from brain4k.data import Data

        self.config = stage_config
        self.inputs = [Data(name, config, config['data'][name]) for name in stage_config['inputs']]
        self.outputs = [Data(name, config, config['data'][name]) for name in stage_config['outputs']]
        self.parameters = config['transforms'][stage_config['transform']].get('parameters', {})
        files = config['transforms'][stage_config['transform']].get('files', {})
        self.files = {name: Data(data_name, config, config['data'][data_name]) for name, data_name in files.iteritems()}

    def chain(self, actions):
        for action in actions:
            if not hasattr(self, action):
                raise ValueError(
                    "{0} does not support action {1}".format(self.name, action)
                )

        return [getattr(self, action)() for action in actions]

    def compute_hash(self):
        from brain4k.data_interfaces import compute_json_hash, compute_file_hash

        stage_hashes = []
        varying_data = self.config.get('accept_variance_in', [])
        data = self.inputs + self.files.values() + self.outputs

        non_varying_data = set([datum.filename for datum in data if datum.name not in varying_data])
        for filename in non_varying_data:
            filehash = compute_file_hash(filename)
            stage_hashes.append(filehash)

        stage_hash = compute_json_hash({'stage_hashes': sorted(stage_hashes)})

        return stage_hash

    def blobs_files_exist(self):
        """
        Before computing the sha1 hash, we might want to check that the
        blob files all exist
        """
        data = set(self.inputs) + set(self.files.values()) + set(self.outputs)
        for datum in data:
            if not os.path.exists(datum.filename):
                return False

        return True

TRANSFORMS = {
    "org.scikit-learn.naive_bayes.MultinomialNB": "brain4k.transforms.sklearn.NaiveBayes",
    "org.berkeleyvision.caffe.bvlc_caffenet": "brain4k.transforms.caffe.BVLCCaffeNet",
    "org.scikit-learn.cross_validation.test_train_split": "brain4k.transforms.sklearn.TestTrainSplit",
    "com.brain4k.transforms.data_join": "brain4k.transforms.b4k.DataJoin",
    "org.scikit-learn.metrics.confusion_matrix": "brain4k.transforms.sklearn.metrics.ConfusionMatrix",
}
<commit_msg>delete all output files if an exception is encountered during the stage.  This is to prevent failed outputs (with accept_variance_in=True) causing the stage to be skipped even though it did not complete on the first run (a common scenario when first running a pipeline).<commit_after>import os
import logging

class PipelineStage(object):

    def __init__(self, stage_config, config):
        from brain4k.data import Data

        self.config = stage_config
        self.inputs = [Data(name, config, config['data'][name]) for name in stage_config['inputs']]
        self.outputs = [Data(name, config, config['data'][name]) for name in stage_config['outputs']]
        self.parameters = config['transforms'][stage_config['transform']].get('parameters', {})
        files = config['transforms'][stage_config['transform']].get('files', {})
        self.files = {name: Data(data_name, config, config['data'][data_name]) for name, data_name in files.iteritems()}

    def chain(self, actions):
        for action in actions:
            if not hasattr(self, action):
                raise ValueError(
                    "{0} does not support action {1}".format(self.name, action)
                )

        try:
            results = [getattr(self, action)() for action in actions]
        except Exception as e:
            logging.exception(
                "Encountered unhandled exception during when"
                " calling {0} on {1}: {2}"
                .format(action, self.name, e)
            )
            logging.exception(
                "Deleting all output blobs for stage"
            )
            for datum in self.outputs:
                os.remove(datum.filename)
            raise
        else:
            return results

    def compute_hash(self):
        from brain4k.data_interfaces import compute_json_hash, compute_file_hash

        stage_hashes = []
        varying_data = self.config.get('accept_variance_in', [])
        data = self.inputs + self.files.values() + self.outputs

        non_varying_data = set([datum.filename for datum in data if datum.name not in varying_data])
        for filename in non_varying_data:
            filehash = compute_file_hash(filename)
            stage_hashes.append(filehash)

        stage_hash = compute_json_hash({'stage_hashes': sorted(stage_hashes)})

        return stage_hash

    def blobs_files_exist(self):
        """
        Before computing the sha1 hash, we might want to check that the
        blob files all exist
        """
        data = set(self.inputs) + set(self.files.values()) + set(self.outputs)
        for datum in data:
            if not os.path.exists(datum.filename):
                return False

        return True

TRANSFORMS = {
    "org.scikit-learn.naive_bayes.MultinomialNB": "brain4k.transforms.sklearn.NaiveBayes",
    "org.berkeleyvision.caffe.bvlc_caffenet": "brain4k.transforms.caffe.BVLCCaffeNet",
    "org.scikit-learn.cross_validation.test_train_split": "brain4k.transforms.sklearn.TestTrainSplit",
    "com.brain4k.transforms.data_join": "brain4k.transforms.b4k.DataJoin",
    "org.scikit-learn.metrics.confusion_matrix": "brain4k.transforms.sklearn.metrics.ConfusionMatrix",
}
<|endoftext|>


<commit_before>import boto
from boto import ec2
import sh
import os
import json
import argparse
import sys
from subprocess import Popen, PIPE, STDOUT, call

EXIT_CODE = 0

def deploy():

    # set paths
    packer_file = 'broker-packer.json'
    tfvar_file = 'broker-vars.tf.json'
    packer_exec_path = '/packer/packerio'
    tf_exec_path = '/terraform/terraform'

    #Set connection
    print('Connecting to AWS via region us-gov-west-1...')
    conn = boto.ec2.connect_to_region(region_name='us-gov-west-1')
    print('Done.')

    parser = argparse.ArgumentParser()

    parser.add_argument("-sbx", "--sandbox", action="store_true", help="Runs deploy for sandbox")
    parser.add_argument("-dev", "--dev", action="store_true", help="Runs deploy for dev")
    parser.add_argument("-stg", "--staging", action="store_true", help="Runs deploy for staging")
    parser.add_argument("-prod", "--prod", action="store_true", help="Runs deploy for prod")
    args = parser.parse_args()
    optionsDict = vars(args)
    noArgs = True

    for arg in optionsDict:
        if(optionsDict[arg]):
            noArgs = False

    if noArgs:
        print ("No environment specified. Please include an argument: --sandbox, --dev, --staging, --prod")
        sys.exit(1)

    if optionsDict["sandbox"]:
        deploy_env = 'sandbox'

    if optionsDict["dev"]:
        deploy_env = 'dev'

    if optionsDict["staging"]:
        deploy_env = 'staging'

    if optionsDict["prod"]:
        deploy_env = 'prod'

    tfvar_file = deploy_env + '-variables.tf.json'

    if optionsDict["sandbox"] or optionsDict["dev"]:

        # Retrieve current Base app AMI (where type=app and current=true)
        print('Retrieving current base app AMI...')
        current_base_ami = conn.get_all_images(filters={
            "tag:current" : "True",
            "tag:base" : "True",
            "tag:type" : "Application"
            })[0].id
        print('Done. Current base app AMI: ' + current_base_ami)

        # Insert current Base app AMI into packer file
        print('Updating Packer file with current base app AMI ' + current_base_ami + '...')
        update_packer_spec(packer_file, current_base_ami)
        print('Done.')

        # Retrieve current App Instance AMIs
        print('Retrieving current app instance AMI(s)...')
        current_app_amis = conn.get_all_images(filters={
            "tag:current" : "True",
            "tag:base" : "False",
            "tag:type" : "Application",
            "tag:environment" : deploy_env
            })
        print('Done. Current app instance AMI(s): '+ '\n'.join(map(str, current_app_amis)) )

        # Build new app instance AMI via Packer
        print('**************************************************************************')
        print('Buiding new app instance AMI via Packer. This may take a few minutes...')
        packer_output = real_time_command([packer_exec_path, 'build', packer_file, '-machine-readable'])
        ami_line = [line for line in packer_output.split('\n') if "amazon-ebs: AMIs were created:" in line][0]
        ami_id = ami_line[ami_line.find('ami-'):ami_line.find('ami-')+12]
        print('Done. Packer AMI created: '+ami_id)

        # Set current=False tag for old App AMIs
        if current_app_amis:
            print('Setting current tag to False on old instance AMIs: \n' + '\n'.join(map(str, current_app_amis)) )
            update_ami_tags(current_app_amis)
            print('Done.')
        else:
            print('No matching old AMIs. Skipping tag update...')

        # Confirm app instance that was created is now only current=True tagged AMI
        if ami_id == conn.get_all_images(filters={
            "tag:current" : "True",
            "tag:base" : "False",
            "tag:type" : "Application",
            "tag:environment" : deploy_env
            })[0].id:
            print('Success! Packer AMI id matches current tagged AMI.')
        else:
            print('Something went wrong. Packer AMI: '+ami_id+'; Tagged AMI: '+conn.get_all_images(filters={"tag:current" : "True", "tag:base" : "False", "tag:type" : "Application", "tag:environment" : deploy_env})[0].id)

        #Add new AMI id to terraform variables
        update_lc_ami(ami_id, tfvar_file)

        #Run terraform
        real_time_command([tf_exec_path, 'plan'])
        real_time_command([tf_exec_path, 'apply'])

    elif optionsDict["prod"]:
        #For prod, we don't build a new artifact, we just run terraform against the current staging AMI
        deploy_env = 'prod'
        # get current staging AMI
        staging_ami = conn.get_all_images(filters={"tag:current" : "True", "tag:base" : "False", "tag:type" : "Application", "tag:environment" : "staging"})[0].id
        # Update terraform variables with staging ami_id
        update_lc_ami(staging_ami, 'variables.tf.json')
        print(staging_ami)
        #Run terraform
        real_time_command([tf_exec_path, 'plan'])
        real_time_command([tf_exec_path, 'apply'])

    global EXIT_CODE
    if EXIT_CODE != 0:
        print('Exiting with a code of {}'.format(EXIT_CODE))
        sys.exit(EXIT_CODE)

def real_time_command(command_to_run):
    process = Popen(command_to_run, stdout=PIPE)
    total_output = ''
    while True:
        output = process.stdout.readline().decode()
        if output == '' and process.poll() is not None:
            break
        if output:
            total_output += output

            # Pretty-print human readable output
            if '-machine-readable' in command_to_run:
                output = output[output.rfind(',') + 1:]
            print(output.strip())

    rc = process.poll()
    global EXIT_CODE
    EXIT_CODE += rc

    return total_output

def update_packer_spec(packer_file='packer.json', current_base_ami=''):
    packer_json = open(packer_file, "r")
    packer_data = json.load(packer_json)
    packer_json.close()

    packer_data['builders'][0]['source_ami'] = current_base_ami

    packer_json = open(packer_file, "w+")
    packer_json.write(json.dumps(packer_data))
    packer_json.close()

    return

def update_lc_ami(new_ami='', tfvar_file='variables.tf.json'):
    tfvar_json = open(tfvar_file, "r")
    tfvar_data = json.load(tfvar_json)
    tfvar_json.close()

    tfvar_data['variable']['aws_amis']['default']['us-gov-west-1'] = new_ami

    tfvar_json = open(tfvar_file, "w+")
    tfvar_json.write(json.dumps(tfvar_data))
    tfvar_json.close()

    return

def update_ami_tags(current_app_amis=False):
    for ami in current_app_amis:
        ami.remove_tag('current','True')
        ami.add_tag('current','False')

    return

if __name__ == '__main__':
    deploy()<commit_msg>Updating only the packer file that is relevant for the specified environment<commit_after>import boto
from boto import ec2
import sh
import os
import json
import argparse
import sys
from subprocess import Popen, PIPE, STDOUT, call

EXIT_CODE = 0

def deploy():

    # set paths
    packer_file = 'broker-packer.json'
    tfvar_file = 'broker-vars.tf.json'
    packer_exec_path = '/packer/packerio'
    tf_exec_path = '/terraform/terraform'

    #Set connection
    print('Connecting to AWS via region us-gov-west-1...')
    conn = boto.ec2.connect_to_region(region_name='us-gov-west-1')
    print('Done.')

    parser = argparse.ArgumentParser()

    parser.add_argument("-sbx", "--sandbox", action="store_true", help="Runs deploy for sandbox")
    parser.add_argument("-dev", "--dev", action="store_true", help="Runs deploy for dev")
    parser.add_argument("-stg", "--staging", action="store_true", help="Runs deploy for staging")
    parser.add_argument("-prod", "--prod", action="store_true", help="Runs deploy for prod")
    args = parser.parse_args()
    optionsDict = vars(args)
    noArgs = True

    for arg in optionsDict:
        if(optionsDict[arg]):
            noArgs = False

    if noArgs:
        print ("No environment specified. Please include an argument: --sandbox, --dev, --staging, --prod")
        sys.exit(1)

    if optionsDict["sandbox"]:
        deploy_env = 'sandbox'

    if optionsDict["dev"]:
        deploy_env = 'dev'

    if optionsDict["staging"]:
        deploy_env = 'staging'

    if optionsDict["prod"]:
        deploy_env = 'prod'

    tfvar_file = deploy_env + '-variables.tf.json'
    packer_file = deploy_env + '-packer.json'

    if optionsDict["sandbox"] or optionsDict["dev"]:

        # Retrieve current Base app AMI (where type=app and current=true)
        print('Retrieving current base app AMI...')
        current_base_ami = conn.get_all_images(filters={
            "tag:current" : "True",
            "tag:base" : "True",
            "tag:type" : "Application"
            })[0].id
        print('Done. Current base app AMI: ' + current_base_ami)

        # Insert current Base app AMI into packer file
        print('Updating Packer file with current base app AMI ' + current_base_ami + '...')
        update_packer_spec(packer_file, current_base_ami)
        print('Done.')

        # Retrieve current App Instance AMIs
        print('Retrieving current app instance AMI(s)...')
        current_app_amis = conn.get_all_images(filters={
            "tag:current" : "True",
            "tag:base" : "False",
            "tag:type" : "Application",
            "tag:environment" : deploy_env
            })
        print('Done. Current app instance AMI(s): '+ '\n'.join(map(str, current_app_amis)) )

        # Build new app instance AMI via Packer
        print('**************************************************************************')
        print('Buiding new app instance AMI via Packer. This may take a few minutes...')
        packer_output = real_time_command([packer_exec_path, 'build', packer_file, '-machine-readable'])
        ami_line = [line for line in packer_output.split('\n') if "amazon-ebs: AMIs were created:" in line][0]
        ami_id = ami_line[ami_line.find('ami-'):ami_line.find('ami-')+12]
        print('Done. Packer AMI created: '+ami_id)

        # Set current=False tag for old App AMIs
        if current_app_amis:
            print('Setting current tag to False on old instance AMIs: \n' + '\n'.join(map(str, current_app_amis)) )
            update_ami_tags(current_app_amis)
            print('Done.')
        else:
            print('No matching old AMIs. Skipping tag update...')

        # Confirm app instance that was created is now only current=True tagged AMI
        if ami_id == conn.get_all_images(filters={
            "tag:current" : "True",
            "tag:base" : "False",
            "tag:type" : "Application",
            "tag:environment" : deploy_env
            })[0].id:
            print('Success! Packer AMI id matches current tagged AMI.')
        else:
            print('Something went wrong. Packer AMI: '+ami_id+'; Tagged AMI: '+conn.get_all_images(filters={"tag:current" : "True", "tag:base" : "False", "tag:type" : "Application", "tag:environment" : deploy_env})[0].id)

        #Add new AMI id to terraform variables
        update_lc_ami(ami_id, tfvar_file)

        #Run terraform
        real_time_command([tf_exec_path, 'plan'])
        real_time_command([tf_exec_path, 'apply'])

    elif optionsDict["prod"]:
        #For prod, we don't build a new artifact, we just run terraform against the current staging AMI
        deploy_env = 'prod'
        # get current staging AMI
        staging_ami = conn.get_all_images(filters={"tag:current" : "True", "tag:base" : "False", "tag:type" : "Application", "tag:environment" : "staging"})[0].id
        # Update terraform variables with staging ami_id
        update_lc_ami(staging_ami, 'variables.tf.json')
        print(staging_ami)
        #Run terraform
        real_time_command([tf_exec_path, 'plan'])
        real_time_command([tf_exec_path, 'apply'])

    global EXIT_CODE
    if EXIT_CODE != 0:
        print('Exiting with a code of {}'.format(EXIT_CODE))
        sys.exit(EXIT_CODE)

def real_time_command(command_to_run):
    process = Popen(command_to_run, stdout=PIPE)
    total_output = ''
    while True:
        output = process.stdout.readline().decode()
        if output == '' and process.poll() is not None:
            break
        if output:
            total_output += output

            # Pretty-print human readable output
            if '-machine-readable' in command_to_run:
                output = output[output.rfind(',') + 1:]
            print(output.strip())

    rc = process.poll()
    global EXIT_CODE
    EXIT_CODE += rc

    return total_output

def update_packer_spec(packer_file='packer.json', current_base_ami=''):
    packer_json = open(packer_file, "r")
    packer_data = json.load(packer_json)
    packer_json.close()

    packer_data['builders'][0]['source_ami'] = current_base_ami

    packer_json = open(packer_file, "w+")
    packer_json.write(json.dumps(packer_data))
    packer_json.close()

    return

def update_lc_ami(new_ami='', tfvar_file='variables.tf.json'):
    tfvar_json = open(tfvar_file, "r")
    tfvar_data = json.load(tfvar_json)
    tfvar_json.close()

    tfvar_data['variable']['aws_amis']['default']['us-gov-west-1'] = new_ami

    tfvar_json = open(tfvar_file, "w+")
    tfvar_json.write(json.dumps(tfvar_data))
    tfvar_json.close()

    return

def update_ami_tags(current_app_amis=False):
    for ami in current_app_amis:
        ami.remove_tag('current','True')
        ami.add_tag('current','False')

    return

if __name__ == '__main__':
    deploy()<|endoftext|>


<commit_before>import boto
from boto import ec2
import sh
import os
import json
import argparse
import sys
from subprocess import Popen, PIPE, STDOUT, call

EXIT_CODE = 0
# set global boto connection
conn = boto.ec2.connect_to_region(region_name='us-gov-west-1')

def deploy():

    # set paths
    packer_file = 'broker-packer.json'
    tfvar_file = 'broker-vars.tf.json'
    packer_exec_path = '/packer/packerio'
    tf_exec_path = '/terraform/terraform'

    parser = argparse.ArgumentParser()

    parser.add_argument("-schemasandbox", "--schemasandbox", action="store_true", help="Runs deploy for schemasandbox")
    parser.add_argument("-sbx", "--sandbox", action="store_true", help="Runs deploy for sandbox")
    parser.add_argument("-dev", "--dev", action="store_true", help="Runs deploy for dev")
    parser.add_argument("-stg", "--staging", action="store_true", help="Runs deploy for staging")
    parser.add_argument("-prod", "--prod", action="store_true", help="Runs deploy for prod")
    args = parser.parse_args()
    optionsDict = vars(args)
    noArgs = True

    for arg in optionsDict:
        if(optionsDict[arg]):
            noArgs = False

    if noArgs:
        print ("No environment specified. Please include an argument: --sandbox, --dev, --staging, --prod")
        sys.exit(1)

    if optionsDict["schemasandbox"]:
        deploy_env = 'schemasandbox'

    if optionsDict["sandbox"]:
        deploy_env = 'sandbox'

    if optionsDict["dev"]:
        deploy_env = 'dev'

    if optionsDict["staging"]:
        deploy_env = 'staging'

    if optionsDict["prod"]:
        deploy_env = 'prod'

    tfvar_file = deploy_env + '-variables.tf.json'
    packer_file = deploy_env + '-packer.json'

    if optionsDict["sandbox"]:
        current_base_ami = get_current_base_ami()
        update_lc_ami(current_base_ami, tfvar_file, 'sbx')

        ## Possible to-do: only run terraform when active instances are not running on current base ami
        # val_instance = get_running_instance(deploy_env='sbx', component='Validator')
        # api_instance = get_running_instance(deploy_env='sbx', component='API')
        # if val_instance.image_id != api_instance.image_id:
        #     print("Error, current environment not configured correctly.")
        #     EXIT_CODE = 1
        #     sys.exit(EXIT_CODE)
        # elif val_instance.image_id != current_base_ami: # temporarily forcing false to test
        #     # Create hosts file for ansible and exit with 0 code
        #     with open('hosts','w') as f:
        #         f.write('[validator]\n')
        #         f.write(val_instance.private_ip_address + '\n')
        #         f.write('\n')
        #         f.write('[api]\n')
        #         f.write(api_instance.private_ip_address + '\n')
        #     print("Current base AMI in use. Hosts file has been built for deploying via ansible.")
        #     EXIT_CODE = 0
        #     sys.exit(EXIT_CODE)

        # else:
        #     # Run terraform
        #     update_lc_ami(current_base_ami, 'temp.tf.json', 'sbx')

        #Run terraform

        real_time_command([tf_exec_path, 'plan'])
        real_time_command([tf_exec_path, 'apply'])

    elif optionsDict["dev"] or optionsDict["staging"]:

        # Retrieve current Base app AMI (where type=app and current=true)
        print('Retrieving current base app AMI...')

        current_base_ami = get_current_base_ami()

        # Insert current Base app AMI into packer file
        print('Updating Packer file with current base app AMI ' + current_base_ami + '...')
        update_packer_spec(packer_file, current_base_ami)
        print('Done.')

        # Retrieve current App Instance AMIs
        print('Retrieving current app instance AMI(s)...')
        current_app_amis = get_current_app_instance_amis(deploy_env)

        print('Done. Current app instance AMI(s): '+ '\n'.join(map(str, current_app_amis)) )

        # Build new app instance AMI via Packer
        print('**************************************************************************')
        print('Buiding new app instance AMI via Packer. This may take a few minutes...')
        packer_output = real_time_command([packer_exec_path, 'build', packer_file, '-machine-readable'])
        ami_line = [line for line in packer_output.split('\n') if "amazon-ebs: AMIs were created:" in line][0]
        ami_id = ami_line[ami_line.find('ami-'):ami_line.find('ami-')+12]
        print('Done. Packer AMI created: '+ami_id)

        # Set current=False tag for old App AMIs
        if current_app_amis:
            print('Setting current tag to False on old instance AMIs: \n' + '\n'.join(map(str, current_app_amis)) )
            update_ami_tags(current_app_amis)
            print('Done.')
        else:
            print('No matching old AMIs. Skipping tag update...')

        # Confirm app instance that was created is now only current=True tagged AMI
        if ami_id == conn.get_all_images(filters={
            "tag:current" : "True",
            "tag:base" : "False",
            "tag:type" : "Application",
            "tag:environment" : deploy_env
            })[0].id:
            print('Success! Packer AMI id matches current tagged AMI.')
        else:
            print('Something went wrong. Packer AMI: '+ami_id+'; Tagged AMI: '+conn.get_all_images(filters={"tag:current" : "True", "tag:base" : "False", "tag:type" : "Application", "tag:environment" : deploy_env})[0].id)

        #Add new AMI id to terraform variables
        update_lc_ami(ami_id, tfvar_file, deploy_env)

        #Run terraform
        real_time_command([tf_exec_path, 'plan'])
        real_time_command([tf_exec_path, 'apply'])

    elif optionsDict["prod"]:
        #For prod, we don't build a new artifact, we just run terraform against the current staging AMI
        deploy_env = 'prod'
        # get current staging AMI
        staging_ami = conn.get_all_images(filters={"tag:current" : "True", "tag:base" : "False", "tag:type" : "Application", "tag:environment" : "staging"})[0].id
        # Update terraform variables with staging ami_id
        update_lc_ami(staging_ami, tfvar_file, deploy_env)
        print(staging_ami)
        #Run terraform
        real_time_command([tf_exec_path, 'plan'])
        real_time_command([tf_exec_path, 'apply'])

    global EXIT_CODE
    if EXIT_CODE != 0:
        print('Exiting with a code of {}'.format(EXIT_CODE))
        sys.exit(EXIT_CODE)

def get_running_instance(deploy_env='sbx', component='Validator'):
    reservations = conn.get_all_instances(filters={
        "tag:Application" : "Broker",
        "tag:Component" : component,
        "tag:Environment" : deploy_env
        })
    if len(reservations) != 1:
        print("Error, current environment not configured correctly.")
        EXIT_CODE = 1
        return
    else:
        return reservations[0].instances[0]

def get_current_base_ami():
    base_ami = conn.get_all_images(filters={
        "tag:current" : "True",
        "tag:base" : "True",
        "tag:type" : "Application"
        })[0].id

    return base_ami

def get_current_app_instance_amis(deploy_env):
    app_instance_amis = conn.get_all_images(filters={
        "tag:current" : "True",
        "tag:base" : "False",
        "tag:type" : "Application",
        "tag:environment" : deploy_env
        })
    return app_instance_amis

def real_time_command(command_to_run):
    process = Popen(command_to_run, stdout=PIPE)
    total_output = ''
    while True:
        output = process.stdout.readline().decode()
        if output == '' and process.poll() is not None:
            break
        if output:
            total_output += output

            if '-machine-readable' in command_to_run:
                output = output[output.rfind(',') + 1:]
            print(output.strip())

    rc = process.poll()
    global EXIT_CODE
    EXIT_CODE += rc

    return total_output

def update_packer_spec(packer_file='packer.json', current_base_ami=''):
    packer_json = open(packer_file, "r")
    packer_data = json.load(packer_json)
    packer_json.close()

    packer_data['builders'][0]['source_ami'] = current_base_ami

    packer_json = open(packer_file, "w+")
    packer_json.write(json.dumps(packer_data))
    packer_json.close()

    return

def update_lc_ami(new_ami='', tfvar_file='variables.tf.json', deploy_env='sbx'):
    tfvar_json = open(tfvar_file, "r")
    tfvar_data = json.load(tfvar_json)
    tfvar_json.close()
    if deploy_env == 'sbx':
        tfvar_data['variable']['base_ami']['default'] = new_ami
    else:
        tfvar_data['variable']['aws_amis']['default']['us-gov-west-1'] = new_ami
    tfvar_json = open(tfvar_file, "w+")
    tfvar_json.write(json.dumps(tfvar_data))
    tfvar_json.close()

    return

def update_ami_tags(current_app_amis=False):
    for ami in current_app_amis:
        ami.remove_tag('current','True')
        ami.add_tag('current','False')

    return

if __name__ == '__main__':
    deploy()<commit_msg>Update broker-deploy.py<commit_after>import boto
from boto import ec2
import sh
import os
import json
import argparse
import sys
from subprocess import Popen, PIPE, STDOUT, call

EXIT_CODE = 0
# set global boto connection
conn = boto.ec2.connect_to_region(region_name='us-gov-west-1')

def deploy():

    # set paths
    packer_file = 'broker-packer.json'
    tfvar_file = 'broker-vars.tf.json'
    packer_exec_path = '/packer/packerio'
    tf_exec_path = '/terraform/terraform'

    parser = argparse.ArgumentParser()

    parser.add_argument("-schemasandbox", "--schemasandbox", action="store_true", help="Runs deploy for schemasandbox")
    parser.add_argument("-sbx", "--sandbox", action="store_true", help="Runs deploy for sandbox")
    parser.add_argument("-dev", "--dev", action="store_true", help="Runs deploy for dev")
    parser.add_argument("-stg", "--staging", action="store_true", help="Runs deploy for staging")
    parser.add_argument("-prod", "--prod", action="store_true", help="Runs deploy for prod")
    args = parser.parse_args()
    optionsDict = vars(args)
    noArgs = True

    for arg in optionsDict:
        if(optionsDict[arg]):
            noArgs = False

    if noArgs:
        print ("No environment specified. Please include an argument: --schemasandbox, --sandbox, --dev, --staging, --prod")
        sys.exit(1)

    if optionsDict["schemasandbox"]:
        deploy_env = 'schemasandbox'

    if optionsDict["sandbox"]:
        deploy_env = 'sandbox'

    if optionsDict["dev"]:
        deploy_env = 'dev'

    if optionsDict["staging"]:
        deploy_env = 'staging'

    if optionsDict["prod"]:
        deploy_env = 'prod'

    tfvar_file = deploy_env + '-variables.tf.json'
    packer_file = deploy_env + '-packer.json'

    if optionsDict["sandbox"]:
        current_base_ami = get_current_base_ami()
        update_lc_ami(current_base_ami, tfvar_file, 'sbx')

        ## Possible to-do: only run terraform when active instances are not running on current base ami
        # val_instance = get_running_instance(deploy_env='sbx', component='Validator')
        # api_instance = get_running_instance(deploy_env='sbx', component='API')
        # if val_instance.image_id != api_instance.image_id:
        #     print("Error, current environment not configured correctly.")
        #     EXIT_CODE = 1
        #     sys.exit(EXIT_CODE)
        # elif val_instance.image_id != current_base_ami: # temporarily forcing false to test
        #     # Create hosts file for ansible and exit with 0 code
        #     with open('hosts','w') as f:
        #         f.write('[validator]\n')
        #         f.write(val_instance.private_ip_address + '\n')
        #         f.write('\n')
        #         f.write('[api]\n')
        #         f.write(api_instance.private_ip_address + '\n')
        #     print("Current base AMI in use. Hosts file has been built for deploying via ansible.")
        #     EXIT_CODE = 0
        #     sys.exit(EXIT_CODE)

        # else:
        #     # Run terraform
        #     update_lc_ami(current_base_ami, 'temp.tf.json', 'sbx')

        #Run terraform

        real_time_command([tf_exec_path, 'plan'])
        real_time_command([tf_exec_path, 'apply'])

    elif optionsDict["schemasandbox"] or optionsDict["dev"] or optionsDict["staging"]:

        # Retrieve current Base app AMI (where type=app and current=true)
        print('Retrieving current base app AMI...')

        current_base_ami = get_current_base_ami()

        # Insert current Base app AMI into packer file
        print('Updating Packer file with current base app AMI ' + current_base_ami + '...')
        update_packer_spec(packer_file, current_base_ami)
        print('Done.')

        # Retrieve current App Instance AMIs
        print('Retrieving current app instance AMI(s)...')
        current_app_amis = get_current_app_instance_amis(deploy_env)

        print('Done. Current app instance AMI(s): '+ '\n'.join(map(str, current_app_amis)) )

        # Build new app instance AMI via Packer
        print('**************************************************************************')
        print('Buiding new app instance AMI via Packer. This may take a few minutes...')
        packer_output = real_time_command([packer_exec_path, 'build', packer_file, '-machine-readable'])
        ami_line = [line for line in packer_output.split('\n') if "amazon-ebs: AMIs were created:" in line][0]
        ami_id = ami_line[ami_line.find('ami-'):ami_line.find('ami-')+12]
        print('Done. Packer AMI created: '+ami_id)

        # Set current=False tag for old App AMIs
        if current_app_amis:
            print('Setting current tag to False on old instance AMIs: \n' + '\n'.join(map(str, current_app_amis)) )
            update_ami_tags(current_app_amis)
            print('Done.')
        else:
            print('No matching old AMIs. Skipping tag update...')

        # Confirm app instance that was created is now only current=True tagged AMI
        if ami_id == conn.get_all_images(filters={
            "tag:current" : "True",
            "tag:base" : "False",
            "tag:type" : "Application",
            "tag:environment" : deploy_env
            })[0].id:
            print('Success! Packer AMI id matches current tagged AMI.')
        else:
            print('Something went wrong. Packer AMI: '+ami_id+'; Tagged AMI: '+conn.get_all_images(filters={"tag:current" : "True", "tag:base" : "False", "tag:type" : "Application", "tag:environment" : deploy_env})[0].id)

        #Add new AMI id to terraform variables
        update_lc_ami(ami_id, tfvar_file, deploy_env)

        #Run terraform
        real_time_command([tf_exec_path, 'plan'])
        real_time_command([tf_exec_path, 'apply'])

    elif optionsDict["prod"]:
        #For prod, we don't build a new artifact, we just run terraform against the current staging AMI
        deploy_env = 'prod'
        # get current staging AMI
        staging_ami = conn.get_all_images(filters={"tag:current" : "True", "tag:base" : "False", "tag:type" : "Application", "tag:environment" : "staging"})[0].id
        # Update terraform variables with staging ami_id
        update_lc_ami(staging_ami, tfvar_file, deploy_env)
        print(staging_ami)
        #Run terraform
        real_time_command([tf_exec_path, 'plan'])
        real_time_command([tf_exec_path, 'apply'])

    global EXIT_CODE
    if EXIT_CODE != 0:
        print('Exiting with a code of {}'.format(EXIT_CODE))
        sys.exit(EXIT_CODE)

def get_running_instance(deploy_env='sbx', component='Validator'):
    reservations = conn.get_all_instances(filters={
        "tag:Application" : "Broker",
        "tag:Component" : component,
        "tag:Environment" : deploy_env
        })
    if len(reservations) != 1:
        print("Error, current environment not configured correctly.")
        EXIT_CODE = 1
        return
    else:
        return reservations[0].instances[0]

def get_current_base_ami():
    base_ami = conn.get_all_images(filters={
        "tag:current" : "True",
        "tag:base" : "True",
        "tag:type" : "Application"
        })[0].id

    return base_ami

def get_current_app_instance_amis(deploy_env):
    app_instance_amis = conn.get_all_images(filters={
        "tag:current" : "True",
        "tag:base" : "False",
        "tag:type" : "Application",
        "tag:environment" : deploy_env
        })
    return app_instance_amis

def real_time_command(command_to_run):
    process = Popen(command_to_run, stdout=PIPE)
    total_output = ''
    while True:
        output = process.stdout.readline().decode()
        if output == '' and process.poll() is not None:
            break
        if output:
            total_output += output

            if '-machine-readable' in command_to_run:
                output = output[output.rfind(',') + 1:]
            print(output.strip())

    rc = process.poll()
    global EXIT_CODE
    EXIT_CODE += rc

    return total_output

def update_packer_spec(packer_file='packer.json', current_base_ami=''):
    packer_json = open(packer_file, "r")
    packer_data = json.load(packer_json)
    packer_json.close()

    packer_data['builders'][0]['source_ami'] = current_base_ami

    packer_json = open(packer_file, "w+")
    packer_json.write(json.dumps(packer_data))
    packer_json.close()

    return

def update_lc_ami(new_ami='', tfvar_file='variables.tf.json', deploy_env='sbx'):
    tfvar_json = open(tfvar_file, "r")
    tfvar_data = json.load(tfvar_json)
    tfvar_json.close()
    if deploy_env == 'sbx':
        tfvar_data['variable']['base_ami']['default'] = new_ami
    else:
        tfvar_data['variable']['aws_amis']['default']['us-gov-west-1'] = new_ami
    tfvar_json = open(tfvar_file, "w+")
    tfvar_json.write(json.dumps(tfvar_data))
    tfvar_json.close()

    return

def update_ami_tags(current_app_amis=False):
    for ami in current_app_amis:
        ami.remove_tag('current','True')
        ami.add_tag('current','False')

    return

if __name__ == '__main__':
    deploy()
<|endoftext|>


<commit_before># Copyright 2015 The Meson development team

# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at

#     http://www.apache.org/licenses/LICENSE-2.0

# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

'''This module provides helper functions for RPM related
functionality such as generating RPM spec file.'''

import build
import datetime
import mlog

class RPMModule:

    def generate_spec(self, state, args, kwargs):
        proj = state.project_name.replace(' ', '_').replace('\t', '_')
        so_installed = False
        devel_subpkg = False
        files = []
        files_devel = []
        for target in state.targets.values():
            if isinstance(target, build.Executable) and target.need_install:
                files.append('%%{_bindir}/%s' % target.get_filename())
            elif isinstance(target, build.SharedLibrary) and target.need_install:
                files.append('%%{_libdir}/%s' % target.get_filename())
                for alias in target.get_aliaslist():
                    if alias.endswith('.so'):
                        files_devel.append('%%{_libdir}/%s' % alias)
                    else:
                        files.append('%%{_libdir}/%s' % alias)
                so_installed = True
            elif isinstance(target, build.StaticLibrary) and target.need_install:
                files.append('%%{_libdir}/%s' % target.get_filename())
                mlog.log('Warning, installing static libs (',
                         mlog.bold(target.get_filename()),
                         ') not recommended')
        if len(files_devel) > 0:
            devel_subpkg = True
        fn = open('%s.spec' % proj, 'w+')
        fn.write('Name: %s\n' % proj)
        fn.write('\n')
        for dep in state.environment.coredata.deps:
            fn.write('BuildRequires: pkgconfig(%s)\n' % dep)
        for lib in state.environment.coredata.ext_libs.values():
            fn.write('BuildRequires: %s # FIXME\n' % lib.fullpath)
            mlog.log('Warning, replace', mlog.bold(lib.fullpath), 'with real package.',
                     'You can use following command to find package which contains this lib:',
                     mlog.bold('dnf provides %s' % lib.fullpath))
        for prog in state.environment.coredata.ext_progs.values():
            fn.write('BuildRequires: %s\n' % ' '.join(prog.fullpath))
        fn.write('BuildRequires: meson\n')
        fn.write('\n')
        fn.write('%description\n')
        fn.write('\n')
        if devel_subpkg:
            fn.write('%package devel\n')
            fn.write('Requires: %{name}%{?_isa} = %{version}-%{release}\n')
            fn.write('\n')
            fn.write('%description devel\n')
            fn.write('\n')
        fn.write('%prep\n')
        fn.write('%autosetup\n')
        fn.write('rm -rf build && mkdir build\n')
        fn.write('\n')
        fn.write('%build\n')
        fn.write('pushd build\n')
        fn.write('  meson ..\n')
        fn.write('  ninja-build -v\n')
        fn.write('popd\n')
        fn.write('\n')
        fn.write('%install\n')
        fn.write('pushd build\n')
        fn.write('  DESTDIR=%{buildroot} ninja-build -v install\n')
        fn.write('popd\n')
        fn.write('\n')
        fn.write('%files\n')
        for f in files:
            fn.write('%s\n' % f)
        fn.write('\n')
        if devel_subpkg:
            fn.write('%files devel\n')
            for f in files_devel:
                fn.write('%s\n' % f)
        fn.write('\n')
        if so_installed:
                fn.write('%post -p /sbin/ldconfig\n')
                fn.write('\n')
                fn.write('%postun -p /sbin/ldconfig\n')
        fn.write('\n')
        fn.write('%changelog\n')
        fn.write('* %s meson <meson@example.com> - \n' % datetime.date.today().strftime('%a %b %d %Y'))
        fn.write('- \n')
        fn.write('\n')
        fn.close()

def initialize():
    return RPMModule()
<commit_msg>modules/rpm: ignore static libs<commit_after># Copyright 2015 The Meson development team

# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at

#     http://www.apache.org/licenses/LICENSE-2.0

# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

'''This module provides helper functions for RPM related
functionality such as generating RPM spec file.'''

import build
import datetime
import mlog

class RPMModule:

    def generate_spec(self, state, args, kwargs):
        proj = state.project_name.replace(' ', '_').replace('\t', '_')
        so_installed = False
        devel_subpkg = False
        files = []
        files_devel = []
        to_delete = []
        for target in state.targets.values():
            if isinstance(target, build.Executable) and target.need_install:
                files.append('%%{_bindir}/%s' % target.get_filename())
            elif isinstance(target, build.SharedLibrary) and target.need_install:
                files.append('%%{_libdir}/%s' % target.get_filename())
                for alias in target.get_aliaslist():
                    if alias.endswith('.so'):
                        files_devel.append('%%{_libdir}/%s' % alias)
                    else:
                        files.append('%%{_libdir}/%s' % alias)
                so_installed = True
            elif isinstance(target, build.StaticLibrary) and target.need_install:
                to_delete.append('%%{buildroot}%%{_libdir}/%s' % target.get_filename())
                mlog.log('Ignoring', mlog.bold(target.get_filename()),
                         'because packaging static libs not recommended')
        if len(files_devel) > 0:
            devel_subpkg = True
        fn = open('%s.spec' % proj, 'w+')
        fn.write('Name: %s\n' % proj)
        fn.write('\n')
        for dep in state.environment.coredata.deps:
            fn.write('BuildRequires: pkgconfig(%s)\n' % dep)
        for lib in state.environment.coredata.ext_libs.values():
            fn.write('BuildRequires: %s # FIXME\n' % lib.fullpath)
            mlog.log('Warning, replace', mlog.bold(lib.fullpath), 'with real package.',
                     'You can use following command to find package which contains this lib:',
                     mlog.bold('dnf provides %s' % lib.fullpath))
        for prog in state.environment.coredata.ext_progs.values():
            fn.write('BuildRequires: %s\n' % ' '.join(prog.fullpath))
        fn.write('BuildRequires: meson\n')
        fn.write('\n')
        fn.write('%description\n')
        fn.write('\n')
        if devel_subpkg:
            fn.write('%package devel\n')
            fn.write('Requires: %{name}%{?_isa} = %{version}-%{release}\n')
            fn.write('\n')
            fn.write('%description devel\n')
            fn.write('\n')
        fn.write('%prep\n')
        fn.write('%autosetup\n')
        fn.write('rm -rf build && mkdir build\n')
        fn.write('\n')
        fn.write('%build\n')
        fn.write('pushd build\n')
        fn.write('  meson ..\n')
        fn.write('  ninja-build -v\n')
        fn.write('popd\n')
        fn.write('\n')
        fn.write('%install\n')
        fn.write('pushd build\n')
        fn.write('  DESTDIR=%{buildroot} ninja-build -v install\n')
        fn.write('popd\n')
        if len(to_delete) > 0:
            fn.write('rm -rf %s\n' % ' '.join(to_delete))
        fn.write('\n')
        fn.write('%files\n')
        for f in files:
            fn.write('%s\n' % f)
        fn.write('\n')
        if devel_subpkg:
            fn.write('%files devel\n')
            for f in files_devel:
                fn.write('%s\n' % f)
        fn.write('\n')
        if so_installed:
                fn.write('%post -p /sbin/ldconfig\n')
                fn.write('\n')
                fn.write('%postun -p /sbin/ldconfig\n')
        fn.write('\n')
        fn.write('%changelog\n')
        fn.write('* %s meson <meson@example.com> - \n' % datetime.date.today().strftime('%a %b %d %Y'))
        fn.write('- \n')
        fn.write('\n')
        fn.close()

def initialize():
    return RPMModule()
<|endoftext|>


<commit_before>import datetime
import io
import os
import time
from lxml import etree
try:
    import cPickle as pickle    # NOQA
except:
    import pickle # NOQA

LAST_USED = {}
USE_RATE = 60
PICKLEFILE = os.path.expanduser("~/.pommerss")
CHECK_INTERVAL = 5 * 60
try:
    FEEDCACHE = pickle.load(io.open(PICKLEFILE, 'r+b'))
except:
    FEEDCACHE = {}

class RSSException(Exception):
    pass

class RSSItem(object):
    def __init__(self, node):
        """ Expects an etree-node, so don't try anything else. """
        self.title = node.find("//title").text
        self.link = node.find("//link").text
        self.guid = node.find("//guid").text
        self.link = node.find("//link").text
        self.description = node.find("//description").text
        self.pubdate = node.find("//pubdate").text

class RSSFeed(object):
    def __init__(self, feed):
        self.feed = feed
        self.items = []

    def update(self):
        self.items = []

        print "RSS: Pulling", self.feed
        doc = etree.parse(self.feed)
        items = doc.findall('//items')
        for item in items:
            self.items.append(RSSItem(item))

def pubmsg(connection, event):
    global LAST_USED
    global USE_RATE

    message = event._arguments[0].strip()
    if message[0] != '!':
        return
    (command, remainer) = message.split(" ", 1)

    if command in ('!feeds', '!rss-feeds'):
        uid = '/'.join([connection.server, event.target()])

        # Throttling.
        if uid in LAST_USED and not older_than_rate(LAST_USED[uid],
                                 datetime.datetime.now(),
                                 USE_RATE):
            return
        LAST_USED[uid] = datetime.datetime.now()

        feedlist = [feed for feed, data
                         in FEEDCACHE.iteritems()
                         if uid in data['channels']]
        if len(feedlist) > 0:
            connection.privmsg(event.target(), "Lijstje!")
            for feed in feedlist:
                connection.privmsg(event.target(), str(feed))
                time.sleep(1)
        else:
            connection.privmsg(event.target(),
                               'Nope. Go away.')

def privmsg(connection, event):
    message = event._arguments[0].strip()
    if message[0] != '!':
        return
    (command, remainder) = message.split(" ", 1)

    if command in ('!add-feed',):
        (feed, channels) = remainder.split(" ", 1)
        try:
            add_feed(connection,
                     nick_from_source(event.source()),
                     feed,
                     channels.split())
        except RSSException:
            return
        except Exception as e:
            print "Exception during adding RSS feed:", e

def cycle(connections):
    global FEEDCACHE

    now = datetime.datetime.now()
    feeds = get_feeds_to_check(now)
    print "Checking", feeds
    for feed in feeds:
        FEEDCACHE[feed]['lastchecked'] = now
        # feed.update()

        if feed.has_new(now):
            spam(feed.items[0], FEEDCACHE[feed], connections)
    store_feedcache()

def store_feedcache():
    global FEEDCACHE
    global PICKLEFILE
    pickle.dump(FEEDCACHE, io.open(PICKLEFILE, 'w+b'))

def spam(feed, feeddata, connections):
    print feed
    print feeddata
    print connections

# Supporting functions
def older_than_rate(t1, t2, rate):
    diff = t2 - t1
    return diff.days > 0 or diff.seconds > rate

def get_feeds_to_check(now):
    global FEEDCACHE
    global CHECK_INTERVAL

    to_check = []

    check_time = datetime.timedelta(seconds=CHECK_INTERVAL)
    for feed, data in FEEDCACHE.iteritems():
        if data['lastchecked'] + check_time < now:
            to_check.append(feed)

    return to_check

def add_feed(connection, source, feed, channels):
    connection.privmsg(source, ">>" + str(channels))
    if not channels:
        connection.privmsg(source, "I want channels")
    if not feed in FEEDCACHE:
        if not check_feed(feed):
            raise RSSException("Not a RSS feed")
        add_feed_to_cache(feed)
    add_channels_to_feed([gen_uid(connection.server, channel)
                          for channel
                          in channels], feed)
    store_feedcache()
    connection.privmsg(source, u"added.")

def gen_uid(server, channel):
    return '/'.join([server, channel])

def check_feed(feed):
    try:
        doc = etree.parse(feed)
        return (unicode(doc.getroot().tag) == u'rss')
    except Exception as e:
        print "Exception:", e
        return False

def add_feed_to_cache(feed):
    global FEEDCACHE
    if feed in FEEDCACHE:
        return

    empty_dict = dict()
    empty_dict['lastchecked'] = datetime.datetime(2011, 1, 1)
    empty_dict['lastguid'] = ''
    empty_dict['channels'] = set()

    FEEDCACHE[feed] = empty_dict

def add_channels_to_feed(channels, feed):
    try:
        FEEDCACHE[feed]['channels'].update(channels)
    except Exception as e:
        print e

def nick_from_source(source):
    return source.split('!')[0]
<commit_msg>First version for RSS feeds?<commit_after>import datetime
import io
import os
import time
from lxml import etree
try:
    import cPickle as pickle    # NOQA
except:
    import pickle # NOQA

LAST_USED = {}
USE_RATE = 60
PICKLEFILE = os.path.expanduser("~/.pommerss")
CHECK_INTERVAL = 0.5 * 60
try:
    FEEDCACHE = pickle.load(io.open(PICKLEFILE, 'r+b'))
except:
    FEEDCACHE = {}

class RSSException(Exception):
    pass

class RSSItem(object):
    def __init__(self, node):
        """ Expects an etree-node, so don't try anything else. """
        self.title = node.find("title").text
        self.link = node.find("link").text
        self.guid = node.find("guid").text
        self.link = node.find("link").text
        self.description = node.find("description").text
        self.pubdate = node.find("pubDate").text

class RSSFeed(object):
    def __init__(self, feed):
        self.feed = feed
        self.items = []

    def update(self):
        self.items = []

        print "RSS: Pulling", self.feed
        doc = etree.parse(self.feed)
        items = doc.findall('//item')
        for item in items:
            self.items.append(RSSItem(item))

def pubmsg(connection, event):
    global LAST_USED
    global USE_RATE

    message = event._arguments[0].strip()
    if message[0] != '!':
        return
    (command, remainer) = message.split(" ", 1)

    if command in ('!feeds', '!rss-feeds'):
        uid = '/'.join([connection.server, event.target()])

        # Throttling.
        if uid in LAST_USED and not older_than_rate(LAST_USED[uid],
                                 datetime.datetime.now(),
                                 USE_RATE):
            return
        LAST_USED[uid] = datetime.datetime.now()

        feedlist = [feed for feed, data
                         in FEEDCACHE.iteritems()
                         if uid in data['channels']]
        if len(feedlist) > 0:
            connection.privmsg(event.target(), "Lijstje!")
            for feed in feedlist:
                connection.privmsg(event.target(), str(feed))
                time.sleep(1)
        else:
            connection.privmsg(event.target(),
                               'Nope. Go away.')

def privmsg(connection, event):
    message = event._arguments[0].strip()
    if message[0] != '!':
        return
    (command, remainder) = message.split(" ", 1)

    if command in ('!add-feed',):
        (feed, channels) = remainder.split(" ", 1)
        try:
            add_feed(connection,
                     nick_from_source(event.source()),
                     feed,
                     channels.split())
        except RSSException:
            return
        except Exception as e:
            print "Exception during adding RSS feed:", e
    elif command in ('!clear-cache',):
        global FEEDCACHE
        FEEDCACHE = {}

def cycle(connections):
    global FEEDCACHE

    now = datetime.datetime.now()
    feeds = get_feeds_to_check(now)
    if feeds:
        print "Checking", feeds
    for feed in feeds:
        FEEDCACHE[feed]['lastchecked'] = now

        feed_ = RSSFeed(feed)
        feed_.update()

        if len(feed_.items) == 0:
            continue

        if feed_.items[0].guid > FEEDCACHE[feed]['lastguid']:
            FEEDCACHE[feed]['lastguid'] = feed_.items[0].guid
            spam(feed_.items[0], FEEDCACHE[feed], connections)
    store_feedcache()

def store_feedcache():
    global FEEDCACHE
    global PICKLEFILE
    pickle.dump(FEEDCACHE, io.open(PICKLEFILE, 'w+b'))

def spam(feed, feeddata, connections):
    c_hash = dict([(c.server, c) for c in connections])
    message = '%s - %s' % (feed.title, feed.link)
    for channel in feeddata['channels']:
        server, channel_name = channel.split('/', 1)
        if server in c_hash:
            c = c_hash[server]
            c.privmsg(channel_name, message)

# Supporting functions
def older_than_rate(t1, t2, rate):
    diff = t2 - t1
    return diff.days > 0 or diff.seconds > rate

def get_feeds_to_check(now):
    global FEEDCACHE
    global CHECK_INTERVAL

    to_check = []

    check_time = datetime.timedelta(seconds=CHECK_INTERVAL)
    for feed, data in FEEDCACHE.iteritems():
        if data['lastchecked'] + check_time < now:
            to_check.append(feed)

    return to_check

def add_feed(connection, source, feed, channels):
    connection.privmsg(source, ">>" + str(channels))
    if not channels:
        connection.privmsg(source, "I want channels")
    if not feed in FEEDCACHE:
        if not check_feed(feed):
            raise RSSException("Not a RSS feed")
        add_feed_to_cache(feed)
    add_channels_to_feed([gen_uid(connection.server, channel)
                          for channel
                          in channels], feed)
    store_feedcache()
    connection.privmsg(source, u"added.")

def gen_uid(server, channel):
    return '/'.join([server, channel])

def check_feed(feed):
    try:
        doc = etree.parse(feed)
        return (unicode(doc.getroot().tag) == u'rss')
    except Exception as e:
        print "Exception:", e
        return False

def add_feed_to_cache(feed):
    global FEEDCACHE
    if feed in FEEDCACHE:
        return

    empty_dict = dict()
    empty_dict['lastchecked'] = datetime.datetime(2011, 1, 1)
    empty_dict['lastguid'] = ''
    empty_dict['channels'] = set()

    FEEDCACHE[feed] = empty_dict

def add_channels_to_feed(channels, feed):
    try:
        FEEDCACHE[feed]['channels'].update(channels)
    except Exception as e:
        print e

def nick_from_source(source):
    return source.split('!')[0]
<|endoftext|>


<commit_before>"""Algorithm for simulate a 2048 game using Monte-Carlo method."""

import random, _2048

SIMULATE_TIMES = 100000
DIRECTIONS = ('UP', 'DOWN', 'LEFT', 'RIGHT')

def simulate_to_end(game):
    while game.get_state():
        dircts = list(DIRECTIONS)
        for i in xrange(3):
            c = random.choice(dircts)
            if game.move(c):
                break
            dircts.remove(c)

    return game.get_score()

def score_sum(game,direction):
    score = 0
    temp = game.clone()
    temp.move(direction)
    for i in xrange(SIMULATE_TIMES):
        score += simulate_to_end(temp)
    return score

def monte_carlo(game):
    scores = {}
    biggest = 0
    directions = list(DIRECTIONS)
    for d in DIRECTIONS:
        test = game.clone()
        if not test.move(d):
            directions.remove(d)
    for direction in directions:
        temp = game.clone()
        score = score_sum(temp, direction)
        if score > biggest:
            biggest = score
            best = direction
        scores[direction] = score
    print scores
    if len(set(scores)) == 1:
        return False
    else:
        return best

if __name__ == '__main__':
    a_game = _2048.Gameplay()
    print monte_carlo(a_game)<commit_msg>0.7.1<commit_after>"""Algorithm for simulating a 2048 game using Monte-Carlo method."""

import random, _2048

SIMULATE_TIMES = 100000
DIRECTIONS = ('UP', 'DOWN', 'LEFT', 'RIGHT')

def simulate_to_end(game):
    while game.get_state():
        dircts = list(DIRECTIONS)
        for i in xrange(3):
            c = random.choice(dircts)
            if game.move(c):
                break
            dircts.remove(c)

    return game.get_score()

def score_sum(game,direction):
    score = 0
    temp = game.clone()
    temp.move(direction)
    for i in xrange(SIMULATE_TIMES):
        score += simulate_to_end(temp)
    return score

def monte_carlo(game):
    scores = {}
    biggest = 0
    directions = list(DIRECTIONS)
    for d in DIRECTIONS:
        test = game.clone()
        if not test.move(d):
            directions.remove(d)
    for direction in directions:
        temp = game.clone()
        score = score_sum(temp, direction)
        if score > biggest:
            biggest = score
            best = direction
        scores[direction] = score
    print scores
    if len(set(scores)) == 1:
        return False
    else:
        return best

if __name__ == '__main__':
    a_game = _2048.Gameplay()
    print monte_carlo(a_game)<|endoftext|>


<commit_before>import asyncio
import struct
import socket
import sys
import traceback
from . import defaults
from . import io
from .log import logger

@asyncio.coroutine
def cmd_connect(socks_req, reader, writer, params):
    _command, _addr_type, address, port = socks_req

    logger('socks').debug('Connecting to %s:%d', address, port)

    remote_rw = yield from io.do_connect(address, port)
    if remote_rw is None:
        writer.close()
        return

    bs = params.get('bs', defaults.BLOCK_SIZE)
    try:
        yield from handshake_done(socks_req, reader, writer)
        yield from io.do_streaming(reader, writer, remote_rw[0], remote_rw[1], bs)
    finally:
        remote_rw[1].close()

@asyncio.coroutine
def cmd_udp_relay(socks_req, reader, writer, params):
    _command, _addr_type, address, port = socks_req

    logger('socks').debug('Relaying UDP packets to %s:%d', address, port)

    loop = asyncio.get_event_loop()
    relay_transport, _relay_protocol = \
            yield from loop.create_datagram_endpoint(
                    lambda: io.UDPRelayProtocol(writer),
                    remote_addr=(address, port))

    try:
        yield from handshake_done(socks_req, reader, writer)
        while True:
            data_len_buf = yield from reader.readexactly(2)
            data_len, = struct.unpack('>H', data_len_buf)
            data = yield from reader.readexactly(data_len)
            relay_transport.sendto(data)
    except asyncio.IncompleteReadError:
        logger('socks').debug('UDP_RELAY connection terminated')
    except:
        logger('socks').debug(traceback.format_exc())
    finally:
        relay_transport.close()

supported_cmds = {
    0x01: cmd_connect,
    0x81: cmd_udp_relay,
}

class VersionNotSupportedError(Exception):
    pass

def check_version(ver):
    if ver != 0x05:
        logger('socks').debug('Protocol version %d not supported', ver)
        raise VersionNotSupportedError('Bad version: {}'.format(ver))

@asyncio.coroutine
def recv_auth_method_list(reader):
    header = yield from reader.readexactly(2)
    check_version(header[0])
    methods = header[1]
    if methods > 0:
        method_list = yield from reader.readexactly(methods)
    else:
        method_list = b''
    return method_list

@asyncio.coroutine
def recv_request(reader):
    header = yield from reader.readexactly(4)
    check_version(header[0])

    if header[3] == 0x01:       # ip
        ip_binary = yield from reader.readexactly(4)
        address = '{}.{}.{}.{}'.format(
            ip_binary[0], ip_binary[1], ip_binary[2], ip_binary[3])
    elif header[3] == 0x03:     # domain name
        domain_len = (yield from reader.readexactly(1))[0]
        address = (yield from reader.readexactly(domain_len)).decode()
    elif header[3] == 0x04:     # ipv6
        ip_binary = yield from reader.readexactly(16)
        address = ('{:02x}{:02x}:' * 7 + '{:02x}{:02x}').format(*ip_binary)

    port_binary = yield from reader.readexactly(2)
    port = (port_binary[0] << 8) + port_binary[1]

    #       command,   addr_type
    return (header[1], header[3], address, port)

@asyncio.coroutine
def handshake(reader, writer):
    try:
        method_list = yield from recv_auth_method_list(reader)

        if 0x0 not in method_list:
            logger('socks').debug('No proper method found')
            yield from io.sync_write(writer, b'\x05\xff')
            return None

        yield from io.sync_write(writer, b'\x05\x00')

        command, addr_type, address, port = \
                yield from recv_request(reader)

        if command not in supported_cmds:
            logger('socks').debug('Bad command: %d', command)
            yield from io.sync_write(
                writer, b'\x05\x07\x00\x01\x00\x00\x00\x00\x00\x00')
            return None
    except:
        logger('socks').debug(traceback.format_exc())
        return None

    return (command, addr_type, address, port)

@asyncio.coroutine
def handshake_done(socks_req, reader, writer):
    yield from io.sync_write(
        writer, b'\x05\x00\x00\x01\x00\x00\x00\x00\x00\x00')

@asyncio.coroutine
def open_connection(proxy, cmd, dst):
    reader, writer = yield from asyncio.open_connection(proxy[0], proxy[1])

    try:
        writer.write(b'\x05\x01\x00')
        yield from writer.drain()

        res = yield from reader.readexactly(2)
        assert res == b'\x05\x00'

        if isinstance(dst[0], str):
            addr_buf = socket.inet_aton(dst[0])
            port_buf = dst[1].to_bytes(2, 'big')
        else:
            addr_buf = dst[0].to_bytes(4, sys.byteorder)
            port_buf = dst[1].to_bytes(2, sys.byteorder)
        writer.write(b'\x05' + cmd.to_bytes(1, sys.byteorder) +
                b'\x00\x01' + addr_buf + port_buf)

        res = yield from reader.readexactly(10)
        assert res == b'\x05\x00\x00\x01\x00\x00\x00\x00\x00\x00'
    except:
        writer.close()
        raise

    return (reader, writer)
<commit_msg>Handle unknown address types for socks requests<commit_after>import asyncio
import struct
import socket
import sys
import traceback
from . import defaults
from . import io
from .log import logger

@asyncio.coroutine
def cmd_connect(socks_req, reader, writer, params):
    _command, _addr_type, address, port = socks_req

    logger('socks').debug('Connecting to %s:%d', address, port)

    remote_rw = yield from io.do_connect(address, port)
    if remote_rw is None:
        writer.close()
        return

    bs = params.get('bs', defaults.BLOCK_SIZE)
    try:
        yield from handshake_done(socks_req, reader, writer)
        yield from io.do_streaming(reader, writer, remote_rw[0], remote_rw[1], bs)
    finally:
        remote_rw[1].close()

@asyncio.coroutine
def cmd_udp_relay(socks_req, reader, writer, params):
    _command, _addr_type, address, port = socks_req

    logger('socks').debug('Relaying UDP packets to %s:%d', address, port)

    loop = asyncio.get_event_loop()
    relay_transport, _relay_protocol = \
            yield from loop.create_datagram_endpoint(
                    lambda: io.UDPRelayProtocol(writer),
                    remote_addr=(address, port))

    try:
        yield from handshake_done(socks_req, reader, writer)
        while True:
            data_len_buf = yield from reader.readexactly(2)
            data_len, = struct.unpack('>H', data_len_buf)
            data = yield from reader.readexactly(data_len)
            relay_transport.sendto(data)
    except asyncio.IncompleteReadError:
        logger('socks').debug('UDP_RELAY connection terminated')
    except:
        logger('socks').debug(traceback.format_exc())
    finally:
        relay_transport.close()

supported_cmds = {
    0x01: cmd_connect,
    0x81: cmd_udp_relay,
}

class VersionNotSupportedError(Exception):
    pass

class AddressTypeNotSupportedError(Exception):
    pass

def check_version(ver):
    if ver != 0x05:
        logger('socks').debug('Protocol version %d not supported', ver)
        raise VersionNotSupportedError('Bad version: {}'.format(ver))

@asyncio.coroutine
def recv_auth_method_list(reader):
    header = yield from reader.readexactly(2)
    check_version(header[0])
    methods = header[1]
    if methods > 0:
        method_list = yield from reader.readexactly(methods)
    else:
        method_list = b''
    return method_list

@asyncio.coroutine
def recv_request(reader):
    header = yield from reader.readexactly(4)
    check_version(header[0])

    if header[3] == 0x01:       # ip
        ip_binary = yield from reader.readexactly(4)
        address = '{}.{}.{}.{}'.format(
            ip_binary[0], ip_binary[1], ip_binary[2], ip_binary[3])
    elif header[3] == 0x03:     # domain name
        domain_len = (yield from reader.readexactly(1))[0]
        address = (yield from reader.readexactly(domain_len)).decode()
    elif header[3] == 0x04:     # ipv6
        ip_binary = yield from reader.readexactly(16)
        address = ('{:02x}{:02x}:' * 7 + '{:02x}{:02x}').format(*ip_binary)
    else:
        raise AddressTypeNotSupportedError(
                'Bad address type: {}'.format(header[3]))

    port_binary = yield from reader.readexactly(2)
    port = (port_binary[0] << 8) + port_binary[1]

    #       command,   addr_type
    return (header[1], header[3], address, port)

@asyncio.coroutine
def handshake(reader, writer):
    try:
        method_list = yield from recv_auth_method_list(reader)

        if 0x0 not in method_list:
            logger('socks').debug('No proper method found')
            yield from io.sync_write(writer, b'\x05\xff')
            return None

        yield from io.sync_write(writer, b'\x05\x00')

        command, addr_type, address, port = \
                yield from recv_request(reader)

        if command not in supported_cmds:
            logger('socks').debug('Bad command: %d', command)
            yield from io.sync_write(
                writer, b'\x05\x07\x00\x01\x00\x00\x00\x00\x00\x00')
            return None
    except:
        logger('socks').debug(traceback.format_exc())
        return None

    return (command, addr_type, address, port)

@asyncio.coroutine
def handshake_done(socks_req, reader, writer):
    yield from io.sync_write(
        writer, b'\x05\x00\x00\x01\x00\x00\x00\x00\x00\x00')

@asyncio.coroutine
def open_connection(proxy, cmd, dst):
    reader, writer = yield from asyncio.open_connection(proxy[0], proxy[1])

    try:
        writer.write(b'\x05\x01\x00')
        yield from writer.drain()

        res = yield from reader.readexactly(2)
        assert res == b'\x05\x00'

        if isinstance(dst[0], str):
            addr_buf = socket.inet_aton(dst[0])
            port_buf = dst[1].to_bytes(2, 'big')
        else:
            addr_buf = dst[0].to_bytes(4, sys.byteorder)
            port_buf = dst[1].to_bytes(2, sys.byteorder)
        writer.write(b'\x05' + cmd.to_bytes(1, sys.byteorder) +
                b'\x00\x01' + addr_buf + port_buf)

        res = yield from reader.readexactly(10)
        assert res == b'\x05\x00\x00\x01\x00\x00\x00\x00\x00\x00'
    except:
        writer.close()
        raise

    return (reader, writer)
<|endoftext|>


<commit_before>VERSION = "0.0.14"
<commit_msg>version bump<commit_after>VERSION = "0.0.17"
<|endoftext|>


<commit_before>#!/usr/bin/env python

import logging
import glob
import os
import os.path
import sys
import re
import shlex
import subprocess
import datetime
import imp

import psycopg2
import psycopg2.extensions
import psycopg2.extras

import click

log = logging.getLogger('schematool')

DEFAULT_CONFIG_MODULE_NAME = 'mschematool_config'

### Loading and processing configuration

class Config(object):

    def __init__(self, verbose, config_path):
        self.verbose = verbose
        self.config_path = config_path
        self._module = None

    def _setup_logging(self):
        global log
        log.setLevel(logging.DEBUG)
        formatter = logging.Formatter('%(asctime)-15s %(message)s')
        console_handler = logging.StreamHandler(sys.stderr)
        console_handler.setFormatter(formatter)
        console_handler.setLevel(logging.DEBUG)
        if self.verbose:
            log.addHandler(console_handler)
        if hasattr(self.module, 'LOG_FILE'):
            file_handler = logging.FileHandler(self.module.LOG_FILE)
            file_handler.setLevel(logging.DEBUG)
            file_handler.setFormatter(formatter)
            log.addHandler(file_handler)

    def _load_config(self):
        if self._module is not None:
            return

        if not os.path.exists(self.config_path):
            msg = 'Configuration module %r does not exist' % self.config_path
            sys.stderr.write(msg + '\n')
            log.critical(msg)
            raise Exception(msg)

        try:
            self._module = imp.load_source('mschematool_config', self.config_path)
        except ImportError:
            msg = 'Cannot import mschematool config module'
            sys.stderr.write(msg + '\n')
            log.critical(msg)
            raise

        self._setup_logging()

    @property
    def module(self):
        self._load_config()
        return self._module

### Utility functions

def _simplify_whitespace(s):
    return ' '.join(s.split())

#### Migrations repositories

class MigrationsRepository(object):
    """A repository of migrations is a place where all available migrations are stored
    (for example a directory with migrations as files).
    """

    def get_migrations(self, exclude=None):
        """Return a list of all migrations. In a common case a migration will be a filename,
        without a leading directory part.

        :param exclude: a list or set of migrations to exclude from the result
        """
        raise NotImplementedError()

    def generate_migration_name(self, name, type='sql'):
        """Returns a name of a new migration. It will usually be a filename with
        a valid and unique name.

        :param name: human-readable name of a migration
        :param type: 'sql' or 'py'
        """
        return os.path.join(self.dir,
                            'm{datestr}_{name}.{type}'.format(
                                datestr=datetime.datetime.now().strftime('%Y%m%d%H%M%S'),
                                name=name.replace(' ', '_'),
                                type=type))

    def migration_type(self, migration):
        """Recognize migration type based on a migration (usually a filename).

        :return: 'sql' or 'py'
        """
        if migration.endswith('.sql'):
            return 'sql'
        if migration.endswith('.py'):
            return 'py'
        assert False, 'Invalid migration %r' % migration

class DirRepository(MigrationsRepository):
    """:class:`MigrationsRepository` implementation with migrations being files
    inside a directory ``dir``. Example filenames:
    - m20140615132455_init.sql
    - m20140615135414_insert3.py
    """

    MIGRATION_PATTERN_SQL = 'm*.sql'
    MIGRATION_PATTERN_PY = 'm*.py'

    def __init__(self, dir):
        self.dir = dir

    def _get_all_filenames(self):
        filenames = glob.glob(os.path.join(self.dir, self.MIGRATION_PATTERN_SQL)) + \
            glob.glob(os.path.join(self.dir, self.MIGRATION_PATTERN_PY))
        filenames.sort()
        return filenames

    def get_migrations(self, exclude=None):
        filenames = self._get_all_filenames()
        if exclude:
            filenames = set(filenames) - set(exclude)
            filenames = sorted(filenames)
        return filenames

#### Database-independent interface for migration-related operations

class MigrationsExecutor(object):
    """A class that executes migrations and stores information about execution.
    It will usually store this information inside a database for which migrations
    are tracked.

    :attr:`MigrationsExecutor.engine` is an engine name that can be referenced from
    a config module.

    :param db_config: a dictionary with configuration for a single dbnick.
    :patam repository: :class:`MigrationsRepository` implementation.
    """

    engine = 'unknown'

    def __init__(self, db_config, repository):
        self.db_config = db_config
        self.repository = repository

    def initialize(self):
        """Initialize resources needed for tracking migrations. It will usually
        create a database table for storing information about executed migrations.
        """
        raise NotImplementedError()

    def fetch_executed_migrations(self):
        """Return a list of executed migrations (filenames).
        """
        raise NotImplementedError()

    def execute_python_migration(self, migration, module):
        """Execute a single Python migration and store information about it.

        :param migration: migration (filename) to be executed
        :param module: Python module imported from the migration
        """
        raise NotImplementedError()

    def execute_sql_migration(self, migration):
        """Execute SQL migration from an SQL file.

        :param migration: migration (filename) to be executed
        """
        raise NotImplementedError()

    def execute_migration(self, migration_file_relative):
        """This recognizes migration type and executes either
        :method:`execute_python_migration` or :method:`execute_sql_migration`
        """
        migration_file = os.path.join(self.db_config['migrations_dir'], migration_file_relative)
        m_type = self.repository.migration_type(migration_file)
        if m_type == 'sql':
            return self.execute_sql_migration(migration_file)
        if m_type == 'py':
            module = imp.load_source('migration_module', migration_file)
            return self.execute_python_migration(migration_file, module)
        assert False, 'Unknown migration type %s' % migration_file

## Postgres

class PostgresLoggingDictCursor(psycopg2.extras.DictCursor):
    """Postgres cursor subclass: log all SQL executed in the database.
    """

    def __init__(self, *args, **kwargs):
        psycopg2.extras.DictCursor.__init__(self, *args, **kwargs)

    def execute(self, sql, args=None):
        global log
        if log.isEnabledFor(logging.INFO):
            realsql = self.mogrify(sql, args)
            log.info('Executing SQL: <<%s>>', _simplify_whitespace(realsql))
        try:
            psycopg2.extras.DictCursor.execute(self, sql, args)
        except:
            log.exception('Exception while executing SQL')
            raise

class PostgresMigrations(MigrationsExecutor):

    engine = 'postgres'

    TABLE = 'migration'

    def __init__(self, db_config, repository):
        MigrationsExecutor.__init__(self, db_config, repository)
        self.conn = psycopg2.connect(self.db_config['dsn'])

    def cursor(self):
        return self.conn.cursor(cursor_factory=PostgresLoggingDictCursor)

    def initialize(self):
        with self.cursor() as cur:
            cur.execute("""CREATE TABLE {table} (
                migration_file TEXT,
                execution_time TIMESTAMP DEFAULT CURRENT_TIMESTAMP
            )""".format(table=self.TABLE))
            cur.connection.commit()

    def fetch_executed_migrations(self):
        with self.cursor() as cur:
            cur.execute("""SELECT migration_file FROM {table}
            ORDER BY execution_time""".format(table=self.TABLE))
            return [row[0] for row in cur.fetchall()]

    def _migration_success(self, migration_file):
        with self.cursor() as cur:
            cur.execute("""INSERT INTO {table} (migration_file) VALUES (%s)""".format(table=self.TABLE),
                    [migration_file])

    def execute_python_migration(self, migration_file, module):
        assert hasattr(module, 'migrate'), 'Python module must have `migrate` function accepting ' \
            'a database connection'
        module.migrate(self.conn)
        self._migration_success(migration_file)
        self.conn.commit()

    # https://bitbucket.org/andrewgodwin/south/src/74742a1ba41ce6e9ea56cc694c824b7a93934ac6/south/db/generic.py?at=default
    def _sqlfile_to_statements(self, sql, regex=r"(?mx) ([^';]* (?:'[^']*'[^';]*)*)",
            comment_regex=r"(?mx) (?:^\s*$)|(?:--.*$)"):
        """
        Takes a SQL file and executes it as many separate statements.
        (Some backends, such as Postgres, don't work otherwise.)
        """
        # Be warned: This function is full of dark magic. Make sure you really
        # know regexes before trying to edit it.
        # First, strip comments
        sql = "\n".join([x.strip().replace("%", "%%") for x in re.split(comment_regex, sql) if x.strip()])
        # Now execute each statement
        return re.split(regex, sql)[1:][::2]

    def execute_sql_migration(self, migration_file):
        with open(migration_file) as f:
            content = f.read()
        for statement in self._sqlfile_to_statements(content):
            with self.cursor() as cur:
                cur.execute(statement)
        self._migration_success(migration_file)
        self.conn.commit()

MIGRATIONS_IMPLS = [PostgresMigrations]
ENGINE_TO_IMPL = {m.engine: m for m in MIGRATIONS_IMPLS}

### Integrating all the classes

class MSchemaTool(object):

    def __init__(self, config, dbnick):
        assert dbnick in config.module.DATABASES, 'Not found in DATABASES in config: %s' % dbnick
        self.config = config
        self.dbnick = dbnick
        self.db_config = config.module.DATABASES[dbnick]
        self.repository = DirRepository(self.db_config['migrations_dir'])
        self.migrations = ENGINE_TO_IMPL[self.db_config['engine']](self.db_config, self.repository)

    def not_executed_migration_files(self):
        return self.repository.get_migrations(exclude=self.migrations.fetch_executed_migrations())

    def execute_after_sync(self):
        after_sync = self.db_config.get('after_sync')
        if not after_sync:
            return
        msg = 'Executing after_sync command %r' % after_sync
        log.info(msg)
        click.echo(msg)
        os.system(after_sync)

#### Commands

HELP = """Example usage:

$ ./mschematool.py my_db initdb

A database nickname defined in the configuration module must be passed as the first argument.
After it a command must be specified.

"""

@click.group(help=HELP)
@click.option('--config', type=click.Path(exists=True, dir_okay=False), envvar='MSCHEMATOOL_CONFIG', help='Path to configuration module, e.g. "mydir/mschematool_config.py". Environment variable MSCHEMATOOL_CONFIG can be specified instead.', required=True)
@click.option('--verbose/--no-verbose', default=False, help='Print executed SQL? Default: no.')
@click.argument('dbnick', type=str)
@click.pass_context
def main(ctx, config, verbose, dbnick):
    config_obj = Config(verbose, config)
    run_ctx = MSchemaTool(config_obj, dbnick)
    ctx.obj = run_ctx

@main.command(help='Creates a DB table used for tracking migrations.')
@click.pass_context
def init_db(ctx):
    ctx.obj.migrations.initialize()

@main.command(help='Show synced migrations.')
@click.pass_context
def synced(ctx):
    migrations = ctx.obj.migrations.fetch_executed_migrations()
    for migration in migrations:
        click.echo(migration)

@main.command(help='Show migrations available for syncing.')
@click.pass_context
def to_sync(ctx):
    migrations = ctx.obj.not_executed_migration_files()
    for migration in migrations:
        click.echo(migration)

@main.command(help='Sync all available migrations.')
@click.pass_context
def sync(ctx):
    to_execute = ctx.obj.not_executed_migration_files()
    if not to_execute:
        click.echo('No migrations to sync')
        return
    for migration_file in to_execute:
        msg = 'Executing %s' % migration_file
        log.info(msg)
        click.echo(msg)
        ctx.obj.migrations.execute_migration(migration_file)
    ctx.obj.execute_after_sync()

@main.command(help='Sync a single migration, without syncing older ones.')
@click.argument('migration_file', type=str)
@click.pass_context
def force_sync_single(ctx, migration_file):
    msg = 'Force executing %s' % migration_file
    log.info(msg)
    click.echo(msg)
    ctx.obj.migrations.execute_migration(migration_file)
    ctx.obj.execute_after_sync()

@main.command(help='Print a filename for a new migration.')
@click.argument('name', type=str)
@click.argument('migration_type', type=click.Choice(['sql', 'py']), default='sql')
@click.pass_context
def print_new(ctx, name, migration_type):
    """Prints filename of a new migration"""
    click.echo(ctx.obj.repository.generate_migration_name(name, migration_type))

@main.command(help='Show latest synced migration.')
@click.pass_context
def latest_synced(ctx):
    migrations = ctx.obj.migrations.fetch_executed_migrations()
    if not migrations:
        click.echo('No synced migrations')
    else:
        click.echo(migrations[-1])

if __name__ == '__main__':
    main()

<commit_msg>Shorter column names. Misc fixes.<commit_after>#!/usr/bin/env python

import logging
import glob
import os
import os.path
import sys
import re
import shlex
import subprocess
import datetime
import imp

import psycopg2
import psycopg2.extensions
import psycopg2.extras

import click

log = logging.getLogger('schematool')

DEFAULT_CONFIG_MODULE_NAME = 'mschematool_config'

### Loading and processing configuration

class Config(object):

    def __init__(self, verbose, config_path):
        self.verbose = verbose
        self.config_path = config_path
        self._module = None

    def _setup_logging(self):
        global log
        log.setLevel(logging.DEBUG)
        formatter = logging.Formatter('%(asctime)-15s %(message)s')
        console_handler = logging.StreamHandler(sys.stderr)
        console_handler.setFormatter(formatter)
        console_handler.setLevel(logging.DEBUG)
        if self.verbose:
            log.addHandler(console_handler)
        if hasattr(self.module, 'LOG_FILE'):
            file_handler = logging.FileHandler(self.module.LOG_FILE)
            file_handler.setLevel(logging.DEBUG)
            file_handler.setFormatter(formatter)
            log.addHandler(file_handler)

    def _load_config(self):
        if self._module is not None:
            return

        if not os.path.exists(self.config_path):
            msg = 'Configuration module %r does not exist' % self.config_path
            sys.stderr.write(msg + '\n')
            log.critical(msg)
            raise Exception(msg)

        try:
            self._module = imp.load_source('mschematool_config', self.config_path)
        except ImportError:
            msg = 'Cannot import mschematool config module'
            sys.stderr.write(msg + '\n')
            log.critical(msg)
            raise

        self._setup_logging()

    @property
    def module(self):
        self._load_config()
        return self._module

### Utility functions

def _simplify_whitespace(s):
    return ' '.join(s.split())

#### Migrations repositories

class MigrationsRepository(object):
    """A repository of migrations is a place where all available migrations are stored
    (for example a directory with migrations as files).
    """

    def get_migrations(self, exclude=None):
        """Return a list of all migrations. In a common case a migration will be a filename,
        without a leading directory part.

        :param exclude: a list or set of migrations to exclude from the result
        """
        raise NotImplementedError()

    def generate_migration_name(self, name, type='sql'):
        """Returns a name of a new migration. It will usually be a filename with
        a valid and unique name.

        :param name: human-readable name of a migration
        :param type: 'sql' or 'py'
        """
        return os.path.join(self.dir,
                            'm{datestr}_{name}.{type}'.format(
                                datestr=datetime.datetime.now().strftime('%Y%m%d%H%M%S'),
                                name=name.replace(' ', '_'),
                                type=type))

    def migration_type(self, migration):
        """Recognize migration type based on a migration (usually a filename).

        :return: 'sql' or 'py'
        """
        if migration.endswith('.sql'):
            return 'sql'
        if migration.endswith('.py'):
            return 'py'
        assert False, 'Invalid migration %r' % migration

class DirRepository(MigrationsRepository):
    """:class:`MigrationsRepository` implementation with migrations being files
    inside a directory ``dir``. Example filenames:
    - m20140615132455_init.sql
    - m20140615135414_insert3.py
    """

    MIGRATION_PATTERN_SQL = 'm*.sql'
    MIGRATION_PATTERN_PY = 'm*.py'

    def __init__(self, dir):
        self.dir = dir

    def _get_all_filenames(self):
        filenames = glob.glob(os.path.join(self.dir, self.MIGRATION_PATTERN_SQL)) + \
            glob.glob(os.path.join(self.dir, self.MIGRATION_PATTERN_PY))
        filenames.sort()
        return filenames

    def get_migrations(self, exclude=None):
        filenames = self._get_all_filenames()
        filenames = [os.path.split(fn)[1] for fn in filenames]
        if exclude:
            filenames = set(filenames) - set(exclude)
            filenames = sorted(filenames)
        return filenames

#### Database-independent interface for migration-related operations

class MigrationsExecutor(object):
    """A class that executes migrations and stores information about execution.
    It will usually store this information inside a database for which migrations
    are tracked.

    :attr:`MigrationsExecutor.engine` is an engine name that can be referenced from
    a config module.

    :param db_config: a dictionary with configuration for a single dbnick.
    :patam repository: :class:`MigrationsRepository` implementation.
    """

    engine = 'unknown'

    def __init__(self, db_config, repository):
        self.db_config = db_config
        self.repository = repository

    def initialize(self):
        """Initialize resources needed for tracking migrations. It will usually
        create a database table for storing information about executed migrations.
        """
        raise NotImplementedError()

    def fetch_executed_migrations(self):
        """Return a list of executed migrations (filenames).
        """
        raise NotImplementedError()

    def execute_python_migration(self, migration, module):
        """Execute a single Python migration and store information about it.

        :param migration: migration (filename) to be executed
        :param module: Python module imported from the migration
        """
        raise NotImplementedError()

    def execute_sql_migration(self, migration):
        """Execute SQL migration from an SQL file.

        :param migration: migration (filename) to be executed
        """
        raise NotImplementedError()

    def execute_migration(self, migration_file_relative):
        """This recognizes migration type and executes either
        :method:`execute_python_migration` or :method:`execute_sql_migration`
        """
        migration_file = os.path.join(self.db_config['migrations_dir'], migration_file_relative)
        m_type = self.repository.migration_type(migration_file)
        if m_type == 'sql':
            return self.execute_sql_migration(migration_file)
        if m_type == 'py':
            module = imp.load_source('migration_module', migration_file)
            return self.execute_python_migration(migration_file, module)
        assert False, 'Unknown migration type %s' % migration_file

## Postgres

class PostgresLoggingDictCursor(psycopg2.extras.DictCursor):
    """Postgres cursor subclass: log all SQL executed in the database.
    """

    def __init__(self, *args, **kwargs):
        psycopg2.extras.DictCursor.__init__(self, *args, **kwargs)

    def execute(self, sql, args=None):
        global log
        if log.isEnabledFor(logging.INFO):
            realsql = self.mogrify(sql, args)
            log.info('Executing SQL: <<%s>>', _simplify_whitespace(realsql))
        try:
            psycopg2.extras.DictCursor.execute(self, sql, args)
        except:
            log.exception('Exception while executing SQL')
            raise

class PostgresMigrations(MigrationsExecutor):

    engine = 'postgres'

    TABLE = 'migration'

    def __init__(self, db_config, repository):
        MigrationsExecutor.__init__(self, db_config, repository)
        self.conn = psycopg2.connect(self.db_config['dsn'])

    def cursor(self):
        return self.conn.cursor(cursor_factory=PostgresLoggingDictCursor)

    def initialize(self):
        with self.cursor() as cur:
            cur.execute("""CREATE TABLE {table} (
                file TEXT,
                executed TIMESTAMP DEFAULT CURRENT_TIMESTAMP
            )""".format(table=self.TABLE))
            cur.connection.commit()

    def fetch_executed_migrations(self):
        with self.cursor() as cur:
            cur.execute("""SELECT file FROM {table}
            ORDER BY executed""".format(table=self.TABLE))
            return [row[0] for row in cur.fetchall()]

    def _migration_success(self, migration_file):
        migration = os.path.split(migration_file)[1]
        with self.cursor() as cur:
            cur.execute("""INSERT INTO {table} (file) VALUES (%s)""".format(table=self.TABLE),
                    [migration])

    def execute_python_migration(self, migration_file, module):
        assert hasattr(module, 'migrate'), 'Python module must have `migrate` function accepting ' \
            'a database connection'
        module.migrate(self.conn)
        self._migration_success(migration_file)
        self.conn.commit()

    # https://bitbucket.org/andrewgodwin/south/src/74742a1ba41ce6e9ea56cc694c824b7a93934ac6/south/db/generic.py?at=default
    def _sqlfile_to_statements(self, sql, regex=r"(?mx) ([^';]* (?:'[^']*'[^';]*)*)",
            comment_regex=r"(?mx) (?:^\s*$)|(?:--.*$)"):
        """
        Takes a SQL file and executes it as many separate statements.
        (Some backends, such as Postgres, don't work otherwise.)
        """
        # Be warned: This function is full of dark magic. Make sure you really
        # know regexes before trying to edit it.
        # First, strip comments
        sql = "\n".join([x.strip().replace("%", "%%") for x in re.split(comment_regex, sql) if x.strip()])
        # Now execute each statement
        return re.split(regex, sql)[1:][::2]

    def execute_sql_migration(self, migration_file):
        with open(migration_file) as f:
            content = f.read()
        for statement in self._sqlfile_to_statements(content):
            with self.cursor() as cur:
                cur.execute(statement)
        self._migration_success(migration_file)
        self.conn.commit()

MIGRATIONS_IMPLS = [PostgresMigrations]
ENGINE_TO_IMPL = {m.engine: m for m in MIGRATIONS_IMPLS}

### Integrating all the classes

class MSchemaTool(object):

    def __init__(self, config, dbnick):
        assert dbnick in config.module.DATABASES, 'Not found in DATABASES in config: %s' % dbnick
        self.config = config
        self.dbnick = dbnick
        self.db_config = config.module.DATABASES[dbnick]
        self.repository = DirRepository(self.db_config['migrations_dir'])
        self.migrations = ENGINE_TO_IMPL[self.db_config['engine']](self.db_config, self.repository)

    def not_executed_migration_files(self):
        return self.repository.get_migrations(exclude=self.migrations.fetch_executed_migrations())

    def execute_after_sync(self):
        after_sync = self.db_config.get('after_sync')
        if not after_sync:
            return
        msg = 'Executing after_sync command %r' % after_sync
        log.info(msg)
        click.echo(msg)
        os.system(after_sync)

#### Commands

HELP = """Example usage:

$ ./mschematool.py default init_db

A database nickname defined in the configuration module must be passed as the first argument.
After it a command must be specified.

"""

@click.group(help=HELP)
@click.option('--config', type=click.Path(exists=True, dir_okay=False), envvar='MSCHEMATOOL_CONFIG', help='Path to configuration module, e.g. "mydir/mschematool_config.py". Environment variable MSCHEMATOOL_CONFIG can be specified instead.', required=True)
@click.option('--verbose/--no-verbose', default=False, help='Print executed SQL? Default: no.')
@click.argument('dbnick', type=str)
@click.pass_context
def main(ctx, config, verbose, dbnick):
    config_obj = Config(verbose, config)
    run_ctx = MSchemaTool(config_obj, dbnick)
    ctx.obj = run_ctx

@main.command(help='Creates a DB table used for tracking migrations.')
@click.pass_context
def init_db(ctx):
    ctx.obj.migrations.initialize()

@main.command(help='Show synced migrations.')
@click.pass_context
def synced(ctx):
    migrations = ctx.obj.migrations.fetch_executed_migrations()
    for migration in migrations:
        click.echo(migration)

@main.command(help='Show migrations available for syncing.')
@click.pass_context
def to_sync(ctx):
    migrations = ctx.obj.not_executed_migration_files()
    for migration in migrations:
        click.echo(migration)

@main.command(help='Sync all available migrations.')
@click.pass_context
def sync(ctx):
    to_execute = ctx.obj.not_executed_migration_files()
    if not to_execute:
        click.echo('No migrations to sync')
        return
    for migration_file in to_execute:
        msg = 'Executing %s' % migration_file
        log.info(msg)
        click.echo(msg)
        ctx.obj.migrations.execute_migration(migration_file)
    ctx.obj.execute_after_sync()

@main.command(help='Sync a single migration, without syncing older ones.')
@click.argument('migration_file', type=str)
@click.pass_context
def force_sync_single(ctx, migration_file):
    if migration_file in ctx.obj.migrations.fetch_executed_migrations():
        click.echo('This migration is already executed')
        return
    msg = 'Force executing %s' % migration_file
    log.info(msg)
    click.echo(msg)
    ctx.obj.migrations.execute_migration(migration_file)
    ctx.obj.execute_after_sync()

@main.command(help='Print a filename for a new migration.')
@click.argument('name', type=str)
@click.argument('migration_type', type=click.Choice(['sql', 'py']), default='sql')
@click.pass_context
def print_new(ctx, name, migration_type):
    """Prints filename of a new migration"""
    click.echo(ctx.obj.repository.generate_migration_name(name, migration_type))

@main.command(help='Show latest synced migration.')
@click.pass_context
def latest_synced(ctx):
    migrations = ctx.obj.migrations.fetch_executed_migrations()
    if not migrations:
        click.echo('No synced migrations')
    else:
        click.echo(migrations[-1])

if __name__ == '__main__':
    main()

<|endoftext|>