<commit_before>from datetime import date

from django import forms
from django.contrib import messages
from django.contrib.admin.views.decorators import staff_member_required
from django.contrib.auth.decorators import login_required
from django.contrib.auth.models import User
from django.http import HttpResponseRedirect
from django.shortcuts import redirect, render
from django.utils.translation import ugettext_lazy, ugettext as _

from towel.forms import stripped_formfield_callback

from zivinetz.models import (Drudge, ExpenseReport, Assignment, WaitList,
    Codeword)
from zivinetz.views.decorators import drudge_required

class AssignmentForm(forms.ModelForm):
    codeword = forms.CharField(label=ugettext_lazy('Codeword'))

    formfield_callback = stripped_formfield_callback

    class Meta:
        model = Assignment
        fields = ('specification', 'regional_office', 'date_from',
            'date_until', 'part_of_long_assignment')

    def clean_codeword(self):
        codeword = self.cleaned_data.get('codeword')
        if codeword != Codeword.objects.word(key='einsatz'):
            raise forms.ValidationError(_('Codeword is incorrect.'))
        return codeword

    def clean(self):
        data = super(AssignmentForm, self).clean()

        if data.get('date_from') and data.get('date_until'):
            if data['date_from'] >= data['date_until'] or data['date_from'] < date.today():
                raise forms.ValidationError(_('Date period is invalid.'))

        return data

class WaitListForm(forms.ModelForm):
    codeword = forms.CharField(label=ugettext_lazy('Codeword'))

    formfield_callback = stripped_formfield_callback

    class Meta:
        model = WaitList
        exclude = ('drudge', 'created')

    def clean_codeword(self):
        codeword = self.cleaned_data.get('codeword')
        if codeword != Codeword.objects.word(key='warteliste'):
            raise forms.ValidationError(_('Codeword is incorrect.'))
        return codeword

    def clean(self):
        data = super(WaitListForm, self).clean()

        if data.get('assignment_date_from') and data.get('assignment_date_until'):
            if data['assignment_date_from'] >= data['assignment_date_until'] or data['assignment_date_from'] < date.today():
                raise forms.ValidationError(_('Date period is invalid.'))

        return data

@drudge_required
def dashboard(request, drudge):

    aform_initial = {
        'regional_office': drudge.regional_office_id,
        }

    aform = AssignmentForm(initial=aform_initial)
    wform = WaitListForm()

    if request.method == 'POST':
        if 'assignment' in request.POST:
            aform = AssignmentForm(request.POST)
            if aform.is_valid():
                assignment = aform.save(commit=False)
                assignment.drudge = drudge
                assignment.save()
                messages.success(request, _('Successfully saved new assignment.'))

                return HttpResponseRedirect('.')

        if 'waitlist' in request.POST:
            wform = WaitListForm(request.POST)
            if wform.is_valid():
                waitlist = wform.save(commit=False)
                waitlist.drudge = drudge
                waitlist.save()
                messages.success(request, _('Successfully saved new waitlist entry.'))

                return HttpResponseRedirect('.')

    return render(request, 'zivinetz/drudge_dashboard.html', {
        'drudge': drudge,

        'assignment_form': aform,
        'waitlist_form': wform,

        'assignments': drudge.assignments.order_by('-date_from'),
        'expense_reports': ExpenseReport.objects.filter(assignment__drudge=drudge).order_by('-date_from'),
        'waitlist': WaitList.objects.filter(drudge=drudge),
        })

class UserForm(forms.ModelForm):
    formfield_callback = stripped_formfield_callback
    error_css_class = 'error'
    required_css_class = 'required'

    first_name = forms.CharField(label=_('first name'))
    last_name = forms.CharField(label=_('last name'))

    class Meta:
        model = User
        fields = ('first_name', 'last_name')

class DrudgeForm(forms.ModelForm):
    formfield_callback = stripped_formfield_callback
    error_css_class = 'error'
    required_css_class = 'required'

    class Meta:
        model = Drudge
        exclude = ('user',)

@login_required
def profile(request):
    try:
        drudge = Drudge.objects.get(user=request.user)
    except Drudge.DoesNotExist:
        drudge = None

    if request.method == 'POST':
        form = DrudgeForm(request.POST, instance=drudge)
        form2 = UserForm(request.POST, instance=request.user)

        if form.is_valid() and form2.is_valid():
            drudge = form.save(commit=False)
            drudge.user = request.user
            drudge.save()

            form2.save()

            messages.success(request, _('Successfully saved profile information.'))

            return HttpResponseRedirect('.')
    else:
        form = DrudgeForm(instance=drudge)
        form2 = UserForm(instance=request.user)

    return render(request, 'zivinetz/drudge_profile.html', {
        'object': drudge,
        'form': form,
        'form2': form2,
        'title': _('Edit profile'),
        })
<commit_msg>Do not display pending expense reports in dashboard<commit_after>from datetime import date

from django import forms
from django.contrib import messages
from django.contrib.admin.views.decorators import staff_member_required
from django.contrib.auth.decorators import login_required
from django.contrib.auth.models import User
from django.http import HttpResponseRedirect
from django.shortcuts import redirect, render
from django.utils.translation import ugettext_lazy, ugettext as _

from towel.forms import stripped_formfield_callback

from zivinetz.models import (Drudge, ExpenseReport, Assignment, WaitList,
    Codeword)
from zivinetz.views.decorators import drudge_required

class AssignmentForm(forms.ModelForm):
    codeword = forms.CharField(label=ugettext_lazy('Codeword'))

    formfield_callback = stripped_formfield_callback

    class Meta:
        model = Assignment
        fields = ('specification', 'regional_office', 'date_from',
            'date_until', 'part_of_long_assignment')

    def clean_codeword(self):
        codeword = self.cleaned_data.get('codeword')
        if codeword != Codeword.objects.word(key='einsatz'):
            raise forms.ValidationError(_('Codeword is incorrect.'))
        return codeword

    def clean(self):
        data = super(AssignmentForm, self).clean()

        if data.get('date_from') and data.get('date_until'):
            if data['date_from'] >= data['date_until'] or data['date_from'] < date.today():
                raise forms.ValidationError(_('Date period is invalid.'))

        return data

class WaitListForm(forms.ModelForm):
    codeword = forms.CharField(label=ugettext_lazy('Codeword'))

    formfield_callback = stripped_formfield_callback

    class Meta:
        model = WaitList
        exclude = ('drudge', 'created')

    def clean_codeword(self):
        codeword = self.cleaned_data.get('codeword')
        if codeword != Codeword.objects.word(key='warteliste'):
            raise forms.ValidationError(_('Codeword is incorrect.'))
        return codeword

    def clean(self):
        data = super(WaitListForm, self).clean()

        if data.get('assignment_date_from') and data.get('assignment_date_until'):
            if data['assignment_date_from'] >= data['assignment_date_until'] or data['assignment_date_from'] < date.today():
                raise forms.ValidationError(_('Date period is invalid.'))

        return data

@drudge_required
def dashboard(request, drudge):

    aform_initial = {
        'regional_office': drudge.regional_office_id,
        }

    aform = AssignmentForm(initial=aform_initial)
    wform = WaitListForm()

    if request.method == 'POST':
        if 'assignment' in request.POST:
            aform = AssignmentForm(request.POST)
            if aform.is_valid():
                assignment = aform.save(commit=False)
                assignment.drudge = drudge
                assignment.save()
                messages.success(request, _('Successfully saved new assignment.'))

                return HttpResponseRedirect('.')

        if 'waitlist' in request.POST:
            wform = WaitListForm(request.POST)
            if wform.is_valid():
                waitlist = wform.save(commit=False)
                waitlist.drudge = drudge
                waitlist.save()
                messages.success(request, _('Successfully saved new waitlist entry.'))

                return HttpResponseRedirect('.')

    return render(request, 'zivinetz/drudge_dashboard.html', {
        'drudge': drudge,

        'assignment_form': aform,
        'waitlist_form': wform,

        'assignments': drudge.assignments.order_by('-date_from'),
        'expense_reports': ExpenseReport.objects.filter(
            assignment__drudge=drudge,
            status__in=(ExpenseReport.FILLED, ExpenseReport.PAID),
            ).order_by('-date_from'),
        'waitlist': WaitList.objects.filter(drudge=drudge),
        })

class UserForm(forms.ModelForm):
    formfield_callback = stripped_formfield_callback
    error_css_class = 'error'
    required_css_class = 'required'

    first_name = forms.CharField(label=_('first name'))
    last_name = forms.CharField(label=_('last name'))

    class Meta:
        model = User
        fields = ('first_name', 'last_name')

class DrudgeForm(forms.ModelForm):
    formfield_callback = stripped_formfield_callback
    error_css_class = 'error'
    required_css_class = 'required'

    class Meta:
        model = Drudge
        exclude = ('user',)

@login_required
def profile(request):
    try:
        drudge = Drudge.objects.get(user=request.user)
    except Drudge.DoesNotExist:
        drudge = None

    if request.method == 'POST':
        form = DrudgeForm(request.POST, instance=drudge)
        form2 = UserForm(request.POST, instance=request.user)

        if form.is_valid() and form2.is_valid():
            drudge = form.save(commit=False)
            drudge.user = request.user
            drudge.save()

            form2.save()

            messages.success(request, _('Successfully saved profile information.'))

            return HttpResponseRedirect('.')
    else:
        form = DrudgeForm(instance=drudge)
        form2 = UserForm(instance=request.user)

    return render(request, 'zivinetz/drudge_profile.html', {
        'object': drudge,
        'form': form,
        'form2': form2,
        'title': _('Edit profile'),
        })
<|endoftext|>


<commit_before># Generate .dmp file from .csv
import argparse
import csv
import itertools
import struct
import sys

def header():
    header = "name: sdss_dr7_stars\n"+ \
             "encoding: binary\n"+ \
             "field: location nc_dim_quadtree_{0}\n"+ \
             "metadata: tbin 2016-01-01_00:00:00_3600s\n"+ \
             "field: time nc_dim_time_2\n"
    header = header.format(LEVEL)

    # This includes the count dimension, not includes time dimension
    for i in range(count):
        header = header + 'field: dim' + str(i) + ' nc_var_float_8' + '\n'
    sys.stdout.write(header+'\n')

def body(filepath):
    global flag
    data_schema = []
    with open(filepath, 'r') as csvfile:
        reader = csv.reader(csvfile, delimiter=',')
        for row in reader:
            key = [float(row[6]), float(row[7])]
            v = []
            v.append(float(row[6])-float(row[11]))
            v.append(float(row[7])-float(row[12]))
            v.append(float(row[8])-float(row[13]))
            v.append(float(row[9])-float(row[14]))
            v.append(float(row[10])-float(row[15]))
            v.append(float(row[16]))
            v.append(float(row[17]))
            v.append(float(row[18]))
            v.append(float(row[19]))
            v.append(float(row[20]))

            var = [1]+v
            if flag is True:
                data_schema.append('count')
                for i in range(len(v)):
                    data_schema.append(str(i))

            comb = itertools.combinations_with_replacement(range(DIMS), 2)
            for c in comb:
                var.append(v[int(c[0])]*v[int(c[1])])
                if flag is True:
                    data_schema.append(str(c[0])+'*'+str(c[1]))

            if flag is True:
                print("NanoCube variable dimensions: "+str(count))
                print("NanoCube quadtree level: "+str(LEVEL))
                print('Variable Schema: {}'.format(data_schema))
                flag = False
                if args.s is True:
                    sys.exit(0)

            # Dump
            pack_str = '<iiH' + 'd'*count

            resolution = 2**LEVEL
            xMin = uExtent[0]
            yMin = gExtent[0]
            xRange = uExtent[1]-uExtent[0]
            yRange = gExtent[1]-gExtent[0]

            xTile = int(resolution*((key[0]*1.0-xMin)/xRange))
            yTile = int(resolution*((key[1]*1.0-yMin)/yRange))
            binStr = struct.pack(pack_str,xTile,yTile,0,*var)
            sys.stdout.write(binStr)

if __name__ == '__main__':
    parser = argparse.ArgumentParser(description='Generate testing data')
    parser.add_argument('filepath', help='file path')
    parser.add_argument('-l', type=int, default=15, help='quadtree level')
    parser.add_argument('-s', action='store_true', default=False, help='only show schema of dmp file')
    args = parser.parse_args()

    LEVEL = args.l
    flag = args.s

    DIMS = 10
    count = DIMS*(DIMS+1)/2 + DIMS + 1

    uExtent = [12.0,33.0]
    gExtent = [10.0,33.0]
    rExtent = [10.0,31.0]
    iExtent = [9.0,31.0]
    zExtent = [8.0,29.0]

    if args.s is False:
        header()

    body(args.filepath)
<commit_msg>ommit missing data<commit_after># Generate .dmp file from .csv
import argparse
import csv
import itertools
import struct
import sys

def isfloat(value):
    try:
        float(value)
        return True
    except ValueError:
        return False

def noNull(row):
    legal = True
    for i in range(6, 21):
        if isfloat(row[i]) is False:
            legal = False
    return legal

def header():
    header = "name: sdss_dr7_stars\n"+ \
             "encoding: binary\n"+ \
             "field: location nc_dim_quadtree_{0}\n"+ \
             "metadata: tbin 2016-01-01_00:00:00_3600s\n"+ \
             "field: time nc_dim_time_2\n"
    header = header.format(LEVEL)

    # This includes the count dimension, not includes time dimension
    for i in range(count):
        header = header + 'field: dim' + str(i) + ' nc_var_float_8' + '\n'
    sys.stdout.write(header+'\n')

def body(filepath):
    global flag
    data_schema = []
    filteredCount = 0
    legalCount = 0
    with open(filepath, 'r') as csvfile:
        reader = csv.reader(csvfile, delimiter=',')
        for row in reader:
            if noNull(row) is False:
                filteredCount += 1
            else:
                legalCount += 1
                key = [float(row[6]), float(row[7])]
                v = []
                v.append(float(row[6])-float(row[11]))
                v.append(float(row[7])-float(row[12]))
                v.append(float(row[8])-float(row[13]))
                v.append(float(row[9])-float(row[14]))
                v.append(float(row[10])-float(row[15]))
                v.append(float(row[16]))
                v.append(float(row[17]))
                v.append(float(row[18]))
                v.append(float(row[19]))
                v.append(float(row[20]))

                var = [1]+v
                if flag is True:
                    data_schema.append('count')
                    for i in range(len(v)):
                        data_schema.append(str(i))

                comb = itertools.combinations_with_replacement(range(DIMS), 2)
                for c in comb:
                    var.append(v[int(c[0])]*v[int(c[1])])
                    if flag is True:
                        data_schema.append(str(c[0])+'*'+str(c[1]))

                if flag is True:
                    print("NanoCube variable dimensions: "+str(count))
                    print("NanoCube quadtree level: "+str(LEVEL))
                    print('Variable Schema: {}'.format(data_schema))
                    flag = False
                    if args.s is True:
                        sys.exit(0)

                # Dump
                pack_str = '<iiH' + 'd'*count

                resolution = 2**LEVEL
                xMin = uExtent[0]
                yMin = gExtent[0]
                xRange = uExtent[1]-uExtent[0]
                yRange = gExtent[1]-gExtent[0]

                xTile = int(resolution*((key[0]*1.0-xMin)/xRange))
                yTile = int(resolution*((key[1]*1.0-yMin)/yRange))
                binStr = struct.pack(pack_str,xTile,yTile,0,*var)
                sys.stdout.write(binStr)

    with open(filepath+'.stat', 'w') as f:
        f.write('Original total Count: {}\n'.format(filteredCount+legalCount))
        f.write('Valid row count: {}\n'.format(legalCount))
        f.write('Has missing value: {}\n'.format(filteredCount))

if __name__ == '__main__':
    parser = argparse.ArgumentParser(description='Generate testing data')
    parser.add_argument('filepath', help='file path')
    parser.add_argument('-l', type=int, default=15, help='quadtree level')
    parser.add_argument('-s', action='store_true', default=False, help='only show schema of dmp file')
    args = parser.parse_args()

    LEVEL = args.l
    flag = args.s

    DIMS = 10
    count = DIMS*(DIMS+1)/2 + DIMS + 1

    uExtent = [12.0,33.0]
    gExtent = [10.0,33.0]
    rExtent = [10.0,31.0]
    iExtent = [9.0,31.0]
    zExtent = [8.0,29.0]

    if args.s is False:
        header()

    body(args.filepath)
<|endoftext|>


<commit_before># Paola - paolap@utas.edu.au 23rd August 2013
# This script produces a file listing all the CMIP5 ensembles available on dcc.nci.org.au that satisfy the given constraints.
# The CMIP5 replica dataset is stored under
# /projects/ua6/unofficial_ESG_replica/tmp/tree
# The file list used as input is updated every Monday by Lawson Hanson (CAWCR).
# It creates a csv file called CMIP5_files_in_tree.csv that lists variable, MIP code, model, ensemble, version and path on dcc for each matching ensemble
# Example of how to run on dcc.nci.org.au
#
#    module load python/2.7.1
#    python search_CMIP5_replica.py -f "mon" -v "ua" -v "tas" -e "piControl" -e "historical" -e "rcp45" -m "bcc-csm1-1" -m "CCSM4" output.csv
#
#  Notes concerning the above example:
#	- multiple arguments can be passed to "-v", "-e", "-m", "-t" and -f";
#	- to pass multiple arguments, declare the option multiple times (as above);
#	- you can pass a different name for the output file, just by listing as last argument (output.csv in the example);
#	- all arguments are optional;
#	- failing to set any constraint will result in the entire dataset being selected.
#
#  To quickly view/browse the contents/properties of the resulting file you could do this (in bash):
#
#    cut -d"," -f1-6 CMIP5_files_in_tree.csv | sort
#    wc CMIP5_files_in_tree.csv
#
#  Notes regarding these bash commands:
#	- the output is "sorted" by the first column, which is by variable.
#	- to sort by other columns, you need to rearrange the column order.
#	eg. Sort by mip (col 2), then experiment (col 4).
#  	cut -d"," -f2,4 CMIP5_files_in_tree.csv > cols_sorting.txt
#	cut -d"," -f1,3,5,6- CMIP5_files_in_tree.csv > cols_other.txt
#	NB: "6-" means "column 6 and all remaining columns"
#	paste -d"," cols_sorting.txt cols_other.txt > CMIP5_files_in_tree_rearranged.csv
#
# Feel free to copy and modify this script and please report any bugs
#
#   NB: This code reads the paths of files under the unofficial replica tree from
#       /home/548/lih548/esg-tree-LATEST-paths.txt
#       This file is updated every Monday and kindly shared by Lawson Hanson from CAWCR
#       If you are having problems accessing it or need a more recently updated list, please let me know

import os, datetime, glob, re
import sys, getopt   # these are needed to accept external arguments

## helper functions

def help():
    ''' Print out a help message and exit '''
    print '''Takes the following arguments:\n
              -v / --variable    CMIP5 variable  ex tas\n
              -m / --model       CMIP5 model     ex GFDL-CM3\n
              -e / -- experiment CMIP5 experiment ex historical\n
              -t / --mip_table   CMIP5 MIP table   ex Amon\n
              -f / --frequency   valid values are: day, mon, yr, 3hr, 6hr, subhr, fx, clim\n
              -h / --help display this message and exit \n
              output file , this should always come last, arguments passed after this will be ignored\n
              All other arguments  can be repeated, for example to select two variables  -v tas -v tasmin. All arguments are optional, failing to input any argument will return the entire dataset.\n
              The script return all the ensembles satifying the constraints\n
              [var1 OR var2 OR ..] AND [model1 OR model2 OR ..] AND [exp1 OR exp2 OR ...] \n
              AND [mip1 OR mip2 OR ...] \n
              Frequency add all the correspondent mip_tables to the mip_table list\n
              If a constraint isn't specified for one of the fields automatically all values for that field will be selected.
    '''
    sys.exit()

def file_details(fname):
    ''' Split the filename in variable, MIP code, model, experiment, ensemble (period is excluded) '''
    namebits = fname.split('_')
    if len(namebits) >= 5:
      details = namebits[0:5]
    else:
      details = []
    return details

def match_constraints(details,cons):
    ''' If the variable and experiment match the constraints add file to list '''
    con_num = len(cons)
    i = 0
    details_set = set(details)
    add_file = True
    while i <= con_num -1:
       if len(cons[i]) > 0:
         if list(details_set.intersection(cons[i])) == []:
            add_file = False
            break
       i +=1
    return add_file

def find_string(bits,string):
    ''' Returns matching string if found in directory structure '''
    dummy = filter(lambda el: re.findall( string, el), bits)
    if len(dummy) == 0:
        return 'not_specified'
    else:
        return dummy[0]

def assign_frequency(frq):
    ''' Append the cmip5 mip tables corresponding to the input frequency to the listmip0 '''
    global mip0
    if frq == 'day':
       mip0 = mip0 + ['day', 'cfDay', 'dayExtras']
    elif frq == 'mon':
       mip0 = mip0 + ['Omon', 'OmonExtras', 'Amon', 'AmonExtras', 'Lmon',
                 'LmonExtras', 'OImon', 'LImon', 'cfMon', 'aero', 'cfOff']
    elif frq == '3hr':
       mip0 = mip0 + ['3hr', '3hrLev', 'cf3hr', 'cfSites']
    elif frq == '6hr':
       mip0 = mip0 + ['6hr', '6hrPlev', '6hrLev']
    elif frq == 'monClim':
       mip0 = mip0 + ['Oclim', 'Lclim', 'Aclim', 'LIclim']
    elif frq == 'yr':
       mip0 = mip0 + ['Oyr', 'OyrExtras']
    elif frq == 'fx':
       mip0 = mip0 + ['fx']
    elif frq == 'subhr':
       mip0 = mip0 + ['cfSites']

# Main program starts here
#set up input file and selected variable (or group of variables) and experiment
#infile is updated every Monday and contains a list of all files replicated on dcc
infile = '/home/548/lih548/esg-tree-LATEST-paths.txt'
# assign default values to constraints
var0 = []
exp0 = []
mod0 = []
mip0 = []
outfile = 'CMIP5_files_in_tree.csv'

# assign constraints from arguments list
letters = 'v:m:e:t:f:h' # the : means an argument needs to be passed after the letter
#the = means that a value is expected after the keyword
keywords = ['variable=', 'model=', 'experiment=', 'mip_table=', 'frequency=', 'help=']
opts, extraparams = getopt.getopt(sys.argv[1:],letters,keywords)
# starts at the second element of argv since the first one is the script name
# extraparams are extra arguments passed after all option/keywords are assigned
# in this case the output file
# opts is a list containing the pair "option"/"value"
for o,p in opts:
  if o in ['-v','--variable']:
     var0.append(p)
  elif o in ['-m','--model']:
     mod0.append(p)
  elif o in ['-e','--experiment']:
     exp0.append(p)
  elif o in ['-t','--mip']:
     mip0.append(p)
  elif o in ['-f','--frequency']:
     frq = p
     assign_frequency(frq)
  elif o in ['-h','--help']:
     help()
for p in extraparams:
    outfile = p

# join constraints in a list
constraints = [var0, mod0, exp0, mip0]
for i in range(len(constraints)):
    print keywords[i] + ":  " + str(constraints[i])
print 'Output file: ' + outfile

### needs to include something to take care of all decadals
inf = open(infile, 'r')
outf = open(outfile, 'w')
line1 = 'variable,mip_table,model,experiment,ensemble,version,path\n'
outf.write(line1)

# read all the file paths line by line
lines = inf.readlines()

# define a valid pattern for version
version = '[a-z]*201[0-9][0-1][0-9][0-3][0-9]'
# out_lines is a set of unique lines to add as output, 1 line for each ensemble
out_lines = set()
# loops through all the files
for filepath in lines[:]:
    filepath.replace('\n','')
    bits = filepath.split('/')
    fname = (bits[-1])[:-1]
# call file_details to retrieve experiment, variable, model etc. from filename
    details = file_details(fname)
# make sure details list isn't empty
    if len(details) > 0:
# if file details satisfies the constraint file added to output
      if match_constraints(details,constraints):
         vers = find_string(bits[:-1], version)
         newpath = "/projects/ua6/unofficial-ESG-replica/tmp/tree/" + '/'.join(bits[:-1])
         slist = details + [vers, newpath]
         sline = ','.join(slist)
         out_lines.add(sline)

# write to output file
for sline in out_lines:
    outf.write(sline+"\n")

# close input and output file
inf.close
outf.close

<commit_msg>	modified:   search_CMIP5_replica.py         fix a bug with the --hep keyword and tidied up the comments<commit_after># Paola - paolap@utas.edu.au 23rd August 2013
# This script produces a file listing all the CMIP5 ensembles available on dcc.nci.org.au that satisfy the given constraints.
# The CMIP5 replica dataset is stored under
# /projects/ua6/unofficial_ESG_replica/tmp/tree
# The file list used as input is updated every Monday by Lawson Hanson (CAWCR).
# It creates a csv file called CMIP5_files_in_tree.csv that lists variable, MIP code, model, ensemble, version and path on dcc for each matching ensemble
# Example of how to run on dcc.nci.org.au
#
#    module load python/2.7.1
#    python search_CMIP5_replica.py -f "mon" -v "ua" -v "tas" -e "piControl" -e "historical" -e "rcp45" -m "bcc-csm1-1" -m "CCSM4" output.csv
#
# Notes concerning the above example:
#  - multiple arguments can be passed to "-v", "-e", "-m", "-t" and -f";
#  - to pass multiple arguments, declare the option multiple times (as above);
#  - you can pass a different name for the output file, just by listing as
#    last argument (output.csv in the example);
#  - all arguments are optional;
#  - failing to set any constraint will result in the entire dataset being
#    selected.
#
# To quickly view the output you can use the following bash commands:
#
#   cut -d"," -f1-6 CMIP5_files_in_tree.csv | sort
#   wc CMIP5_files_in_tree.csv
#
# Notes regarding these bash commands:
#   - the output is "sorted" by the first column, which is by variable;
#   - to sort by other columns, you need to rearrange the column order
#    eg. Sort by mip (col 2), then experiment (col 4):
#    cut -d"," -f2,4 CMIP5_files_in_tree.csv > cols_sorting.txt
#    cut -d"," -f1,3,5,6- CMIP5_files_in_tree.csv > cols_other.txt
#    NB: "6-" means "column 6 and all remaining columns"
#    paste -d"," cols_sorting.txt cols_other.txt > CMIP5_files_in_tree_rearranged.csv
#
# Feel free to copy and modify this script and please report any bugs
#
# NB: This code reads the paths of files under the unofficial replica tree from
#       /home/548/lih548/esg-tree-LATEST-paths.txt
#  This file is updated every Monday and kindly shared by Lawson Hanson (CAWCR)
#  If you are having problems accessing it or need a more recently updated list,
#  please let me know

import os, datetime, glob, re
import sys, getopt   # these are needed to accept external arguments

## helper functions

def help():
    ''' Print out a help message and exit '''
    print '''\n
 Takes the following arguments:\n
   -v / --variable    CMIP5 variable  ex tas\n
   -m / --model       CMIP5 model     ex GFDL-CM3\n
   -e / -- experiment CMIP5 experiment ex historical\n
   -t / --mip_table   CMIP5 MIP table   ex Amon\n
   -f / --frequency   valid values are: day, mon, yr, 3hr, 6hr, subhr, fx, clim\n
   -h / --help        display this message and exit \n
   output_file        this should always come last, arguments passed after this\n
                      will be ignored\n
 All other arguments  can be repeated, for example to select two variables:\n
  -v tas -v tasmin .\n
 All arguments are optional, failing to input any argument will return the\n
 entire dataset.\n
 The script returns all the ensembles satifying the constraints\n
  [var1 OR var2 OR ..] AND [model1 OR model2 OR ..] AND [exp1 OR exp2 OR ...] \n
   AND [mip1 OR mip2 OR ...] \n
 Frequency adds all the correspondent mip_tables to the mip_table list\n
 If a constraint isn't specified for one of the fields automatically all values\n
 for that field will be selected.
    '''
    sys.exit()

def file_details(fname):
    ''' Split the filename in variable, MIP code, model, experiment, ensemble (period is excluded) '''
    namebits = fname.split('_')
    if len(namebits) >= 5:
      details = namebits[0:5]
    else:
      details = []
    return details

def match_constraints(details,cons):
    ''' If the variable and experiment match the constraints add file to list '''
    con_num = len(cons)
    i = 0
    details_set = set(details)
    add_file = True
    while i <= con_num -1:
       if len(cons[i]) > 0:
         if list(details_set.intersection(cons[i])) == []:
            add_file = False
            break
       i +=1
    return add_file

def find_string(bits,string):
    ''' Returns matching string if found in directory structure '''
    dummy = filter(lambda el: re.findall( string, el), bits)
    if len(dummy) == 0:
        return 'not_specified'
    else:
        return dummy[0]

def assign_frequency(frq):
    ''' Append the cmip5 mip tables corresponding to the input frequency to the listmip0 '''
    global mip0
    if frq == 'day':
       mip0 = mip0 + ['day', 'cfDay', 'dayExtras']
    elif frq == 'mon':
       mip0 = mip0 + ['Omon', 'OmonExtras', 'Amon', 'AmonExtras', 'Lmon',
                 'LmonExtras', 'OImon', 'LImon', 'cfMon', 'aero', 'cfOff']
    elif frq == '3hr':
       mip0 = mip0 + ['3hr', '3hrLev', 'cf3hr', 'cfSites']
    elif frq == '6hr':
       mip0 = mip0 + ['6hr', '6hrPlev', '6hrLev']
    elif frq == 'monClim':
       mip0 = mip0 + ['Oclim', 'Lclim', 'Aclim', 'LIclim']
    elif frq == 'yr':
       mip0 = mip0 + ['Oyr', 'OyrExtras']
    elif frq == 'fx':
       mip0 = mip0 + ['fx']
    elif frq == 'subhr':
       mip0 = mip0 + ['cfSites']

# Main program starts here
#set up input file and selected variable (or group of variables) and experiment
#infile is updated every Monday and contains a list of all files replicated on dcc
infile = '/home/548/lih548/esg-tree-LATEST-paths.txt'
# assign default values to constraints
var0 = []
exp0 = []
mod0 = []
mip0 = []
outfile = 'CMIP5_files_in_tree.csv'

# assign constraints from arguments list
letters = 'v:m:e:t:f:h' # the : means an argument needs to be passed after the letter
#the = means that a value is expected after the keyword
keywords = ['variable=', 'model=', 'experiment=', 'mip_table=', 'frequency=', 'help']
opts, extraparams = getopt.getopt(sys.argv[1:],letters,keywords)
# starts at the second element of argv since the first one is the script name
# extraparams are extra arguments passed after all option/keywords are assigned
# in this case the output file
# opts is a list containing the pair "option"/"value"
for o,p in opts:
  if o in ['-v','--variable']:
     var0.append(p)
  elif o in ['-m','--model']:
     mod0.append(p)
  elif o in ['-e','--experiment']:
     exp0.append(p)
  elif o in ['-t','--mip']:
     mip0.append(p)
  elif o in ['-f','--frequency']:
     frq = p
     assign_frequency(frq)
  elif o in ['-h','--help']:
     help()
for p in extraparams:
    outfile = p

# join constraints in a list
constraints = [var0, mod0, exp0, mip0]
for i in range(len(constraints)):
    print keywords[i] + ":  " + str(constraints[i])
print 'Output file: ' + outfile

### needs to include something to take care of all decadals
inf = open(infile, 'r')
outf = open(outfile, 'w')
line1 = 'variable,mip_table,model,experiment,ensemble,version,path\n'
outf.write(line1)

# read all the file paths line by line
lines = inf.readlines()

# define a valid pattern for version
version = '[a-z]*201[0-9][0-1][0-9][0-3][0-9]'
# out_lines is a set of unique lines to add as output, 1 line for each ensemble
out_lines = set()
# loops through all the files
for filepath in lines[:]:
    filepath.replace('\n','')
    bits = filepath.split('/')
    fname = (bits[-1])[:-1]
# call file_details to retrieve experiment, variable, model etc. from filename
    details = file_details(fname)
# make sure details list isn't empty
    if len(details) > 0:
# if file details satisfies the constraint file added to output
      if match_constraints(details,constraints):
         vers = find_string(bits[:-1], version)
         newpath = "/projects/ua6/unofficial-ESG-replica/tmp/tree/" + '/'.join(bits[:-1])
         slist = details + [vers, newpath]
         sline = ','.join(slist)
         out_lines.add(sline)

# write to output file
for sline in out_lines:
    outf.write(sline+"\n")

# close input and output file
inf.close
outf.close

<|endoftext|>


<commit_before># -*- coding: utf-8 -*-
#
# Copyright (c) 2015 -- Lars Heuer - Semagia <http://www.semagia.com/>.
# All rights reserved.
#
# License: BSD License
#

__version__ = '0.1.6'
<commit_msg>Incremented version<commit_after># -*- coding: utf-8 -*-
#
# Copyright (c) 2016 -- Lars Heuer - Semagia <http://www.semagia.com/>.
# All rights reserved.
#
# License: BSD License
#

__version__ = '0.1.7'
<|endoftext|>


<commit_before>import argparse
import datetime
import logging
import os
import re
import json
import tempfile
import boto3
import requests

from dataactcore.config import CONFIG_BROKER
from dataactcore.interfaces.db import GlobalDB
from dataactcore.logging import configure_logging
from dataactcore.models.domainModels import DUNS
from dataactcore.utils.duns import (parse_duns_file, update_duns, parse_exec_comp_file, update_missing_parent_names,
                                    backfill_uei, request_sam_csv_api, is_nonexistent_file_error)
from dataactvalidator.health_check import create_app

logger = logging.getLogger(__name__)

SAM_FILE_FORMAT = 'SAM_{data_type}_UTF-8_{period}{version}_%Y%m%d.ZIP'
DATA_TYPES = {
    'DUNS': 'FOUO',
    'Executive Compensation': 'EXECCOMP'
}
PERIODS = ['MONTHLY', 'DAILY']
VERSIONS = {
    'v1': '',  # V1 files simply exclude the version
    'v2': '_V2'
}
S3_ARCHIVE = CONFIG_BROKER['sam']['duns']['csv_archive_bucket']
S3_ARCHIVE_PATH = '{data_type}/{version}/{file_name}'

def load_from_sam(data_type, sess, historic, local=None, metrics=None, reload_date=None):
    """ Process the script arguments to figure out which files to process in which order

        Args:
            data_type: data type to load (DUNS or executive compensation)
            sess: the database connection
            historic: whether to load in monthly file and daily files after, or just the latest daily files
            local: path to local directory to process, if None, it will go though the remote SAM service
            metrics: dictionary representing metrics data for the load
            reload_date: specific date to force reload from
    """
    if not metrics:
        metrics = {}

    # Figure out what files we have available based on our local or remote setup
    if local:
        local_files = os.listdir(local)
        monthly_v1_files = sorted([monthly_file for monthly_file in local_files
                                   if re.match('SAM_{}_UTF-8_MONTHLY_\d+\.ZIP'.format(DATA_TYPES[data_type]),
                                               monthly_file.upper())])
        monthly_v2_files = sorted([monthly_file for monthly_file in local_files
                                   if re.match('SAM_{}_UTF-8_MONTHLY_V2_\d+\.ZIP'.format(DATA_TYPES[data_type]),
                                               monthly_file.upper())])
        daily_v1_files = sorted([daily_file for daily_file in local_files
                                 if re.match('SAM_{}_UTF-8_DAILY_\d+\.ZIP'.format(DATA_TYPES[data_type]),
                                             daily_file.upper())])
        daily_v2_files = sorted([daily_file for daily_file in local_files
                                 if re.match('SAM_{}_UTF-8_DAILY_V2_\d+\.ZIP'.format(DATA_TYPES[data_type]),
                                             daily_file.upper())])
    else:
        # TODO: the SAM API currently doesn't list available files and doesnt include historic ones,
        #       so we're pulling files from the CSV_ARCHIVE_BUCKET bucket up and then use the API.
        #       Rework this if SAM includes these historic files in the API and list what files are available
        monthly_v1_files = list_s3_archive_files(data_type, 'MONTHLY', 'v1')
        monthly_v2_files = list_s3_archive_files(data_type, 'MONTHLY', 'v2')
        daily_v1_files = list_s3_archive_files(data_type, 'DAILY', 'v1')
        daily_v2_files = list_s3_archive_files(data_type, 'DAILY', 'v2')

    # Extracting the dates from these to figure out which files to process where
    # For both monthly and daily files, we only want to process v1 files until the equivalent v2 files are available
    monthly_v1_dates = extract_dates_from_list(monthly_v1_files, data_type, 'MONTHLY', 'v1')
    monthly_v2_dates = extract_dates_from_list(monthly_v2_files, data_type, 'MONTHLY', 'v2')
    monthly_v1_dates = [monthly_v1_date for monthly_v1_date in monthly_v1_dates
                        if monthly_v1_date not in monthly_v2_dates]
    if historic:
        earliest_date = sorted(monthly_v1_dates + monthly_v2_dates)[0]

    daily_v1_dates = extract_dates_from_list(daily_v1_files, data_type, 'DAILY', 'v1')
    daily_v2_dates = extract_dates_from_list(daily_v2_files, data_type, 'DAILY', 'v2')
    daily_v1_dates = [daily_v1_dates for daily_v1_dates in daily_v1_dates
                      if daily_v1_dates not in daily_v2_dates]
    latest_date = sorted(daily_v1_dates + daily_v2_dates)[-1]

    # For any dates after the latest date we have in the archive, use the API
    daily_v2_api_dates = [latest_date + datetime.timedelta(days=i)
                          for i in range(1, (datetime.date.today() - latest_date).days + 1)]

    # determine which daily files to load in by setting the start load date
    if historic:
        load_date = earliest_date
    elif reload_date:
        # a bit redundant but also date validation
        load_date = datetime.datetime.strptime(reload_date, '%Y-%m-%d').date()
    else:
        sam_field = DUNS.last_sam_mod_date if data_type == 'DUNS' else DUNS.last_exec_comp_mod_date
        load_date = sess.query(sam_field).filter(sam_field.isnot(None)).order_by(sam_field.desc()).first()
        if not load_date:
            field = 'sam' if data_type == 'duns' else 'executive compenstation'
            raise Exception('No last {} mod date found in DUNS table. Please run historic loader first.'.format(field))
        load_date = load_date[0]

    # only load in the daily files after the load date
    daily_v1_dates = list(filter(lambda daily_date: daily_date >= load_date, daily_v1_dates))
    daily_v2_dates = list(filter(lambda daily_date: daily_date >= load_date, daily_v2_dates))
    daily_v2_api_dates = list(filter(lambda daily_date: daily_date >= load_date, daily_v2_api_dates))

    if historic:
        # load in the earliest monthly file and all daily files after
        version = 'v1' if earliest_date in monthly_v1_dates else 'v2'
        process_sam_file(data_type, 'MONTHLY', version, earliest_date, sess, local=local, metrics=metrics)
    for daily_v1_date in daily_v1_dates:
        process_sam_file(data_type, 'DAILY', 'v1', daily_v1_date, sess, local=local, metrics=metrics)
    for daily_v2_date in daily_v2_dates:
        process_sam_file(data_type, 'DAILY', 'v2', daily_v2_date, sess, local=local, metrics=metrics)
    if not local:
        for daily_api_v2_date in daily_v2_api_dates:
            try:
                process_sam_file(data_type, 'DAILY', 'v2', daily_api_v2_date, sess, local=local, api=True,
                                 metrics=metrics)
            except requests.exceptions.HTTPError as e:
                if is_nonexistent_file_error(e):
                    logger.warning('No file found for {}, continuing'.format(daily_api_v2_date))
                    continue
                else:
                    logger.exception(e.response.content.decode('utf-8'))
                    raise e

    if data_type == 'DUNS':
        updated_date = datetime.date.today()
        metrics['parent_rows_updated'] = update_missing_parent_names(sess, updated_date=updated_date)
        metrics['parent_update_date'] = str(updated_date)

        if historic:
            backfill_uei(sess, DUNS)

def extract_dates_from_list(sam_files, data_type, period, version):
    """ Given a list of SAM files, extract the dates the files refer to

        Args:
            sam_files: list of sam file names to extract dates from
            data_type: data type to load (DUNS or executive compensation)
            period: monthly or daily
            version: v1 or v2

        Returns:
            sorted list of dates corresponding to the files
    """
    sam_filename_format = SAM_FILE_FORMAT.format(data_type=DATA_TYPES[data_type], period=period,
                                                 version=VERSIONS[version])
    return sorted([datetime.datetime.strptime(sam_file, sam_filename_format).date() for sam_file in sam_files])

def list_s3_archive_files(data_type, period, version):
    """ Given the requested fields, provide a list of available files from the remote S3 archive

        Args:
            data_type: data type to load (DUNS or executive compensation)
            period: monthly or daily
            version: v1 or v2

        Returns:
            list of available files in the S3 archive
    """
    s3_resource = boto3.resource('s3', region_name='us-gov-west-1')
    archive_bucket = s3_resource.Bucket(S3_ARCHIVE)
    file_name = SAM_FILE_FORMAT[:30].format(data_type=DATA_TYPES[data_type], period=period)
    prefix = S3_ARCHIVE_PATH.format(data_type=data_type, version=version, file_name=file_name)
    return [os.path.basename(object.key) for object in archive_bucket.objects.filter(Prefix=prefix)]

def download_sam_file(root_dir, file_name, api=False):
    """ Downloads the requested DUNS file to root_dir

        Args:
            root_dir: the folder containing the DUNS file
            file_name: the name of the SAM file
            api: whether to use the SAM CSV API or not

        Raises:
            requests.exceptions.HTTPError if the SAM HTTP API doesnt have the file requested
    """
    logger.info('Pulling {} via {}'.format(file_name, 'API' if api else 'archive'))
    if api:
        request_sam_csv_api(root_dir, file_name)
    else:
        s3_client = boto3.client('s3', region_name='us-gov-west-1')
        reverse_map = {v: k for k, v in DATA_TYPES.items()}
        data_type = reverse_map[file_name.split('_')[1]]
        version = 'v2' if 'V2' in file_name else 'v1'
        key = S3_ARCHIVE_PATH.format(data_type=data_type, version=version, file_name=file_name)
        s3_client.download_file(S3_ARCHIVE, key, os.path.join(root_dir, file_name))

def process_sam_file(data_type, period, version, date, sess, local=None, api=False, metrics=None):
    """ Process the SAM file found locally or remotely

        Args:
            data_type: data type to load (DUNS or executive compensation)
            period: monthly or daily
            version: v1 or v2
            sess: the database connection
            local: path to local directory to process, if None, it will go though the remote SAM service
            api: whether to use the SAM CSV API or not
            metrics: dictionary representing metrics data for the load

        Raises:
            requests.exceptions.HTTPError if the SAM HTTP API doesnt have the file requested
    """
    if not metrics:
        metrics = {}

    root_dir = local if local else tempfile.gettempdir()
    file_name_format = SAM_FILE_FORMAT.format(data_type=DATA_TYPES[data_type], period=period, version=VERSIONS[version])
    file_name = date.strftime(file_name_format)
    if not local:
        download_sam_file(root_dir, file_name, api=api)

    file_path = os.path.join(root_dir, file_name)
    if data_type == 'DUNS':
        add_update_data, delete_data = parse_duns_file(file_path, metrics=metrics)
        if add_update_data is not None:
            update_duns(sess, add_update_data, metrics=metrics)
        if delete_data is not None:
            update_duns(sess, delete_data, metrics=metrics, deletes=True)
    else:
        exec_comp_data = parse_exec_comp_file(file_path, metrics=metrics)
        update_duns(sess, exec_comp_data, metrics=metrics)
    if not local:
        os.remove(file_path)

def get_parser():
    """ Generates list of command-line arguments

        Returns:
            argument parser to be used for commandline
    """
    parser = argparse.ArgumentParser(description='Get data from SAM and update duns/exec comp tables')
    parser.add_argument("-t", "--data_type", choices=['duns', 'exec_comp', 'both'], default='both',
                        help='Select data type to load')
    scope = parser.add_mutually_exclusive_group(required=True)
    scope.add_argument("-a", "--historic", action="store_true", help='Reload from the first monthly file on')
    scope.add_argument("-u", "--update", action="store_true", help='Load daily files since latest last_sam_mod_date')
    environ = parser.add_mutually_exclusive_group(required=True)
    environ.add_argument("-l", "--local", type=str, default=None, help='Local directory to work from')
    environ.add_argument("-r", "--remote", action="store_true", help='Work from a remote directory (SAM)')
    parser.add_argument("-f", "--reload_date", type=str, default=None, help='Force update from a specific date'
                                                                            ' (YYYY-MM-DD)')
    return parser

if __name__ == '__main__':
    now = datetime.datetime.now()

    configure_logging()
    parser = get_parser()
    args = parser.parse_args()

    data_type = args.data_type
    historic = args.historic
    update = args.update
    local = args.local
    reload_date = args.reload_date

    metrics = {
        'script_name': 'load_duns_exec_comp.py',
        'start_time': str(now),
        'files_processed': [],
        'records_received': 0,
        'records_processed': 0,
        'adds_received': 0,
        'updates_received': 0,
        'deletes_received': 0,
        'added_duns': [],
        'updated_duns': [],
        'records_added': 0,
        'records_updated': 0,
        'parent_rows_updated': 0,
        'parent_update_date': None
    }

    with create_app().app_context():
        sess = GlobalDB.db().session
        if data_type in ('duns', 'both'):
            load_from_sam('DUNS', sess, historic, local, metrics=metrics, reload_date=reload_date)
        if data_type in ('exec_comp', 'both'):
            load_from_sam('Executive Compensation', sess, historic, local, metrics=metrics, reload_date=reload_date)
        sess.close()

    metrics['records_added'] = len(set(metrics['added_duns']))
    metrics['records_updated'] = len(set(metrics['updated_duns']) - set(metrics['added_duns']))
    del metrics['added_duns']
    del metrics['updated_duns']

    logger.info('Added {} records and updated {} records'.format(metrics['records_added'], metrics['records_updated']))

    metrics['duration'] = str(datetime.datetime.now() - now)
    with open('load_duns_exec_comp_metrics.json', 'w+') as metrics_file:
        json.dump(metrics, metrics_file)
<commit_msg>Removing backfill_uei script at the end of historical load as it's not entirely useful<commit_after>import argparse
import datetime
import logging
import os
import re
import json
import tempfile
import boto3
import requests

from dataactcore.config import CONFIG_BROKER
from dataactcore.interfaces.db import GlobalDB
from dataactcore.logging import configure_logging
from dataactcore.models.domainModels import DUNS
from dataactcore.utils.duns import (parse_duns_file, update_duns, parse_exec_comp_file, update_missing_parent_names,
                                    request_sam_csv_api, is_nonexistent_file_error)
from dataactvalidator.health_check import create_app

logger = logging.getLogger(__name__)

SAM_FILE_FORMAT = 'SAM_{data_type}_UTF-8_{period}{version}_%Y%m%d.ZIP'
DATA_TYPES = {
    'DUNS': 'FOUO',
    'Executive Compensation': 'EXECCOMP'
}
PERIODS = ['MONTHLY', 'DAILY']
VERSIONS = {
    'v1': '',  # V1 files simply exclude the version
    'v2': '_V2'
}
S3_ARCHIVE = CONFIG_BROKER['sam']['duns']['csv_archive_bucket']
S3_ARCHIVE_PATH = '{data_type}/{version}/{file_name}'

def load_from_sam(data_type, sess, historic, local=None, metrics=None, reload_date=None):
    """ Process the script arguments to figure out which files to process in which order

        Args:
            data_type: data type to load (DUNS or executive compensation)
            sess: the database connection
            historic: whether to load in monthly file and daily files after, or just the latest daily files
            local: path to local directory to process, if None, it will go though the remote SAM service
            metrics: dictionary representing metrics data for the load
            reload_date: specific date to force reload from
    """
    if not metrics:
        metrics = {}

    # Figure out what files we have available based on our local or remote setup
    if local:
        local_files = os.listdir(local)
        monthly_v1_files = sorted([monthly_file for monthly_file in local_files
                                   if re.match('SAM_{}_UTF-8_MONTHLY_\d+\.ZIP'.format(DATA_TYPES[data_type]),
                                               monthly_file.upper())])
        monthly_v2_files = sorted([monthly_file for monthly_file in local_files
                                   if re.match('SAM_{}_UTF-8_MONTHLY_V2_\d+\.ZIP'.format(DATA_TYPES[data_type]),
                                               monthly_file.upper())])
        daily_v1_files = sorted([daily_file for daily_file in local_files
                                 if re.match('SAM_{}_UTF-8_DAILY_\d+\.ZIP'.format(DATA_TYPES[data_type]),
                                             daily_file.upper())])
        daily_v2_files = sorted([daily_file for daily_file in local_files
                                 if re.match('SAM_{}_UTF-8_DAILY_V2_\d+\.ZIP'.format(DATA_TYPES[data_type]),
                                             daily_file.upper())])
    else:
        # TODO: the SAM API currently doesn't list available files and doesnt include historic ones,
        #       so we're pulling files from the CSV_ARCHIVE_BUCKET bucket up and then use the API.
        #       Rework this if SAM includes these historic files in the API and list what files are available
        monthly_v1_files = list_s3_archive_files(data_type, 'MONTHLY', 'v1')
        monthly_v2_files = list_s3_archive_files(data_type, 'MONTHLY', 'v2')
        daily_v1_files = list_s3_archive_files(data_type, 'DAILY', 'v1')
        daily_v2_files = list_s3_archive_files(data_type, 'DAILY', 'v2')

    # Extracting the dates from these to figure out which files to process where
    # For both monthly and daily files, we only want to process v1 files until the equivalent v2 files are available
    monthly_v1_dates = extract_dates_from_list(monthly_v1_files, data_type, 'MONTHLY', 'v1')
    monthly_v2_dates = extract_dates_from_list(monthly_v2_files, data_type, 'MONTHLY', 'v2')
    monthly_v1_dates = [monthly_v1_date for monthly_v1_date in monthly_v1_dates
                        if monthly_v1_date not in monthly_v2_dates]
    if historic:
        earliest_date = sorted(monthly_v1_dates + monthly_v2_dates)[0]

    daily_v1_dates = extract_dates_from_list(daily_v1_files, data_type, 'DAILY', 'v1')
    daily_v2_dates = extract_dates_from_list(daily_v2_files, data_type, 'DAILY', 'v2')
    daily_v1_dates = [daily_v1_dates for daily_v1_dates in daily_v1_dates
                      if daily_v1_dates not in daily_v2_dates]
    latest_date = sorted(daily_v1_dates + daily_v2_dates)[-1]

    # For any dates after the latest date we have in the archive, use the API
    daily_v2_api_dates = [latest_date + datetime.timedelta(days=i)
                          for i in range(1, (datetime.date.today() - latest_date).days + 1)]

    # determine which daily files to load in by setting the start load date
    if historic:
        load_date = earliest_date
    elif reload_date:
        # a bit redundant but also date validation
        load_date = datetime.datetime.strptime(reload_date, '%Y-%m-%d').date()
    else:
        sam_field = DUNS.last_sam_mod_date if data_type == 'DUNS' else DUNS.last_exec_comp_mod_date
        load_date = sess.query(sam_field).filter(sam_field.isnot(None)).order_by(sam_field.desc()).first()
        if not load_date:
            field = 'sam' if data_type == 'duns' else 'executive compenstation'
            raise Exception('No last {} mod date found in DUNS table. Please run historic loader first.'.format(field))
        load_date = load_date[0]

    # only load in the daily files after the load date
    daily_v1_dates = list(filter(lambda daily_date: daily_date >= load_date, daily_v1_dates))
    daily_v2_dates = list(filter(lambda daily_date: daily_date >= load_date, daily_v2_dates))
    daily_v2_api_dates = list(filter(lambda daily_date: daily_date >= load_date, daily_v2_api_dates))

    if historic:
        # load in the earliest monthly file and all daily files after
        version = 'v1' if earliest_date in monthly_v1_dates else 'v2'
        process_sam_file(data_type, 'MONTHLY', version, earliest_date, sess, local=local, metrics=metrics)
    for daily_v1_date in daily_v1_dates:
        process_sam_file(data_type, 'DAILY', 'v1', daily_v1_date, sess, local=local, metrics=metrics)
    for daily_v2_date in daily_v2_dates:
        process_sam_file(data_type, 'DAILY', 'v2', daily_v2_date, sess, local=local, metrics=metrics)
    if not local:
        for daily_api_v2_date in daily_v2_api_dates:
            try:
                process_sam_file(data_type, 'DAILY', 'v2', daily_api_v2_date, sess, local=local, api=True,
                                 metrics=metrics)
            except requests.exceptions.HTTPError as e:
                if is_nonexistent_file_error(e):
                    logger.warning('No file found for {}, continuing'.format(daily_api_v2_date))
                    continue
                else:
                    logger.exception(e.response.content.decode('utf-8'))
                    raise e

    if data_type == 'DUNS':
        updated_date = datetime.date.today()
        metrics['parent_rows_updated'] = update_missing_parent_names(sess, updated_date=updated_date)
        metrics['parent_update_date'] = str(updated_date)

        if historic:
            logger.info('Despite the historical load being done, the UEI will most likely be out of date. '
                        'Please manually update using the UEI crosswalk file and SQL.')

def extract_dates_from_list(sam_files, data_type, period, version):
    """ Given a list of SAM files, extract the dates the files refer to

        Args:
            sam_files: list of sam file names to extract dates from
            data_type: data type to load (DUNS or executive compensation)
            period: monthly or daily
            version: v1 or v2

        Returns:
            sorted list of dates corresponding to the files
    """
    sam_filename_format = SAM_FILE_FORMAT.format(data_type=DATA_TYPES[data_type], period=period,
                                                 version=VERSIONS[version])
    return sorted([datetime.datetime.strptime(sam_file, sam_filename_format).date() for sam_file in sam_files])

def list_s3_archive_files(data_type, period, version):
    """ Given the requested fields, provide a list of available files from the remote S3 archive

        Args:
            data_type: data type to load (DUNS or executive compensation)
            period: monthly or daily
            version: v1 or v2

        Returns:
            list of available files in the S3 archive
    """
    s3_resource = boto3.resource('s3', region_name='us-gov-west-1')
    archive_bucket = s3_resource.Bucket(S3_ARCHIVE)
    file_name = SAM_FILE_FORMAT[:30].format(data_type=DATA_TYPES[data_type], period=period)
    prefix = S3_ARCHIVE_PATH.format(data_type=data_type, version=version, file_name=file_name)
    return [os.path.basename(object.key) for object in archive_bucket.objects.filter(Prefix=prefix)]

def download_sam_file(root_dir, file_name, api=False):
    """ Downloads the requested DUNS file to root_dir

        Args:
            root_dir: the folder containing the DUNS file
            file_name: the name of the SAM file
            api: whether to use the SAM CSV API or not

        Raises:
            requests.exceptions.HTTPError if the SAM HTTP API doesnt have the file requested
    """
    logger.info('Pulling {} via {}'.format(file_name, 'API' if api else 'archive'))
    if api:
        request_sam_csv_api(root_dir, file_name)
    else:
        s3_client = boto3.client('s3', region_name='us-gov-west-1')
        reverse_map = {v: k for k, v in DATA_TYPES.items()}
        data_type = reverse_map[file_name.split('_')[1]]
        version = 'v2' if 'V2' in file_name else 'v1'
        key = S3_ARCHIVE_PATH.format(data_type=data_type, version=version, file_name=file_name)
        s3_client.download_file(S3_ARCHIVE, key, os.path.join(root_dir, file_name))

def process_sam_file(data_type, period, version, date, sess, local=None, api=False, metrics=None):
    """ Process the SAM file found locally or remotely

        Args:
            data_type: data type to load (DUNS or executive compensation)
            period: monthly or daily
            version: v1 or v2
            sess: the database connection
            local: path to local directory to process, if None, it will go though the remote SAM service
            api: whether to use the SAM CSV API or not
            metrics: dictionary representing metrics data for the load

        Raises:
            requests.exceptions.HTTPError if the SAM HTTP API doesnt have the file requested
    """
    if not metrics:
        metrics = {}

    root_dir = local if local else tempfile.gettempdir()
    file_name_format = SAM_FILE_FORMAT.format(data_type=DATA_TYPES[data_type], period=period, version=VERSIONS[version])
    file_name = date.strftime(file_name_format)
    if not local:
        download_sam_file(root_dir, file_name, api=api)

    file_path = os.path.join(root_dir, file_name)
    if data_type == 'DUNS':
        add_update_data, delete_data = parse_duns_file(file_path, metrics=metrics)
        if add_update_data is not None:
            update_duns(sess, add_update_data, metrics=metrics)
        if delete_data is not None:
            update_duns(sess, delete_data, metrics=metrics, deletes=True)
    else:
        exec_comp_data = parse_exec_comp_file(file_path, metrics=metrics)
        update_duns(sess, exec_comp_data, metrics=metrics)
    if not local:
        os.remove(file_path)

def get_parser():
    """ Generates list of command-line arguments

        Returns:
            argument parser to be used for commandline
    """
    parser = argparse.ArgumentParser(description='Get data from SAM and update duns/exec comp tables')
    parser.add_argument("-t", "--data_type", choices=['duns', 'exec_comp', 'both'], default='both',
                        help='Select data type to load')
    scope = parser.add_mutually_exclusive_group(required=True)
    scope.add_argument("-a", "--historic", action="store_true", help='Reload from the first monthly file on')
    scope.add_argument("-u", "--update", action="store_true", help='Load daily files since latest last_sam_mod_date')
    environ = parser.add_mutually_exclusive_group(required=True)
    environ.add_argument("-l", "--local", type=str, default=None, help='Local directory to work from')
    environ.add_argument("-r", "--remote", action="store_true", help='Work from a remote directory (SAM)')
    parser.add_argument("-f", "--reload_date", type=str, default=None, help='Force update from a specific date'
                                                                            ' (YYYY-MM-DD)')
    return parser

if __name__ == '__main__':
    now = datetime.datetime.now()

    configure_logging()
    parser = get_parser()
    args = parser.parse_args()

    data_type = args.data_type
    historic = args.historic
    update = args.update
    local = args.local
    reload_date = args.reload_date

    metrics = {
        'script_name': 'load_duns_exec_comp.py',
        'start_time': str(now),
        'files_processed': [],
        'records_received': 0,
        'records_processed': 0,
        'adds_received': 0,
        'updates_received': 0,
        'deletes_received': 0,
        'added_duns': [],
        'updated_duns': [],
        'records_added': 0,
        'records_updated': 0,
        'parent_rows_updated': 0,
        'parent_update_date': None
    }

    with create_app().app_context():
        sess = GlobalDB.db().session
        if data_type in ('duns', 'both'):
            load_from_sam('DUNS', sess, historic, local, metrics=metrics, reload_date=reload_date)
        if data_type in ('exec_comp', 'both'):
            load_from_sam('Executive Compensation', sess, historic, local, metrics=metrics, reload_date=reload_date)
        sess.close()

    metrics['records_added'] = len(set(metrics['added_duns']))
    metrics['records_updated'] = len(set(metrics['updated_duns']) - set(metrics['added_duns']))
    del metrics['added_duns']
    del metrics['updated_duns']

    logger.info('Added {} records and updated {} records'.format(metrics['records_added'], metrics['records_updated']))

    metrics['duration'] = str(datetime.datetime.now() - now)
    with open('load_duns_exec_comp_metrics.json', 'w+') as metrics_file:
        json.dump(metrics, metrics_file)
<|endoftext|>


<commit_before>import collections
import json
import logging
import multiprocessing
import os
import signal
import sys
import textwrap
from concurrent.futures import FIRST_COMPLETED
from concurrent.futures import wait
from functools import partial
from time import time

import requests
import requests.adapters
from six.moves import queue

from datarobot_batch_scoring.consts import (SENTINEL,
                                            WriterQueueMsg, Batch,
                                            REPORT_INTERVAL,
                                            ProgressQueueMsg)
from datarobot_batch_scoring.reader import (fast_to_csv_chunk,
                                            slow_to_csv_chunk)
from datarobot_batch_scoring.utils import compress, get_rusage, Worker

try:
    from futures import ThreadPoolExecutor
except ImportError:
    from concurrent.futures import ThreadPoolExecutor

logger = logging.getLogger(__name__)

FakeResponse = collections.namedtuple('FakeResponse', 'status_code, text')

class Network(Worker):
    state_names = {
        b"-": "Initial",
        b"I": "Idle",
        b"e": "PreIdle",
        b"E": "PrePreIdle",
        b"R": "Doing Requests",
        b"F": "Pool is Full",
        b"W": "Waiting for Finish",
        b"D": "Done"
    }

    def __init__(self, concurrency, timeout, ui,
                 network_queue,
                 network_deque,
                 writer_queue,
                 progress_queue,
                 abort_flag,
                 network_status,
                 endpoint,
                 headers,
                 user,
                 api_token,
                 pred_name,
                 fast_mode,
                 max_batch_size,
                 compression):

        Worker.__init__(self, network_status)

        self.concurrency = concurrency
        self.timeout = timeout
        self.ui = ui or logger
        self.network_queue = network_queue
        self.network_deque = network_deque
        self.writer_queue = writer_queue
        self.progress_queue = progress_queue
        self.abort_flag = abort_flag
        self.endpoint = endpoint
        self.headers = headers
        self.user = user
        self.api_token = api_token
        self.pred_name = pred_name
        self.fast_mode = fast_mode
        self.max_batch_size = max_batch_size
        self.compression = compression

        self._timeout = timeout
        self.futures = []
        self.concurrency = concurrency

        self._executor = None
        self.session = None
        self.proc = None

        self.n_consumed = 0
        self.n_retried = 0
        self.n_requests = 0

    def send_warning_to_ctx(self, batch, message):
        self.ui.info('CTX WARNING batch_id {} , '
                     'message {}'.format(batch.id, message))
        self.writer_queue.put((WriterQueueMsg.CTX_WARNING, {
            "batch": batch,
            "error": message
        }))

    def send_error_to_ctx(self, batch, message):
        self.ui.info('CTX ERROR batch_id {} , '
                     'message {}'.format(batch.id, message))

        self.writer_queue.put((WriterQueueMsg.CTX_ERROR, {
            "batch": batch,
            "error": message
        }))

    def push_retry(self, batch):
        # we retry a batch - decrement retry counter
        batch = batch._replace(rty_cnt=batch.rty_cnt - 1)
        try:
            self.network_deque.put(batch, block=False)
        except queue.Full:
            msg = 'Dropping {} due to backfill queue full.'.format(
                batch)
            self.ui.error(msg)
            self.send_error_to_ctx(batch, msg)

    def _response_callback(self, r, batch=None, *args, **kw):
        try:
            if r.status_code == 200:
                pickleable_resp = {'elapsed': r.elapsed.total_seconds(),
                                   'text': r.text,
                                   'headers': r.headers}
                self.writer_queue.put((WriterQueueMsg.RESPONSE, {
                    "request": pickleable_resp,
                    "batch": batch
                }))
                return
            elif isinstance(r, FakeResponse):
                if r.status_code == 499:
                    msg = ('batch {} timed out, dropping; '
                           'we lost {} records'
                           ''.format(batch.id, len(batch.data)))
                    self.ui.error(msg)
                    self.send_error_to_ctx(batch, msg)
                    return

                self.ui.debug('Skipping processing response '
                              'because of FakeResponse')
            else:
                try:
                    self.ui.warning('batch {} failed with status code '
                                    '{} message: {}'
                                    ''.format(
                                         batch.id,
                                         r.status_code,
                                         json.loads(r.text)['message']))
                except ValueError:
                    self.ui.warning('batch {} failed with status code: {}'
                                    ''.format(batch.id, r.status_code))

                text = r.text
                msg = ('batch {} failed, status_code:{} '
                       'text:{}'.format(batch.id, r.status_code, text))
                self.ui.error(msg)
                self.send_warning_to_ctx(batch, msg)

            if batch.rty_cnt == 1:
                msg = ('batch {} exceeded retry limit; '
                       'we lost {} records'
                       ''.format(batch.id, len(batch.data)))
                self.ui.error(msg)
                self.send_error_to_ctx(batch, msg)
            else:
                self.ui.warning('retrying failed batch {}, attempts left: {}'
                                .format(batch.id, batch.rty_cnt - 1))
                self.push_retry(batch)

        except Exception as e:
            msg = 'batch {} - dropping due to: {}, {} records lost'.format(
                batch.id, e, batch.rows)
            self.ui.error(msg)
            self.send_error_to_ctx(batch, msg)

    def _request(self, request):

        prepared = self.session.prepare_request(request)
        try:
            self.session.send(prepared, timeout=self._timeout)
        except Exception as exc:
            code = 400
            if isinstance(exc, requests.exceptions.ReadTimeout):
                self.ui.warning(textwrap.dedent("""The server did not send any data
in the allotted amount of time.
You might want to decrease the "--n_concurrent" parameters
or
increase "--timeout" parameter.
"""))
                code = 499
            else:
                self.ui.debug('Exception {}: {}'.format(type(exc), exc))
                raise

            try:
                callback = request.kwargs['hooks']['response']
            except AttributeError:
                callback = request.hooks['response'][0]
            response = FakeResponse(code, 'No Response')
            callback(response)

    def get_batch(self, dry_run=False):
        while True:
            if self.abort_flag.value:
                self.exit_fast(None, None)
                break
            try:
                r = self.network_deque.get_nowait()
                self.ui.debug('Got batch from dequeu: {}'.format(r.id))
                self.n_retried += 1
                yield r
            except queue.Empty:
                try:
                    r = self.network_queue.get(timeout=1)
                    if r.id == SENTINEL.id:
                        break
                    self.n_consumed += 1
                    yield r
                except queue.Empty:
                    if dry_run and self.state in (b"-", b"R"):
                        self.state = b'E'
                    elif self.state == b"E":
                        self.state = b'e'
                    elif self.state == b"e":
                        self.state = b'I'
                        if dry_run:
                            break

                except OSError:
                    self.ui.error('OS Error')
                    break

    def split_batch(self, batch):
            if self.fast_mode:
                chunk_formatter = fast_to_csv_chunk
            else:
                chunk_formatter = slow_to_csv_chunk

            todo = [batch]
            while todo:
                batch = todo.pop(0)
                data = chunk_formatter(batch.data, batch.fieldnames)
                starting_size = sys.getsizeof(data)
                if starting_size < self.max_batch_size:
                    if self.compression:
                        data = compress(data)
                        self.ui.debug(
                            'batch {}-{} transmitting {} byte - space savings '
                            '{}%'.format(batch.id, batch.rows,
                                         sys.getsizeof(data),
                                         '%.2f' % float(1 -
                                                        (sys.getsizeof(data) /
                                                         starting_size))))
                    else:
                        self.ui.debug('batch {}-{} transmitting {} bytes'
                                      ''.format(batch.id, batch.rows,
                                                starting_size))

                    yield (batch, data)
                else:
                    if batch.rows < 2:
                        msg = ('batch {} is single row but bigger '
                               'than limit, skipping. We lost {} '
                               'records'.format(batch.id,
                                                len(batch.data)))
                        self.ui.error(msg)
                        self.send_error_to_ctx(batch, msg)
                        continue

                    msg = ('batch {}-{} is too long: {} bytes,'
                           ' splitting'.format(batch.id, batch.rows,
                                               len(data)))
                    self.ui.debug(msg)
                    self.send_warning_to_ctx(batch, msg)
                    split_point = int(batch.rows/2)

                    data1 = batch.data[:split_point]
                    batch1 = Batch(batch.id, split_point, batch.fieldnames,
                                   data1, batch.rty_cnt)
                    todo.append(batch1)

                    data2 = batch.data[split_point:]
                    batch2 = Batch(batch.id + split_point,
                                   batch.rows - split_point,
                                   batch.fieldnames, data2, batch.rty_cnt)
                    todo.append(batch2)
                    todo.sort()

    def request_cb(self, f):
        futures = [i for i in self.futures if not i.done()]
        self.ui.debug('request finished, pending futures: {}'
                      ''.format(len(futures)))

        if len(futures) == 0:
            self.state = b'E'

    def exit_fast(self, a, b):
        self.state = b'D'
        os._exit(1)

    def perform_requests(self, dry_run=False):
        signal.signal(signal.SIGINT, self.exit_fast)
        signal.signal(signal.SIGTERM, self.exit_fast)

        self.state = b'E'
        for q_batch in self.get_batch(dry_run):
            for (batch, data) in self.split_batch(q_batch):

                if dry_run:
                    if self.state != b"R":
                        self.state = b'R'
                    yield
                    continue

                hook = partial(self._response_callback, batch=batch)
                r = requests.Request(
                    method='POST',
                    url=self.endpoint,
                    headers=self.headers,
                    data=data,
                    auth=(self.user, self.api_token),
                    hooks={'response': hook})

                self.n_requests += 1

                while True:
                    self.futures = [i for i in self.futures if not i.done()]
                    if len(self.futures) < self.concurrency:
                        self.state = b'R'
                        f = self._executor.submit(self._request, r)
                        f.add_done_callback(self.request_cb)
                        self.futures.append(f)
                        break
                    else:
                        self.state = b'F'
                        wait(self.futures, return_when=FIRST_COMPLETED)
                yield
        #  wait for all batches to finish before returning
        self.state = b'W'
        while self.futures:
            f_len = len(self.futures)
            self.futures = [i for i in self.futures if not i.done()]
            if f_len != len(self.futures):
                self.ui.debug('Waiting for final requests to finish. '
                              'remaining requests: {}'
                              ''.format(len(self.futures)))
            wait(self.futures, return_when=FIRST_COMPLETED)
        self.state = b'D'
        yield True

    def __enter__(self):
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        if self.proc and self.proc.is_alive():
            self.proc.terminate()

    def run(self, dry_run=False):
        if dry_run:
            i = 0
            for _ in self.perform_requests(True):
                i += 1

            return i

        self._executor = ThreadPoolExecutor(self.concurrency)
        self.session = requests.Session()
        adapter = requests.adapters.HTTPAdapter(
            pool_connections=self.concurrency, pool_maxsize=self.concurrency)
        self.session.mount('http://', adapter)
        self.session.mount('https://', adapter)

        t0 = time()
        last_report = time()
        i = 0
        r = None
        for r in self.perform_requests():
            if r is not True:
                i += 1
                self.ui.info('{} responses sent | time elapsed {}s'
                             .format(i, time() - t0))

                if time() - last_report > REPORT_INTERVAL:
                    self.progress_queue.put((
                        ProgressQueueMsg.NETWORK_PROGRESS, {
                            "processed": self.n_requests,
                            "retried": self.n_retried,
                            "consumed": self.n_consumed,
                            "rusage": get_rusage(),
                        }))
                    last_report = time()

        self.progress_queue.put((ProgressQueueMsg.NETWORK_DONE, {
            "ret": r,
            "processed": self.n_requests,
            "retried": self.n_retried,
            "consumed": self.n_consumed,
            "rusage": get_rusage(),
        }))

    def go(self, dry_run=False):
        if dry_run:
            return self.run(dry_run)

        self.ui.set_next_UI_name('network')
        self.proc = \
            multiprocessing.Process(target=self.run,
                                    name='Netwrk_Proc')
        self.proc.start()
        return self.proc
<commit_msg>Remove the dry_run logic from the Network<commit_after>import collections
import json
import logging
import multiprocessing
import os
import signal
import sys
import textwrap
from concurrent.futures import FIRST_COMPLETED
from concurrent.futures import wait
from functools import partial
from time import time

import requests
import requests.adapters
from six.moves import queue

from datarobot_batch_scoring.consts import (SENTINEL,
                                            WriterQueueMsg, Batch,
                                            REPORT_INTERVAL,
                                            ProgressQueueMsg)
from datarobot_batch_scoring.reader import (fast_to_csv_chunk,
                                            slow_to_csv_chunk)
from datarobot_batch_scoring.utils import compress, get_rusage, Worker

try:
    from futures import ThreadPoolExecutor
except ImportError:
    from concurrent.futures import ThreadPoolExecutor

logger = logging.getLogger(__name__)

FakeResponse = collections.namedtuple('FakeResponse', 'status_code, text')

class Network(Worker):
    state_names = {
        b"-": "Initial",
        b"I": "Idle",
        b"e": "PreIdle",
        b"E": "PrePreIdle",
        b"R": "Doing Requests",
        b"F": "Pool is Full",
        b"W": "Waiting for Finish",
        b"D": "Done"
    }

    def __init__(self, concurrency, timeout, ui,
                 network_queue,
                 network_deque,
                 writer_queue,
                 progress_queue,
                 abort_flag,
                 network_status,
                 endpoint,
                 headers,
                 user,
                 api_token,
                 pred_name,
                 fast_mode,
                 max_batch_size,
                 compression):

        Worker.__init__(self, network_status)

        self.concurrency = concurrency
        self.timeout = timeout
        self.ui = ui or logger
        self.network_queue = network_queue
        self.network_deque = network_deque
        self.writer_queue = writer_queue
        self.progress_queue = progress_queue
        self.abort_flag = abort_flag
        self.endpoint = endpoint
        self.headers = headers
        self.user = user
        self.api_token = api_token
        self.pred_name = pred_name
        self.fast_mode = fast_mode
        self.max_batch_size = max_batch_size
        self.compression = compression

        self._timeout = timeout
        self.futures = []
        self.concurrency = concurrency

        self._executor = None
        self.session = None
        self.proc = None

        self.n_consumed = 0
        self.n_retried = 0
        self.n_requests = 0

    def send_warning_to_ctx(self, batch, message):
        self.ui.info('CTX WARNING batch_id {} , '
                     'message {}'.format(batch.id, message))
        self.writer_queue.put((WriterQueueMsg.CTX_WARNING, {
            "batch": batch,
            "error": message
        }))

    def send_error_to_ctx(self, batch, message):
        self.ui.info('CTX ERROR batch_id {} , '
                     'message {}'.format(batch.id, message))

        self.writer_queue.put((WriterQueueMsg.CTX_ERROR, {
            "batch": batch,
            "error": message
        }))

    def push_retry(self, batch):
        # we retry a batch - decrement retry counter
        batch = batch._replace(rty_cnt=batch.rty_cnt - 1)
        try:
            self.network_deque.put(batch, block=False)
        except queue.Full:
            msg = 'Dropping {} due to backfill queue full.'.format(
                batch)
            self.ui.error(msg)
            self.send_error_to_ctx(batch, msg)

    def _response_callback(self, r, batch=None, *args, **kw):
        try:
            if r.status_code == 200:
                pickleable_resp = {'elapsed': r.elapsed.total_seconds(),
                                   'text': r.text,
                                   'headers': r.headers}
                self.writer_queue.put((WriterQueueMsg.RESPONSE, {
                    "request": pickleable_resp,
                    "batch": batch
                }))
                return
            elif isinstance(r, FakeResponse):
                if r.status_code == 499:
                    msg = ('batch {} timed out, dropping; '
                           'we lost {} records'
                           ''.format(batch.id, len(batch.data)))
                    self.ui.error(msg)
                    self.send_error_to_ctx(batch, msg)
                    return

                self.ui.debug('Skipping processing response '
                              'because of FakeResponse')
            else:
                try:
                    self.ui.warning('batch {} failed with status code '
                                    '{} message: {}'
                                    ''.format(
                                         batch.id,
                                         r.status_code,
                                         json.loads(r.text)['message']))
                except ValueError:
                    self.ui.warning('batch {} failed with status code: {}'
                                    ''.format(batch.id, r.status_code))

                text = r.text
                msg = ('batch {} failed, status_code:{} '
                       'text:{}'.format(batch.id, r.status_code, text))
                self.ui.error(msg)
                self.send_warning_to_ctx(batch, msg)

            if batch.rty_cnt == 1:
                msg = ('batch {} exceeded retry limit; '
                       'we lost {} records'
                       ''.format(batch.id, len(batch.data)))
                self.ui.error(msg)
                self.send_error_to_ctx(batch, msg)
            else:
                self.ui.warning('retrying failed batch {}, attempts left: {}'
                                .format(batch.id, batch.rty_cnt - 1))
                self.push_retry(batch)

        except Exception as e:
            msg = 'batch {} - dropping due to: {}, {} records lost'.format(
                batch.id, e, batch.rows)
            self.ui.error(msg)
            self.send_error_to_ctx(batch, msg)

    def _request(self, request):

        prepared = self.session.prepare_request(request)
        try:
            self.session.send(prepared, timeout=self._timeout)
        except Exception as exc:
            code = 400
            if isinstance(exc, requests.exceptions.ReadTimeout):
                self.ui.warning(textwrap.dedent("""The server did not send any data
in the allotted amount of time.
You might want to decrease the "--n_concurrent" parameters
or
increase "--timeout" parameter.
"""))
                code = 499
            else:
                self.ui.debug('Exception {}: {}'.format(type(exc), exc))
                raise

            try:
                callback = request.kwargs['hooks']['response']
            except AttributeError:
                callback = request.hooks['response'][0]
            response = FakeResponse(code, 'No Response')
            callback(response)

    def get_batch(self):
        while True:
            if self.abort_flag.value:
                self.exit_fast(None, None)
                break
            try:
                r = self.network_deque.get_nowait()
                self.ui.debug('Got batch from dequeu: {}'.format(r.id))
                self.n_retried += 1
                yield r
            except queue.Empty:
                try:
                    r = self.network_queue.get(timeout=1)
                    if r.id == SENTINEL.id:
                        break
                    self.n_consumed += 1
                    yield r
                except queue.Empty:
                    if self.state == b"E":
                        self.state = b'e'
                    elif self.state == b"e":
                        self.state = b'I'

                except OSError:
                    self.ui.error('OS Error')
                    break

    def split_batch(self, batch):
            if self.fast_mode:
                chunk_formatter = fast_to_csv_chunk
            else:
                chunk_formatter = slow_to_csv_chunk

            todo = [batch]
            while todo:
                batch = todo.pop(0)
                data = chunk_formatter(batch.data, batch.fieldnames)
                starting_size = sys.getsizeof(data)
                if starting_size < self.max_batch_size:
                    if self.compression:
                        data = compress(data)
                        self.ui.debug(
                            'batch {}-{} transmitting {} byte - space savings '
                            '{}%'.format(batch.id, batch.rows,
                                         sys.getsizeof(data),
                                         '%.2f' % float(1 -
                                                        (sys.getsizeof(data) /
                                                         starting_size))))
                    else:
                        self.ui.debug('batch {}-{} transmitting {} bytes'
                                      ''.format(batch.id, batch.rows,
                                                starting_size))

                    yield (batch, data)
                else:
                    if batch.rows < 2:
                        msg = ('batch {} is single row but bigger '
                               'than limit, skipping. We lost {} '
                               'records'.format(batch.id,
                                                len(batch.data)))
                        self.ui.error(msg)
                        self.send_error_to_ctx(batch, msg)
                        continue

                    msg = ('batch {}-{} is too long: {} bytes,'
                           ' splitting'.format(batch.id, batch.rows,
                                               len(data)))
                    self.ui.debug(msg)
                    self.send_warning_to_ctx(batch, msg)
                    split_point = int(batch.rows/2)

                    data1 = batch.data[:split_point]
                    batch1 = Batch(batch.id, split_point, batch.fieldnames,
                                   data1, batch.rty_cnt)
                    todo.append(batch1)

                    data2 = batch.data[split_point:]
                    batch2 = Batch(batch.id + split_point,
                                   batch.rows - split_point,
                                   batch.fieldnames, data2, batch.rty_cnt)
                    todo.append(batch2)
                    todo.sort()

    def request_cb(self, f):
        futures = [i for i in self.futures if not i.done()]
        self.ui.debug('request finished, pending futures: {}'
                      ''.format(len(futures)))

        if len(futures) == 0:
            self.state = b'E'

    def exit_fast(self, a, b):
        self.state = b'D'
        os._exit(1)

    def perform_requests(self):
        signal.signal(signal.SIGINT, self.exit_fast)
        signal.signal(signal.SIGTERM, self.exit_fast)

        self.state = b'E'
        for q_batch in self.get_batch():
            for (batch, data) in self.split_batch(q_batch):

                hook = partial(self._response_callback, batch=batch)
                r = requests.Request(
                    method='POST',
                    url=self.endpoint,
                    headers=self.headers,
                    data=data,
                    auth=(self.user, self.api_token),
                    hooks={'response': hook})

                self.n_requests += 1

                while True:
                    self.futures = [i for i in self.futures if not i.done()]
                    if len(self.futures) < self.concurrency:
                        self.state = b'R'
                        f = self._executor.submit(self._request, r)
                        f.add_done_callback(self.request_cb)
                        self.futures.append(f)
                        break
                    else:
                        self.state = b'F'
                        wait(self.futures, return_when=FIRST_COMPLETED)
                yield
        #  wait for all batches to finish before returning
        self.state = b'W'
        while self.futures:
            f_len = len(self.futures)
            self.futures = [i for i in self.futures if not i.done()]
            if f_len != len(self.futures):
                self.ui.debug('Waiting for final requests to finish. '
                              'remaining requests: {}'
                              ''.format(len(self.futures)))
            wait(self.futures, return_when=FIRST_COMPLETED)
        self.state = b'D'
        yield True

    def __enter__(self):
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        if self.proc and self.proc.is_alive():
            self.proc.terminate()

    def run(self):
        self._executor = ThreadPoolExecutor(self.concurrency)
        self.session = requests.Session()
        adapter = requests.adapters.HTTPAdapter(
            pool_connections=self.concurrency, pool_maxsize=self.concurrency)
        self.session.mount('http://', adapter)
        self.session.mount('https://', adapter)

        t0 = time()
        last_report = time()
        i = 0
        r = None
        for r in self.perform_requests():
            if r is not True:
                i += 1
                self.ui.info('{} responses sent | time elapsed {}s'
                             .format(i, time() - t0))

                if time() - last_report > REPORT_INTERVAL:
                    self.progress_queue.put((
                        ProgressQueueMsg.NETWORK_PROGRESS, {
                            "processed": self.n_requests,
                            "retried": self.n_retried,
                            "consumed": self.n_consumed,
                            "rusage": get_rusage(),
                        }))
                    last_report = time()

        self.progress_queue.put((ProgressQueueMsg.NETWORK_DONE, {
            "ret": r,
            "processed": self.n_requests,
            "retried": self.n_retried,
            "consumed": self.n_consumed,
            "rusage": get_rusage(),
        }))

    def go(self):
        self.ui.set_next_UI_name('network')
        self.proc = \
            multiprocessing.Process(target=self.run,
                                    name='Netwrk_Proc')
        self.proc.start()
        return self.proc
<|endoftext|>