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

""" Python 3 implementation of signal_at_orf function of hwglabr package."""

import helpers
import sys
import os
import time
import pandas as pd
import numpy as np
from scipy.interpolate import InterpolatedUnivariateSpline
from colorama import Fore, Back, Style, init as colinit
colinit(autoreset=True)

def signal_at_orf(wiggle_data, gff_file, save_file=True):
    """
    Given wiggle data generated using the lab's ChIP-seq analysis pipelines (as a
    dictionary of pandas data frames; output of function read_wiggle), this function
    pulls out the ChIP signal over all ORFs in the yeast genome. It collects the
    signal over each ORF plus both flanking regions (1/2 the length of the ORF on
    each side) and scales them all to the same value (1000).

    This means that for two example genes with lengths of 500 bp and 2 kb, flanking
    regions of 250 bp and 1 kb, respectively, will be collected up and downstream.
    Both gene lengths will then be scaled to 500 and all four flanking regions to
    250.

    After scaling, fits a spline function passing through all signal points and uses
    it to output values of the signal at each integer position between 1 and 1000.

    Keyword arguments
    =================
    :param wiggle_data: dictionary of pandas DataFrames
            Input wiggle data (output of function read_wiggle; no default)
    :param gff_file: string
            Path to the gff file providing the ORF coordinates (no default)
    :param save_file: bool, optional
            Whether output should be written to a TSV file (in current working
            directory; defaults to True)
    :return: Either a pandas DataFrame or a TSV file of a table with four columns:
            - chr: chromosome number
            - position: nucleotide coordinate (in normalized total length of 1 kb)
            - signal: ChIP-seq signal at each position (1 to 1000)
            - gene: Systematic gene name
    """

    t0 = time.time()

    # GFF
    gff = helpers.read_gff(gff_file)
    # Wiggle data
    wiggle = helpers.read_wiggle(wiggle_data, use_pbar=False)
    # Check reference genome
    gff_gen = helpers.check_genome(gff.ix[0, 'seqname'])
    wiggle_gen = helpers.check_genome(next(iter(wiggle.keys())))

    if gff_gen == wiggle_gen:
        print('Identified reference genome: ', end="")
        print(Fore.RED + gff_gen)
    else:
        sys.exit('Error: reference genomes in gff file and wiggle data don\'t match.')

    # Some feedback on the provided gff
    features = ', '.join(gff.ix[:, 'feature'].unique())
    print(Fore.YELLOW + 'The following types of features are present in the gff data you provided:')
    print(Fore.YELLOW + '(they will all be included in the analysis)')
    print(Fore.RED + features)

    print()
    print('Collecting signal...')
    print('(Skip genes with missing coordinates and signal in wiggle data)')

    # Create dfs to collect final data for all chrs
    plus_final = pd.DataFrame()
    minus_final = pd.DataFrame()

    # Keep track of total and non-skipped genes
    number_genes = 0
    number_skipped_genes = 0

    # Loop through chrs
    for chrNum, chromData in wiggle.items():
        print(Style.BRIGHT + chrNum + ':')
        # -------------------------------------- Plus strand ------------------------------------- #
        # Create data frame to collect final data for all genes in chr strand
        plus_strand = pd.DataFrame()

        # Get all genes on "+" strand of current chromosome
        chrgff = gff.loc[(gff['seqname'] == chrNum) & (gff['strand'] == '+')]

        gene_count = 0

        # Loop through rows (use itertuples because it's faster than iterrows)
        for row in chrgff.itertuples():
            # Skip if gene coordinates not in ChIPseq data
            if (row.start not in chromData.loc[:, 'position'] or
                        row.end not in chromData.loc[:, 'position']):
                continue

            # Collect flanking regions scaled according to ratio gene length / 1 kb
            gene_leng = row.end - row.start
            start = row.start - (gene_leng // 2)
            end = row.end + (gene_leng // 2)
            full_leng = (end - start) + 1
            gene = row.attribute

            # Pull out signal
            gene_data = chromData.loc[(chromData['position'] >= start) & (chromData['position'] <= end)]

            # Skip if there are discontinuities in the data (missing position:value pairs)
            if gene_data.shape[0] != full_leng:
                continue

            # Normalize to segment length of 1000
            pd.options.mode.chained_assignment = None  # Disable warning about chained assignment
            gene_data['position'] = gene_data['position'] - start + 1
            gene_data['position'] = gene_data['position'] * (1000 / full_leng)

            # Genes of different sizes have different numbers of positions; small genes
            # (<1000bp) cannot produce signal values for all 1000 positions and will have gaps
            # This means that longer genes with more signal values per each position in the
            # sequence of 1000 positions will contribute more to the final output.
            # In order to avoid this, first build an interpolation model of the signal and then
            # use it to project the signal onto the first 1000 integers
            f = InterpolatedUnivariateSpline(gene_data['position'], gene_data['signal'])
            new_positions = np.int_(np.linspace(1, 1000, num=1000, endpoint=True))
            new_signals = f(new_positions)

            # Make data frame for this gene
            gene_data = pd.DataFrame({'chr': chrNum, 'position': new_positions, 'signal': new_signals, 'gene': gene})

            # To collect all genes
            plus_strand = plus_strand.append(gene_data)

            gene_count += 1

        print('... + strand: {0} genes (skipped {1})'.format(gene_count,
                                                             chrgff.shape[0] - gene_count))

        # Keep track of total and non-skipped genes, to print info at the end
        number_genes += chrgff.shape[0]
        number_skipped_genes += chrgff.shape[0] - gene_count

        # To collect all chrs
        plus_final = plus_final.append(plus_strand)

        # ------------------------------------- Minus strand ------------------------------------- #
        # Create data frame to collect final data for all genes in chr strand
        minus_strand = pd.DataFrame()

        # Get all genes on "+" strand of current chromosome
        chrgff = gff.loc[(gff['seqname'] == chrNum) & (gff['strand'] == '-')]

        gene_count = 0

        # Loop through rows (use itertuples because it's faster than iterrows)
        for row in chrgff.itertuples():
            # Skip if gene coordinates not in ChIPseq data
            if (row.start not in chromData.loc[:, 'position'] or
                        row.end not in chromData.loc[:, 'position']):
                continue

            # Collect flanking regions scaled according to ratio gene length / 1 kb
            gene_leng = row.end - row.start
            start = row.start - (gene_leng // 2)
            end = row.end + (gene_leng // 2)
            full_leng = (end - start) + 1
            gene = row.attribute

            # Pull out signal
            gene_data = chromData.loc[(chromData['position'] >= start) & (chromData['position'] <= end)]

            # Skip if there are discontinuities in the data (missing position:value pairs)
            if gene_data.shape[0] != full_leng:
                continue

            # Normalize to segment length of 1000
            pd.options.mode.chained_assignment = None  # Disable warning about chained assignment
            gene_data['position'] = gene_data['position'] - start + 1
            gene_data['position'] = gene_data['position'] * (1000 / full_leng)

            f = InterpolatedUnivariateSpline(gene_data['position'], gene_data['signal'])
            new_positions = np.int_(np.linspace(1, 1000, num=1000, endpoint=True))
            new_signals = f(new_positions)

            # Reverse the order of the position values
            new_positions = (1000 - new_positions) + 1

            # Make data frame for this gene
            gene_data = pd.DataFrame({'chr': chrNum,
                                     'position': new_positions,
                                     'signal': new_signals,
                                     'gene': gene})

            # To collect all genes
            minus_strand = minus_strand.append(gene_data)

            gene_count += 1

        print('... - strand: {0} genes (skipped {1})'.format(gene_count,
                                                             chrgff.shape[0] - gene_count))

        # Keep track of total and non-skipped genes, to print info at the end
        number_genes += chrgff.shape[0]
        number_skipped_genes += chrgff.shape[0] - gene_count

        # To collect all chrs
        minus_final = minus_final.append(minus_strand)

    # Merge '+' and '-' strand data
    merged_strands = plus_final.append(minus_final)

    # Sort by gene and position
    merged_strands = merged_strands.sort_values(['gene', 'position'])

    # Print info on total and non-skipped genes
    print()
    print('------')
    percent_skipped = np.around(number_skipped_genes * 100 / number_genes, decimals=1)
    print(Back.RED + 'Skipped {0} of a total of {1} genes ({2}%)'.format(number_skipped_genes,
                                                       number_genes, percent_skipped))
    print('------')

    if save_file:
        file_name = os.path.basename(os.path.normpath(wiggle_data)) + '_metaORF.tsv'
        print(Fore.YELLOW + 'Saving output to file:')
        print(file_name)
        merged_strands.to_csv(file_name, sep='\t', index=False)
        print()
        print(Fore.CYAN + 'Completed in ', end=" ")
        helpers.print_elapsed_time(t0)
    else:
        print()
        print(Fore.CYAN + 'Completed in ', end=" ")
        helpers.print_elapsed_time(t0)
        return merged_strands

def main():
    print()
    print(Style.BRIGHT + "----------------------------------------------------")
    print(Style.BRIGHT + "          MEAN ChIP-SEQ SIGNAL AT META ORF")
    print(Style.BRIGHT + "----------------------------------------------------")
    print()
    signal_at_orf(wiggle_data=args.wiggle_data, gff_file=args.gff_file, save_file=True)
    print(Style.BRIGHT + "----------------------------------------------------")
    print()

if __name__ == '__main__':
    import argparse
    parser = argparse.ArgumentParser(description="Calculate mean ChIP-seq signal at ORF")
    parser.add_argument('-w', '--wiggle_data', help=("path to a folder containing wiggle data"
                                                     "to get ChIP-seq signal from"), required=True)
    parser.add_argument('-g', '--gff_file', help='path to gff file', required=True)
    args = parser.parse_args()
    main()
<commit_msg>Minor changes to comments<commit_after>#!/usr/bin/env python

""" Python 3 implementation of signal_at_orf function of hwglabr package."""

import helpers
import sys
import os
import time
import pandas as pd
import numpy as np
from scipy.interpolate import InterpolatedUnivariateSpline
from colorama import Fore, Back, Style, init as colinit
colinit(autoreset=True)

def signal_at_orf(wiggle_data, gff_file, save_file=True):
    """
    Given wiggle data generated using the lab's ChIP-seq analysis pipelines (as a
    dictionary of pandas data frames; output of function read_wiggle), this function
    pulls out the ChIP signal over all ORFs in the yeast genome. It collects the
    signal over each ORF plus both flanking regions (1/2 the length of the ORF on
    each side) and scales them all to the same value (1000).

    This means that for two example genes with lengths of 500 bp and 2 kb, flanking
    regions of 250 bp and 1 kb, respectively, will be collected up and downstream.
    Both gene lengths will then be scaled to 500 and all four flanking regions to
    250.

    After scaling, fits a spline function passing through all signal points and uses
    it to output values of the signal at each integer position between 1 and 1000.

    Keyword arguments
    =================
    :param wiggle_data: dictionary of pandas DataFrames
            Input wiggle data (output of function read_wiggle; no default)
    :param gff_file: string
            Path to the gff file providing the ORF coordinates (no default)
    :param save_file: bool, optional
            Whether output should be written to a TSV file (in current working
            directory; defaults to True)
    :return: Either a pandas DataFrame or a TSV file of a table with four columns:
            - chr: chromosome number
            - position: nucleotide coordinate (in normalized total length of 1 kb)
            - signal: ChIP-seq signal at each position (1 to 1000)
            - gene: Systematic gene name
    """

    t0 = time.time()

    # GFF
    gff = helpers.read_gff(gff_file)
    # Wiggle data
    wiggle = helpers.read_wiggle(wiggle_data, use_pbar=False)
    # Check reference genome
    gff_gen = helpers.check_genome(gff.ix[0, 'seqname'])
    wiggle_gen = helpers.check_genome(next(iter(wiggle.keys())))

    if gff_gen == wiggle_gen:
        print('Identified reference genome: ', end="")
        print(Fore.RED + gff_gen)
    else:
        sys.exit('Error: reference genomes in gff file and wiggle data don\'t match.')

    # Some feedback on the provided gff
    features = ', '.join(gff.ix[:, 'feature'].unique())
    print(Fore.YELLOW + 'The following types of features are present in the gff data you provided:')
    print(Fore.YELLOW + '(they will all be included in the analysis)')
    print(Fore.RED + features)

    print()
    print('Collecting signal...')
    print('(Skip genes with missing coordinates and signal in wiggle data)')

    # Create dfs to collect final data for all chrs
    plus_final = pd.DataFrame()
    minus_final = pd.DataFrame()

    # Keep track of total and non-skipped genes
    number_genes = 0
    number_skipped_genes = 0

    # Loop through chrs
    for chrNum, chromData in wiggle.items():
        print(Style.BRIGHT + chrNum + ':')
        # -------------------------------------- Plus strand ------------------------------------- #
        # Create data frame to collect final data for all genes in chr strand
        plus_strand = pd.DataFrame()

        # Get all genes on "+" strand of current chromosome
        chrgff = gff.loc[(gff['seqname'] == chrNum) & (gff['strand'] == '+')]

        gene_count = 0

        # Loop through rows (use itertuples, it's faster than iterrows)
        for row in chrgff.itertuples():
            # Skip if gene coordinates not in ChIPseq data
            if (row.start not in chromData.loc[:, 'position'] or
                        row.end not in chromData.loc[:, 'position']):
                continue

            # Collect flanking regions scaled according to ratio gene length / 1 kb
            gene_leng = row.end - row.start
            start = row.start - (gene_leng // 2)
            end = row.end + (gene_leng // 2)
            full_leng = (end - start) + 1
            gene = row.attribute

            # Pull out signal
            gene_data = chromData.loc[(chromData['position'] >= start) & (chromData['position'] <= end)]

            # Skip if there are discontinuities in the data (missing position:value pairs)
            if gene_data.shape[0] != full_leng:
                continue

            # Normalize to segment length of 1000
            pd.options.mode.chained_assignment = None  # Disable warning about chained assignment
            gene_data['position'] = gene_data['position'] - start + 1
            gene_data['position'] = gene_data['position'] * (1000 / full_leng)

            # Genes of different sizes have different numbers of positions; small genes
            # (<1000bp) cannot produce signal values for all 1000 positions and will have gaps
            # This means that longer genes with more signal values per each position in the
            # sequence of 1000 positions will contribute more to the final output.
            # In order to avoid this, first build an interpolation model of the signal and then
            # use it to project the signal onto the first 1000 integers
            f = InterpolatedUnivariateSpline(gene_data['position'], gene_data['signal'])
            new_positions = np.int_(np.linspace(1, 1000, num=1000, endpoint=True))
            new_signals = f(new_positions)

            # Make data frame for this gene
            gene_data = pd.DataFrame({'chr': chrNum, 'position': new_positions, 'signal': new_signals, 'gene': gene})

            # To collect all genes
            plus_strand = plus_strand.append(gene_data)

            gene_count += 1

        print('... + strand: {0} genes (skipped {1})'.format(gene_count,
                                                             chrgff.shape[0] - gene_count))

        # Keep track of total and non-skipped genes, to print info at the end
        number_genes += chrgff.shape[0]
        number_skipped_genes += chrgff.shape[0] - gene_count

        # To collect all chrs
        plus_final = plus_final.append(plus_strand)

        # ------------------------------------- Minus strand ------------------------------------- #
        # Create data frame to collect final data for all genes in chr strand
        minus_strand = pd.DataFrame()

        # Get all genes on "+" strand of current chromosome
        chrgff = gff.loc[(gff['seqname'] == chrNum) & (gff['strand'] == '-')]

        gene_count = 0

        # Loop through rows (use itertuples because it's faster than iterrows)
        for row in chrgff.itertuples():
            # Skip if gene coordinates not in ChIPseq data
            if (row.start not in chromData.loc[:, 'position'] or
                        row.end not in chromData.loc[:, 'position']):
                continue

            # Collect flanking regions scaled according to ratio gene length / 1 kb
            gene_leng = row.end - row.start
            start = row.start - (gene_leng // 2)
            end = row.end + (gene_leng // 2)
            full_leng = (end - start) + 1
            gene = row.attribute

            # Pull out signal
            gene_data = chromData.loc[(chromData['position'] >= start) & (chromData['position'] <= end)]

            # Skip if there are discontinuities in the data (missing position:value pairs)
            if gene_data.shape[0] != full_leng:
                continue

            # Normalize to segment length of 1000
            pd.options.mode.chained_assignment = None  # Disable warning about chained assignment
            gene_data['position'] = gene_data['position'] - start + 1
            gene_data['position'] = gene_data['position'] * (1000 / full_leng)

            f = InterpolatedUnivariateSpline(gene_data['position'], gene_data['signal'])
            new_positions = np.int_(np.linspace(1, 1000, num=1000, endpoint=True))
            new_signals = f(new_positions)

            # Reverse the order of the position values (to join with Watson strand)
            new_positions = (1000 - new_positions) + 1

            # Make data frame for this gene
            gene_data = pd.DataFrame({'chr': chrNum,
                                     'position': new_positions,
                                     'signal': new_signals,
                                     'gene': gene})

            # To collect all genes
            minus_strand = minus_strand.append(gene_data)

            gene_count += 1

        print('... - strand: {0} genes (skipped {1})'.format(gene_count,
                                                             chrgff.shape[0] - gene_count))

        # Keep track of total and non-skipped genes, to print info at the end
        number_genes += chrgff.shape[0]
        number_skipped_genes += chrgff.shape[0] - gene_count

        # To collect all chrs
        minus_final = minus_final.append(minus_strand)

    # Merge '+' and '-' strand data
    merged_strands = plus_final.append(minus_final)

    # Sort by gene and position
    merged_strands = merged_strands.sort_values(['gene', 'position'])

    # Print info on total and non-skipped genes
    print()
    print('------')
    percent_skipped = np.around(number_skipped_genes * 100 / number_genes, decimals=1)
    print(Back.RED + 'Skipped {0} of a total of {1} genes ({2}%)'.format(number_skipped_genes,
                                                       number_genes, percent_skipped))
    print('------')

    if save_file:
        file_name = os.path.basename(os.path.normpath(wiggle_data)) + '_metaORF.tsv'
        print(Fore.YELLOW + 'Saving output to file:')
        print(file_name)
        merged_strands.to_csv(file_name, sep='\t', index=False)
        print()
        print(Fore.CYAN + 'Completed in ', end=" ")
        helpers.print_elapsed_time(t0)
    else:
        print()
        print(Fore.CYAN + 'Completed in ', end=" ")
        helpers.print_elapsed_time(t0)
        return merged_strands

def main():
    print()
    print(Style.BRIGHT + "----------------------------------------------------")
    print(Style.BRIGHT + "          MEAN ChIP-SEQ SIGNAL AT META ORF")
    print(Style.BRIGHT + "----------------------------------------------------")
    print()
    signal_at_orf(wiggle_data=args.wiggle_data, gff_file=args.gff_file, save_file=True)
    print(Style.BRIGHT + "----------------------------------------------------")
    print()

if __name__ == '__main__':
    import argparse
    parser = argparse.ArgumentParser(description="Calculate mean ChIP-seq signal at ORF")
    parser.add_argument('-w', '--wiggle_data', help=("path to a folder containing wiggle data"
                                                     "to get ChIP-seq signal from"), required=True)
    parser.add_argument('-g', '--gff_file', help='path to gff file', required=True)
    args = parser.parse_args()
    main()
<|endoftext|>


<commit_before># This Source Code Form is subject to the terms of the Mozilla Public
# License, v. 2.0. If a copy of the MPL was not distributed with this
# file, You can obtain one at http://mozilla.org/MPL/2.0/.
"""
Measurements app for collecting call latency and performance stats.
"""
import weakref
from switchy.marks import event_callback
from switchy import utils
from metrics import new_array

class Metrics(object):
    """Collect call oriented measurements

    Only an instance of this class can be loaded as switchy app
    """
    def __init__(self):
        self.log = utils.get_logger(__name__)

    def prepost(self, listener, array=None, pool=None):
        self.listener = listener
        # array can be overriden at app load time
        self._array = array if array else new_array()  # np default buffer
        self.pool = weakref.proxy(pool) if pool else self.listener

    @property
    def array(self):
        return self._array

    @event_callback('CHANNEL_HANGUP')
    def log_stats(self, sess, job):
        """Append measurement data inserting only once per call
        """
        call = sess.call
        # ensure this is NOT the last active session in the call
        if not call.sessions:
            # all other sessions in call have been hungup and this is
            # the last so all measurements should already have been collected
            return

        # TODO: eventually we should use a data type which allows the
        # separation of metrics between those that require two vs. one
        # leg to calculate?
        l = self.listener
        pool = self.pool

        peer = sess.call.get_peer(sess)
        if peer and l.sessions.get(peer.uuid, False):
            first = call.first
            last = call.last
            rollover = self._array.insert((
                first.create_time,  # invite time index
                last.create_time - first.create_time,
                last.answer_time - first.answer_time,
                first.answer_time - first.create_time,
                first.originate_time - job.launch_time if job else 0,
                first.originate_event_time - first.create_time if job else 0,
                pool.count_failed() if pool else 0,
                pool.count_sessions(),
            ))
            if rollover:
                self.log.warn('resetting metric buffer index!')
<commit_msg>Port metrics to use `Session.times` map<commit_after># This Source Code Form is subject to the terms of the Mozilla Public
# License, v. 2.0. If a copy of the MPL was not distributed with this
# file, You can obtain one at http://mozilla.org/MPL/2.0/.
"""
Measurements app for collecting call latency and performance stats.
"""
import weakref
import time
from switchy.marks import event_callback
from switchy import utils
from metrics import new_array

class Metrics(object):
    """Collect call oriented measurements

    Only an instance of this class can be loaded as switchy app
    """
    def __init__(self):
        self.log = utils.get_logger(__name__)

    def prepost(self, listener, array=None, pool=None):
        self.listener = listener
        # array can be overriden at app load time
        self._array = array if array else new_array()  # np default buffer
        self.pool = weakref.proxy(pool) if pool else self.listener

    @property
    def array(self):
        return self._array

    @event_callback('CHANNEL_ORIGINATE')
    def on_originate(self, sess):
        # store local time stamp for originate
        sess.times['originate'] = sess.time
        sess.times['req_originate'] = time.time()

    @event_callback('CHANNEL_HANGUP')
    def log_stats(self, sess, job):
        """Append measurement data inserting only once per call
        """
        call = sess.call
        # ensure this is NOT the last active session in the call
        if not call.sessions:
            # all other sessions in call have been hungup and this is
            # the last so all measurements should already have been collected
            return

        # TODO: eventually we should use a data type which allows the
        # separation of metrics between those that require two vs. one
        # leg to calculate?
        l = self.listener
        pool = self.pool

        peer = sess.call.get_peer(sess)
        if peer and l.sessions.get(peer.uuid, False):
            first = call.first.times
            last = call.last.times
            rollover = self._array.insert((
                first['create'],  # invite time index
                last['create'] - first['create'],
                last['answer'] - first['answer'],
                first['answer'] - first['create'],
                first['req_originate'] - job.launch_time if job else 0,
                first['originate'] - first['create'] if job else 0,
                pool.count_failed() if pool else 0,
                pool.count_sessions(),
            ))
            if rollover:
                self.log.warn('resetting metric buffer index!')
<|endoftext|>


<commit_before>from luigi import Task, WrapperTask, Parameter, IntParameter, ListParameter
from tasks.us.census.tiger import GeoNamesTable, ShorelineClip, SUMLEVELS
from tasks.base_tasks import TempTableTask
from tasks.meta import current_session
from lib.logger import get_logger

LOGGER = get_logger(__name__)

def _union(tables, output):
    unions = ['SELECT * FROM {table}'.format(table=table.qualified_tablename)
              for table in tables]
    return 'CREATE TABLE {output} AS {unions}'.format(
        output=output,
        unions=' UNION '.join(unions))

class USHierarchy(Task):
    year = IntParameter()

    def requires(self):
        levels = self._levels()
        LOGGER.info('Levels: {}'.format(levels))
        return {
            'info': USHierarchyInfoUnion(year=self.year, levels=levels),
            'rel': USHierarchyChildParentsUnion(year=self.year, levels=levels)
        }

    def _levels(self):
        level_infos = [(level, info) for level, info in SUMLEVELS.items()]

        sorted_level_infos = reversed(sorted(level_infos,
                                             key=lambda level_info:
                                             level_info[1][
                                                 'weight']))
        return [level_info[0] for level_info in sorted_level_infos]

    def run(self):
        session = current_session()
        session.execute('ALTER TABLE {rel_table} ADD '
                        'CONSTRAINT ushierarchy_fk_child '
                        'FOREIGN KEY (child_id, child_level) '
                        'REFERENCES {info_table} (geoid, level) '.format(
            rel_table=self.input()['rel'].qualified_tablename,
            info_table=self.input()['info'].qualified_tablename))
        session.execute('ALTER TABLE {rel_table} ADD '
                        'CONSTRAINT ushierarchy_fk_parent '
                        'FOREIGN KEY (parent_id, parent_level) '
                        'REFERENCES {info_table} (geoid, level) '.format(
            rel_table=self.input()['rel'].qualified_tablename,
            info_table=self.input()['info'].qualified_tablename))
        session.commit()

    def complete(self):
        session = current_session()
        sql = "SELECT 1 FROM information_schema.constraint_column_usage " \
              "WHERE table_schema = '{schema}' " \
              "  AND table_name = '{table}' " \
              "  AND constraint_name = '{constraint}'"
        table = self.input()['info']
        check = sql.format(schema=table.schema,
                           table=table.tablename,
                           constraint='ushierarchy_fk_parent')
        result = session.execute(check).fetchall()
        return len(result) > 0

class _YearLevelsTask:
    year = IntParameter()
    levels = ListParameter(significant=False)

class USHierarchyInfoUnion(TempTableTask, _YearLevelsTask):

    def requires(self):
        return [USLevelInfo(year=self.year, geography=level)
                for level in self.levels]

    def run(self):
        session = current_session()
        tablename = self.output().qualified_tablename
        session.execute(_union(self.input(), tablename))
        alter_sql = 'ALTER TABLE {tablename} ADD PRIMARY KEY (geoid, level)'
        session.execute(alter_sql.format(tablename=tablename))
        session.commit()

class USHierarchyChildParentsUnion(TempTableTask, _YearLevelsTask):

    def requires(self):
        child_parents = self._child_parents()
        LOGGER.info('Child-parents: {}'.format(child_parents))
        return [USLevelInclusionHierarchy(year=self.year,
                                          current_geography=child_parent[0],
                                          parent_geographies=child_parent[1])
                for child_parent in child_parents]

    def _child_parents(self):
        child_parents = []
        previous = None
        for idx, level in enumerate(self.levels):
            if previous:
                parents = self.levels[idx:]
                if parents:
                    child_parents.append([previous, parents])
            previous = level
        return child_parents

    def run(self):
        session = current_session()
        tablename = self.output().qualified_tablename
        session.execute(_union(self.input(), tablename))
        alter_sql = 'ALTER TABLE {tablename} ADD PRIMARY KEY ' \
                    '(child_id, child_level, parent_id, parent_level)'
        session.execute(alter_sql.format(tablename=tablename))
        session.commit()

GEOGRAPHIES_ABBREVIATIONS = {
    'school_district_secondary': 'sds',
    'school_district_unified': 'sdu',
    'school_district_elementary': 'sde'
}

def _abbreviation(geography):
    return GEOGRAPHIES_ABBREVIATIONS.get(geography, geography)

def abbr_tablename(target, geographies, year):
    abbrs = [_abbreviation(geography) for geography in geographies]

    if [x for x in zip(abbrs, geographies) if x[0] != x[1]]:
        splits = target.split('_')
        target = '_'.join([splits[0], '_'.join(abbrs), str(year), splits[-1]])

    return target

class USLevelInclusionHierarchy(Task):
    year = IntParameter()
    current_geography = Parameter()
    parent_geographies = ListParameter()

    UNWEIGHTED_CHILD_SQL = """
        SELECT DISTINCT child_id, child_level
        FROM {table}
        WHERE weight = 1
        GROUP BY child_id, child_level
        HAVING count(1) > 1
    """

    def requires(self):
        return {
            'level': USLevelHierarchy(year=self.year,
                                      current_geography=self.current_geography,
                                      parent_geographies=self.parent_geographies),
            'current_geom': ShorelineClip(year=self.year,
                                          geography=self.current_geography),
            'parent_geoms': [ShorelineClip(year=self.year,
                                           geography=parent_geography) for
                             parent_geography in self.parent_geographies]
        }

    def run(self):
        session = current_session()
        sql = '''
            UPDATE {table}
            SET weight = ST_Area(
                ST_Intersection(cgt.the_geom, pgt.the_geom), False)
            FROM
                observatory.{current_geom_table} cgt,
                observatory.{parent_geom_table} pgt
            WHERE cgt.geoid = {table}.child_id
              AND pgt.geoid = {table}.parent_id
              AND (child_id, child_level) IN (
              {unweighted_child_sql}
        )
        '''
        table = self.input()['level'].qualified_tablename
        for parent_geom in self.input()['parent_geoms']:
            session.execute(
                sql.format(
                    table=table,
                    current_geom_table=self.input()['current_geom'].get(
                        session).tablename,
                    parent_geom_table=parent_geom.get(
                        session).tablename,
                    unweighted_child_sql=self.UNWEIGHTED_CHILD_SQL.format(
                        table=table)
                )
            )
        session.commit()

    def complete(self):
        sql = self.UNWEIGHTED_CHILD_SQL.format(
            table=self.input()['level'].qualified_tablename)
        try:
            return len(current_session().execute(sql).fetchall()) == 0
        except:
            # Table doesn't exist yet
            return False

    def output(self):
        return self.input()['level']

class USLevelHierarchy(TempTableTask):
    year = IntParameter()
    current_geography = Parameter()
    parent_geographies = ListParameter()

    def requires(self):
        return {
            'current_info': USLevelInfo(year=self.year,
                                        geography=self.current_geography),
            'current_geom': ShorelineClip(year=self.year,
                                          geography=self.current_geography),
            'parents_infos': [USLevelInfo(year=self.year,
                                          geography=parent_geography) for
                              parent_geography in self.parent_geographies],
            'parents_geoms': [ShorelineClip(year=self.year,
                                            geography=parent_geography) for
                              parent_geography in self.parent_geographies]
        }

    def target_tablename(self):
        return abbr_tablename(super(USLevelHierarchy, self).target_tablename(),
                              [self.current_geography, 'parents'],
                              self.year)

    def run(self):
        session = current_session()
        input_ = self.input()
        current_info_tablename = input_['current_info'].qualified_tablename
        current_geom_table = input_['current_geom'].get(session)
        parent_info_tablename = input_['parents_infos'][0].qualified_tablename
        parent_geom_table = input_['parents_geoms'][0].get(session)

        create_table_sql = '''
        CREATE TABLE {output_table} AS
        {child_parent_sql}
        '''

        # First creation will link child with direct parents and leave nulls
        # for those that don't have.
        session.execute(create_table_sql.format(
            output_table=self.output().qualified_tablename,
            child_parent_sql=self._CHILD_PARENT_SQL.format(
                current_info_table=current_info_tablename,
                current_geom_table=current_geom_table.tablename,
                parent_info_table=parent_info_tablename,
                parent_geom_table=parent_geom_table.tablename)
        ))

        inputs = list(zip(input_['parents_infos'], input_['parents_geoms']))
        for parent_info_geom in inputs[1:]:
            # For those without parents, insert the next ones
            parent_info_tablename = parent_info_geom[0].qualified_tablename
            parent_geom_table = parent_info_geom[1].get(session)
            fill_parents_sql = '''
                INSERT INTO {output_table}
                (child_id, child_level, parent_id, parent_level, weight)
                {child_parent_sql}
                INNER JOIN {output_table} ot ON ot.child_id = cit.geoid
                                            AND ot.child_level = cit.level
                WHERE ot.parent_id IS NULL
            '''
            session.execute(fill_parents_sql.format(
                output_table=self.output().qualified_tablename,
                child_parent_sql=self._CHILD_PARENT_SQL.format(
                    current_info_table=current_info_tablename,
                    current_geom_table=current_geom_table.tablename,
                    parent_info_table=parent_info_tablename,
                    parent_geom_table=parent_geom_table.tablename)))

            # ... and then, delete the rows with null parents for those
            # child that have any parent
            delete_non_orphans = '''
                DELETE FROM {output_table} ot
                WHERE ot.parent_id IS NULL
                AND (child_id, child_level) IN (
                    SELECT child_id, child_level
                    FROM {output_table}
                    WHERE parent_id IS NOT NULL
            )
            '''
            session.execute(delete_non_orphans.format(
                output_table=self.output().qualified_tablename))

        session.commit()

    _CHILD_PARENT_SQL = '''
        SELECT
            cit.geoid AS child_id,
            cit.level AS child_level,
            pgt.geoid AS parent_id,
            pit.level AS parent_level,
            1.0::FLOAT AS weight
        FROM {current_info_table} cit
        INNER JOIN observatory.{current_geom_table} cgt ON cit.geoid = cgt.geoid
        LEFT JOIN observatory.{parent_geom_table} pgt
            ON ST_Within(ST_PointOnSurface(cgt.the_geom), pgt.the_geom)
        LEFT JOIN {parent_info_table} pit ON pgt.geoid = pit.geoid
    '''

class _YearGeographyTask:
    year = IntParameter()
    geography = Parameter()

class USLevelInfo(WrapperTask, _YearGeographyTask):
    GEOGRAPHIES_WITHOUT_GEONAMES = ['zcta5']

    def requires(self):
        if self.geography in self.GEOGRAPHIES_WITHOUT_GEONAMES:
            return USLevelInfoFromShorelineClip(year=self.year,
                                                geography=self.geography)
        else:
            return USLevelInfoFromGeoNames(year=self.year,
                                           geography=self.geography)

    def output(self):
        return self.input()

class USLevelInfoFromShorelineClip(TempTableTask, _YearGeographyTask):

    def requires(self):
        return ShorelineClip(year=self.year, geography=self.geography)

    def run(self):
        session = current_session()

        names_table = self.input().get(session).tablename
        output_table = self.output().qualified_tablename
        output_tablename = self.output().tablename

        sql = '''
            CREATE TABLE {output_table} AS
            SELECT n.geoid geoid, '{geography}' as level, n.geoid as name, {year} as year
            FROM observatory.{names_table} n
        '''
        session.execute(
            sql.format(output_table=output_table,
                       geography=self.geography,
                       year=self.year,
                       names_table=names_table))
        session.execute('''
            CREATE INDEX {output_tablename}_idx ON {output_table} (geoid)
        '''.format(output_table=output_table,
                   output_tablename=output_tablename))
        session.commit()

class USLevelInfoFromGeoNames(TempTableTask, _YearGeographyTask):

    def requires(self):
        return GeoNamesTable(year=self.year, geography=self.geography)

    def target_tablename(self):
        return abbr_tablename(
            super(USLevelInfoFromGeoNames, self).target_tablename(),
            [self.geography], self.year)

    def run(self):
        session = current_session()

        names_table = self.input().get(session).tablename
        output_table = self.output().qualified_tablename
        output_tablename = self.output().tablename

        sql = '''
            CREATE TABLE {output_table} AS
            SELECT n.geoidsc geoid, '{geography}' as level, n.geoname, {year} as year
            FROM observatory.{names_table} n
        '''
        session.execute(
            sql.format(output_table=output_table,
                       geography=self.geography,
                       year=self.year,
                       names_table=names_table))
        session.execute('''
            CREATE INDEX {output_tablename}_idx ON {output_table} (geoid)
        '''.format(output_table=output_table,
                   output_tablename=output_tablename))
        session.commit()
<commit_msg>USLevelHierarchyWeights extraction<commit_after>from luigi import Task, WrapperTask, Parameter, IntParameter, ListParameter
from tasks.us.census.tiger import GeoNamesTable, ShorelineClip, SUMLEVELS
from tasks.base_tasks import TempTableTask
from tasks.meta import current_session
from lib.logger import get_logger

LOGGER = get_logger(__name__)

def _union(tables, output):
    unions = ['SELECT * FROM {table}'.format(table=table.qualified_tablename)
              for table in tables]
    return 'CREATE TABLE {output} AS {unions}'.format(
        output=output,
        unions=' UNION '.join(unions))

class USHierarchy(Task):
    year = IntParameter()

    def requires(self):
        levels = self._levels()
        LOGGER.info('Levels: {}'.format(levels))
        return {
            'info': USHierarchyInfoUnion(year=self.year, levels=levels),
            'rel': USHierarchyChildParentsUnion(year=self.year, levels=levels)
        }

    def _levels(self):
        level_infos = [(level, info) for level, info in SUMLEVELS.items()]

        sorted_level_infos = reversed(sorted(level_infos,
                                             key=lambda level_info:
                                             level_info[1][
                                                 'weight']))
        return [level_info[0] for level_info in sorted_level_infos]

    def run(self):
        session = current_session()
        session.execute('ALTER TABLE {rel_table} ADD '
                        'CONSTRAINT ushierarchy_fk_child '
                        'FOREIGN KEY (child_id, child_level) '
                        'REFERENCES {info_table} (geoid, level) '.format(
            rel_table=self.input()['rel'].qualified_tablename,
            info_table=self.input()['info'].qualified_tablename))
        session.execute('ALTER TABLE {rel_table} ADD '
                        'CONSTRAINT ushierarchy_fk_parent '
                        'FOREIGN KEY (parent_id, parent_level) '
                        'REFERENCES {info_table} (geoid, level) '.format(
            rel_table=self.input()['rel'].qualified_tablename,
            info_table=self.input()['info'].qualified_tablename))
        session.commit()

    def complete(self):
        session = current_session()
        sql = "SELECT 1 FROM information_schema.constraint_column_usage " \
              "WHERE table_schema = '{schema}' " \
              "  AND table_name = '{table}' " \
              "  AND constraint_name = '{constraint}'"
        table = self.input()['info']
        check = sql.format(schema=table.schema,
                           table=table.tablename,
                           constraint='ushierarchy_fk_parent')
        result = session.execute(check).fetchall()
        return len(result) > 0

class _YearLevelsTask:
    year = IntParameter()
    levels = ListParameter(significant=False)

class USHierarchyInfoUnion(TempTableTask, _YearLevelsTask):

    def requires(self):
        return [USLevelInfo(year=self.year, geography=level)
                for level in self.levels]

    def run(self):
        session = current_session()
        tablename = self.output().qualified_tablename
        session.execute(_union(self.input(), tablename))
        alter_sql = 'ALTER TABLE {tablename} ADD PRIMARY KEY (geoid, level)'
        session.execute(alter_sql.format(tablename=tablename))
        session.commit()

class USHierarchyChildParentsUnion(TempTableTask, _YearLevelsTask):

    def requires(self):
        child_parents = self._child_parents()
        LOGGER.info('Child-parents: {}'.format(child_parents))
        return {
            'hierarchy': [USLevelInclusionHierarchy(year=self.year,
                                                    current_geography=
                                                    child_parent[0],
                                                    parent_geographies=
                                                    child_parent[1])
                          for child_parent in child_parents],
            'weight': [USLevelHierarchyWeights(year=self.year,
                                               current_geography=child_parent[
                                                   0],
                                               parent_geographies=child_parent[
                                                   1])
                       for child_parent in child_parents]
        }

    def _child_parents(self):
        child_parents = []
        previous = None
        for idx, level in enumerate(self.levels):
            if previous:
                parents = self.levels[idx:]
                if parents:
                    child_parents.append([previous, parents])
            previous = level
        return child_parents

    def run(self):
        session = current_session()
        tablename = self.output().qualified_tablename
        session.execute(_union(self.input(), tablename))
        alter_sql = 'ALTER TABLE {tablename} ADD PRIMARY KEY ' \
                    '(child_id, child_level, parent_id, parent_level)'
        session.execute(alter_sql.format(tablename=tablename))
        session.commit()

GEOGRAPHIES_ABBREVIATIONS = {
    'school_district_secondary': 'sds',
    'school_district_unified': 'sdu',
    'school_district_elementary': 'sde'
}

def _abbreviation(geography):
    return GEOGRAPHIES_ABBREVIATIONS.get(geography, geography)

def abbr_tablename(target, geographies, year):
    abbrs = [_abbreviation(geography) for geography in geographies]

    if [x for x in zip(abbrs, geographies) if x[0] != x[1]]:
        splits = target.split('_')
        target = '_'.join([splits[0], '_'.join(abbrs), str(year), splits[-1]])

    return target

class USLevelInclusionHierarchy(WrapperTask):
    year = IntParameter()
    current_geography = Parameter()
    parent_geographies = ListParameter()

    def requires(self):
        return {
            'level': USLevelHierarchy(year=self.year,
                                      current_geography=self.current_geography,
                                      parent_geographies=self.parent_geographies)
        }

class USLevelHierarchyWeights(Task):
    year = IntParameter()
    current_geography = Parameter()
    parent_geographies = ListParameter()

    UNWEIGHTED_CHILD_SQL = """
        SELECT DISTINCT child_id, child_level
        FROM {table}
        WHERE weight = 1
        GROUP BY child_id, child_level
        HAVING count(1) > 1
    """

    def requires(self):
        return {
            'level': USLevelHierarchy(year=self.year,
                                      current_geography=self.current_geography,
                                      parent_geographies=self.parent_geographies),
            'current_geom': ShorelineClip(year=self.year,
                                          geography=self.current_geography),
            'parent_geoms': [ShorelineClip(year=self.year,
                                           geography=parent_geography) for
                             parent_geography in self.parent_geographies]
        }

    def run(self):
        session = current_session()
        sql = '''
            UPDATE {table}
            SET weight = ST_Area(
                ST_Intersection(cgt.the_geom, pgt.the_geom), False)
            FROM
                observatory.{current_geom_table} cgt,
                observatory.{parent_geom_table} pgt
            WHERE cgt.geoid = {table}.child_id
              AND pgt.geoid = {table}.parent_id
              AND (child_id, child_level) IN (
              {unweighted_child_sql}
        )
        '''
        table = self.input()['level'].qualified_tablename
        for parent_geom in self.input()['parent_geoms']:
            session.execute(
                sql.format(
                    table=table,
                    current_geom_table=self.input()['current_geom'].get(
                        session).tablename,
                    parent_geom_table=parent_geom.get(
                        session).tablename,
                    unweighted_child_sql=self.UNWEIGHTED_CHILD_SQL.format(
                        table=table)
                )
            )
        session.commit()

    def complete(self):
        sql = self.UNWEIGHTED_CHILD_SQL.format(
            table=self.input()['level'].qualified_tablename)
        try:
            return len(current_session().execute(sql).fetchall()) == 0
        except:
            # Table doesn't exist yet
            return False

    def output(self):
        return self.input()['level']

class USLevelHierarchy(TempTableTask):
    year = IntParameter()
    current_geography = Parameter()
    parent_geographies = ListParameter()

    def requires(self):
        return {
            'current_info': USLevelInfo(year=self.year,
                                        geography=self.current_geography),
            'current_geom': ShorelineClip(year=self.year,
                                          geography=self.current_geography),
            'parents_infos': [USLevelInfo(year=self.year,
                                          geography=parent_geography) for
                              parent_geography in self.parent_geographies],
            'parents_geoms': [ShorelineClip(year=self.year,
                                            geography=parent_geography) for
                              parent_geography in self.parent_geographies]
        }

    def target_tablename(self):
        return abbr_tablename(super(USLevelHierarchy, self).target_tablename(),
                              [self.current_geography, 'parents'],
                              self.year)

    def run(self):
        session = current_session()
        input_ = self.input()
        current_info_tablename = input_['current_info'].qualified_tablename
        current_geom_table = input_['current_geom'].get(session)
        parent_info_tablename = input_['parents_infos'][0].qualified_tablename
        parent_geom_table = input_['parents_geoms'][0].get(session)

        create_table_sql = '''
        CREATE TABLE {output_table} AS
        {child_parent_sql}
        '''

        # First creation will link child with direct parents and leave nulls
        # for those that don't have.
        session.execute(create_table_sql.format(
            output_table=self.output().qualified_tablename,
            child_parent_sql=self._CHILD_PARENT_SQL.format(
                current_info_table=current_info_tablename,
                current_geom_table=current_geom_table.tablename,
                parent_info_table=parent_info_tablename,
                parent_geom_table=parent_geom_table.tablename)
        ))

        inputs = list(zip(input_['parents_infos'], input_['parents_geoms']))
        for parent_info_geom in inputs[1:]:
            # For those without parents, insert the next ones
            parent_info_tablename = parent_info_geom[0].qualified_tablename
            parent_geom_table = parent_info_geom[1].get(session)
            fill_parents_sql = '''
                INSERT INTO {output_table}
                (child_id, child_level, parent_id, parent_level, weight)
                {child_parent_sql}
                INNER JOIN {output_table} ot ON ot.child_id = cit.geoid
                                            AND ot.child_level = cit.level
                WHERE ot.parent_id IS NULL
            '''
            session.execute(fill_parents_sql.format(
                output_table=self.output().qualified_tablename,
                child_parent_sql=self._CHILD_PARENT_SQL.format(
                    current_info_table=current_info_tablename,
                    current_geom_table=current_geom_table.tablename,
                    parent_info_table=parent_info_tablename,
                    parent_geom_table=parent_geom_table.tablename)))

            # ... and then, delete the rows with null parents for those
            # child that have any parent
            delete_non_orphans = '''
                DELETE FROM {output_table} ot
                WHERE ot.parent_id IS NULL
                AND (child_id, child_level) IN (
                    SELECT child_id, child_level
                    FROM {output_table}
                    WHERE parent_id IS NOT NULL
            )
            '''
            session.execute(delete_non_orphans.format(
                output_table=self.output().qualified_tablename))

        session.commit()

    _CHILD_PARENT_SQL = '''
        SELECT
            cit.geoid AS child_id,
            cit.level AS child_level,
            pgt.geoid AS parent_id,
            pit.level AS parent_level,
            1.0::FLOAT AS weight
        FROM {current_info_table} cit
        INNER JOIN observatory.{current_geom_table} cgt ON cit.geoid = cgt.geoid
        LEFT JOIN observatory.{parent_geom_table} pgt
            ON ST_Within(ST_PointOnSurface(cgt.the_geom), pgt.the_geom)
        LEFT JOIN {parent_info_table} pit ON pgt.geoid = pit.geoid
    '''

class _YearGeographyTask:
    year = IntParameter()
    geography = Parameter()

class USLevelInfo(WrapperTask, _YearGeographyTask):
    GEOGRAPHIES_WITHOUT_GEONAMES = ['zcta5']

    def requires(self):
        if self.geography in self.GEOGRAPHIES_WITHOUT_GEONAMES:
            return USLevelInfoFromShorelineClip(year=self.year,
                                                geography=self.geography)
        else:
            return USLevelInfoFromGeoNames(year=self.year,
                                           geography=self.geography)

    def output(self):
        return self.input()

class USLevelInfoFromShorelineClip(TempTableTask, _YearGeographyTask):

    def requires(self):
        return ShorelineClip(year=self.year, geography=self.geography)

    def run(self):
        session = current_session()

        names_table = self.input().get(session).tablename
        output_table = self.output().qualified_tablename
        output_tablename = self.output().tablename

        sql = '''
            CREATE TABLE {output_table} AS
            SELECT n.geoid geoid, '{geography}' as level, n.geoid as name, {year} as year
            FROM observatory.{names_table} n
        '''
        session.execute(
            sql.format(output_table=output_table,
                       geography=self.geography,
                       year=self.year,
                       names_table=names_table))
        session.execute('''
            CREATE INDEX {output_tablename}_idx ON {output_table} (geoid)
        '''.format(output_table=output_table,
                   output_tablename=output_tablename))
        session.commit()

class USLevelInfoFromGeoNames(TempTableTask, _YearGeographyTask):

    def requires(self):
        return GeoNamesTable(year=self.year, geography=self.geography)

    def target_tablename(self):
        return abbr_tablename(
            super(USLevelInfoFromGeoNames, self).target_tablename(),
            [self.geography], self.year)

    def run(self):
        session = current_session()

        names_table = self.input().get(session).tablename
        output_table = self.output().qualified_tablename
        output_tablename = self.output().tablename

        sql = '''
            CREATE TABLE {output_table} AS
            SELECT n.geoidsc geoid, '{geography}' as level, n.geoname, {year} as year
            FROM observatory.{names_table} n
        '''
        session.execute(
            sql.format(output_table=output_table,
                       geography=self.geography,
                       year=self.year,
                       names_table=names_table))
        session.execute('''
            CREATE INDEX {output_tablename}_idx ON {output_table} (geoid)
        '''.format(output_table=output_table,
                   output_tablename=output_tablename))
        session.commit()
<|endoftext|>


<commit_before># Copyright (c) 2015, Djaodjin Inc.
# All rights reserved.
#
# Redistribution and use in source and binary forms, with or without
# modification, are permitted provided that the following conditions are met:
#
# 1. Redistributions of source code must retain the above copyright notice,
#    this list of conditions and the following disclaimer.
# 2. Redistributions in binary form must reproduce the above copyright notice,
#    this list of conditions and the following disclaimer in the documentation
#    and/or other materials provided with the distribution.
#
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO,
# THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR
# PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR
# CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,
# EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,
# PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS;
# OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY,
# WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR
# OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF
# ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.

"""
User Model for the signup app
"""

import datetime, hashlib, logging, random, re

from django.db import models
from django.db import transaction
from django.contrib.auth.models import AbstractUser, UserManager
from django.utils.decorators import method_decorator
from django.utils.translation import ugettext_lazy as _
from django.utils.timezone import now as datetime_now

from signup import settings

LOGGER = logging.getLogger(__name__)

EMAIL_VERIFICATION_RE = re.compile('^%s$' % settings.EMAIL_VERIFICATION_PAT)

class ActivatedUserManager(UserManager):

    @method_decorator(transaction.atomic)
    def create_user(self, username, email=None, is_active=False, **kwargs):
        """
        Create an inactive user with a default username.

        We have a different notion of an active user than Django does.
        For Django when is_active is False, the user cannot be identified
        and requests fall back to Anonymous. That's a problem because
        we want a user who has given us a name and email address to be
        able to use the site. We only require a password for the second
        login. Our definition of inactive is thus a user that has an invalid
        password.
        """
        if not is_active:
            if not username:
                username = email.split('@')[0] \
                    + ''.join(random.choice('0123456789') for count in range(3))
            # Force is_active to True and create an email verification key
            # (see above definition of active user).
            user = super(ActivatedUserManager, self).create_user(
                username, email=email, **kwargs) # Django 1.7 is_active=True
            salt = hashlib.sha1(str(random.random())).hexdigest()[:5]
            if isinstance(username, unicode):
                username = username.encode('utf-8')
            user.email_verification_key = hashlib.sha1(
                salt+username).hexdigest()
            user.save()
        else:
            user = super(ActivatedUserManager, self).create_user(
                username, email=email, is_active=is_active, **kwargs)
        return user

    def find_user(self, email_verification_key):
        """
        Find a user based on a verification key but do not activate the user.
        """
        if EMAIL_VERIFICATION_RE.search(email_verification_key):
            try:
                user = self.get(email_verification_key=email_verification_key)
                if not user.email_verification_key_expired():
                    return user
            except ActivatedUser.DoesNotExist:
                pass # We return None instead here.
        return None

    def activate_user(self, email_verification_key):
        """
        Activate a user whose email address has been verified.
        """
        user = self.find_user(email_verification_key)
        if user:
            LOGGER.info('user #%d activated through code: %s',
                        user.id, user.email_verification_key)
            user.email_verification_key = ActivatedUser.VERIFIED
            user.is_active = True
            user.save()
        return user

class ActivatedUser(AbstractUser):
    """
    A user model that requires activation. We use it to verify
    the email address.
    """
    VERIFIED = "VERIFIED"

    objects = ActivatedUserManager()

    email_verification_key = models.CharField(
        _('email verification key'), max_length=40)

    def email_verification_key_expired(self):
        expiration_date = datetime.timedelta(days=settings.KEY_EXPIRATION)
        return self.email_verification_key == self.VERIFIED or \
               (self.last_login + expiration_date <= datetime_now())
    email_verification_key_expired.boolean = True

    def __unicode__(self):
        return self.username

    @property
    def has_invalid_password(self):
        return self.password.startswith('!')

    @property
    def is_reachable(self):
        """
        Returns True if the user is reachable by email.
        """
        return self.email_verification_key == ActivatedUser.VERIFIED

    @property
    def printable_name(self):
        if self.first_name:
            return self.first_name
        return self.username

    class Meta: #pylint: disable=old-style-class,no-init
        db_table = u'auth_user'
        swappable = 'AUTH_USER_MODEL'
<commit_msg>fix: Django 1.7 compatible (take 2)<commit_after># Copyright (c) 2015, Djaodjin Inc.
# All rights reserved.
#
# Redistribution and use in source and binary forms, with or without
# modification, are permitted provided that the following conditions are met:
#
# 1. Redistributions of source code must retain the above copyright notice,
#    this list of conditions and the following disclaimer.
# 2. Redistributions in binary form must reproduce the above copyright notice,
#    this list of conditions and the following disclaimer in the documentation
#    and/or other materials provided with the distribution.
#
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO,
# THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR
# PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR
# CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,
# EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,
# PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS;
# OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY,
# WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR
# OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF
# ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.

"""
User Model for the signup app
"""

import datetime, hashlib, logging, random, re

from django.db import models
from django.db import transaction
from django.contrib.auth.models import AbstractUser, UserManager
from django.utils.decorators import method_decorator
from django.utils.translation import ugettext_lazy as _
from django.utils.timezone import now as datetime_now

from signup import settings

LOGGER = logging.getLogger(__name__)

EMAIL_VERIFICATION_RE = re.compile('^%s$' % settings.EMAIL_VERIFICATION_PAT)

class ActivatedUserManager(UserManager):

    @method_decorator(transaction.atomic)
    def create_user(self, username, email=None, is_active=False, **kwargs):
        """
        Create an inactive user with a default username.

        We have a different notion of an active user than Django does.
        For Django when is_active is False, the user cannot be identified
        and requests fall back to Anonymous. That's a problem because
        we want a user who has given us a name and email address to be
        able to use the site. We only require a password for the second
        login. Our definition of inactive is thus a user that has an invalid
        password.
        """
        if not is_active:
            if not username:
                username = email.split('@')[0] \
                    + ''.join(random.choice('0123456789') for count in range(3))
            # Force is_active to True and create an email verification key
            # (see above definition of active user).
            user = super(ActivatedUserManager, self).create_user(
                username, email=email, **kwargs) # Django 1.7 is_active=True
            salt = hashlib.sha1(str(random.random())).hexdigest()[:5]
            if isinstance(username, unicode):
                username = username.encode('utf-8')
            user.email_verification_key = hashlib.sha1(
                salt+username).hexdigest()
            user.save()
        else:
            user = super(ActivatedUserManager, self).create_user(
                username, email=email, **kwargs)
        return user

    def find_user(self, email_verification_key):
        """
        Find a user based on a verification key but do not activate the user.
        """
        if EMAIL_VERIFICATION_RE.search(email_verification_key):
            try:
                user = self.get(email_verification_key=email_verification_key)
                if not user.email_verification_key_expired():
                    return user
            except ActivatedUser.DoesNotExist:
                pass # We return None instead here.
        return None

    def activate_user(self, email_verification_key):
        """
        Activate a user whose email address has been verified.
        """
        user = self.find_user(email_verification_key)
        if user:
            LOGGER.info('user #%d activated through code: %s',
                        user.id, user.email_verification_key)
            user.email_verification_key = ActivatedUser.VERIFIED
            user.is_active = True
            user.save()
        return user

class ActivatedUser(AbstractUser):
    """
    A user model that requires activation. We use it to verify
    the email address.
    """
    VERIFIED = "VERIFIED"

    objects = ActivatedUserManager()

    email_verification_key = models.CharField(
        _('email verification key'), max_length=40)

    def email_verification_key_expired(self):
        expiration_date = datetime.timedelta(days=settings.KEY_EXPIRATION)
        return self.email_verification_key == self.VERIFIED or \
               (self.last_login + expiration_date <= datetime_now())
    email_verification_key_expired.boolean = True

    def __unicode__(self):
        return self.username

    @property
    def has_invalid_password(self):
        return self.password.startswith('!')

    @property
    def is_reachable(self):
        """
        Returns True if the user is reachable by email.
        """
        return self.email_verification_key == ActivatedUser.VERIFIED

    @property
    def printable_name(self):
        if self.first_name:
            return self.first_name
        return self.username

    class Meta: #pylint: disable=old-style-class,no-init
        db_table = u'auth_user'
        swappable = 'AUTH_USER_MODEL'
<|endoftext|>


<commit_before># coding: utf-8

__version__ = '0.1.0'
__author__ = 'Daniel Krikun'
__license__ = 'MIT'

import sys
import logging
import msvcrt
import time
from argparse import ArgumentParser
import zmq

from agent_pb2 import *
from ffmpeg_recorder import *

def parse_cmdline_args():
    """Parse command-line arguments."""

    prsr = ArgumentParser(description='Distributed screen recorder agent')
    prsr.add_argument('--version', action='version',
            version='%(prog)s {}'.format(__version__))

    verbosity_group = prsr.add_mutually_exclusive_group()
    verbosity_group.add_argument('-V', '--verbose', action='store_true',
                                 help='be verbose')
    verbosity_group.add_argument('-q', '--silent', action='store_true',
                                 help='be silent')
    prsr.add_argument('--ctrl-address', help='remote control address')
    prsr.add_argument('--status-address', help='status address')
    prsr.add_argument('--nostdin', action='store_true',
            help='disable kbhit control')

    prsr.add_argument('--show-video', action='store_true',
            help='show the recorded video in a window')

    return prsr.parse_args()

def main():
    """Main here."""

    args = parse_cmdline_args()

    # decide on logging severity level
    severity_level = logging.CRITICAL if args.silent else logging.DEBUG \
            if args.verbose else logging.INFO

    # set up logging
    logging.basicConfig(stream=sys.stderr, level=severity_level,
            format='%(message)s')

    # debug-print the parsed arguments
    logging.debug(args)

    # set up zmq machinery
    zctx = zmq.Context()
    zsck_ctrl = zctx.socket(zmq.PULL)

    ctrl_addr = args.ctrl_address or 'tcp://*:17267'
    zsck_ctrl.bind(ctrl_addr)

    zsck_status = zctx.socket(zmq.PUB)
    status_addr = args.status_address or 'tcp://*:17268'
    zsck_status.bind(status_addr)

    recorder = FfmpegRecorder()
    recorder.debug_show_video = args.show_video

    # main loop frequency
    frequency = 30.;

    # last status, to send messages on change only
    recording = recorder.running
    paused = recorder.paused
    has_crashed = recorder.has_crashed

    while True:
        # check if an incoming message is available
        events = zsck_ctrl.poll(1000./frequency)

        if events & zmq.POLLIN:
            zmsg = zsck_ctrl.recv(zmq.NOBLOCK)
            msg.ParseFromString(zmsg)
            logging.debug('recved: %s', msg)

            if msg.opcode == FfmpegControl.RECORD:
                recorder.run()

            elif msg.opcode == FfmpegControl.IDLE:
                recorder.stop()

            elif msg.opcode == FfmpegControl.PAUSE:
                recorder.pause()

            elif msg.opcode == FfmpegControl.UNPAUSE:
                recorder.unpause()

            elif msg.opcode == FfmpegControl.SHUTDOWN:
                recorder.stop()
                break

            # update recording parameters
            if msg.HasField('capture_x'):
                recorder.capture_x = msg.capture_x

            if msg.HasField('capture_y'):
                recorder.capture_y = msg.capture_y

            if msg.HasField('capture_width'):
                recorder.capture_width = msg.capture_width

            if msg.HasField('capture_height'):
                recorder.capture_height = msg.capture_height

            if msg.HasField('capture_fps'):
                recorder.capture_fps = msg.capture_fps

            if msg.HasField('audio_device'):
                recorder.audio_device = msg.audio_device

            if msg.HasField('video_device'):
                recorder.video_device = msg.video_device

            if msg.HasField('scale'):
                recorder.scale = msg.scale

            if msg.HasField('output_file'):
                recorder.output_file = msg.output_file

            if msg.HasField('debug_show_video'):
                recorder.debug_show_video = msg.debug_show_video

        # dispatch kbhit commands
        if not args.nostdin:
            cmd = msvcrt.getch() if msvcrt.kbhit() else 0

            if cmd == 'r':
                recorder.run()

            elif cmd == 't':
                recorder.stop()

            elif cmd == 'p':
                recorder.pause()

            elif cmd == 'P':
                recorder.unpause()

            elif cmd == 'X':
                recorder.stop()
                break

        # publish status
        status = FfmpegStatus()
        dirty = False

        if recorder.running != recording:
            recording = recorder.running
            status.is_recording = recorder.running
            dirty = True

        if recorder.paused != paused:
            paused = recorder.paused
            status.is_paused = recorder.paused
            dirty = True

        if recorder.has_crashed != has_crashed:
            has_crashed = recorder.has_crashed
            status.has_crashed = recorder.has_crashed
            dirty = True

        if dirty:
            zsck_status.send(status.SerializeToString(), zmq.NOBLOCK)

if __name__ == "__main__":
    sys.exit(main())

<commit_msg>recv() non-blocking after polling ok<commit_after># coding: utf-8

__version__ = '0.1.0'
__author__ = 'Daniel Krikun'
__license__ = 'MIT'

import sys
import logging
import msvcrt
import time
from argparse import ArgumentParser
import zmq

from agent_pb2 import *
from ffmpeg_recorder import *

def parse_cmdline_args():
    """Parse command-line arguments."""

    prsr = ArgumentParser(description='Distributed screen recorder agent')
    prsr.add_argument('--version', action='version',
            version='%(prog)s {}'.format(__version__))

    verbosity_group = prsr.add_mutually_exclusive_group()
    verbosity_group.add_argument('-V', '--verbose', action='store_true',
                                 help='be verbose')
    verbosity_group.add_argument('-q', '--silent', action='store_true',
                                 help='be silent')
    prsr.add_argument('--ctrl-address', help='remote control address')
    prsr.add_argument('--status-address', help='status address')
    prsr.add_argument('--nostdin', action='store_true',
            help='disable kbhit control')

    prsr.add_argument('--show-video', action='store_true',
            help='show the recorded video in a window')

    return prsr.parse_args()

def main():
    """Main here."""

    args = parse_cmdline_args()

    # decide on logging severity level
    severity_level = logging.CRITICAL if args.silent else logging.DEBUG \
            if args.verbose else logging.INFO

    # set up logging
    logging.basicConfig(stream=sys.stderr, level=severity_level,
            format='%(message)s')

    # debug-print the parsed arguments
    logging.debug(args)

    # set up zmq machinery
    zctx = zmq.Context()
    zsck_ctrl = zctx.socket(zmq.PULL)

    ctrl_addr = args.ctrl_address or 'tcp://*:17267'
    zsck_ctrl.bind(ctrl_addr)

    zsck_status = zctx.socket(zmq.PUB)
    status_addr = args.status_address or 'tcp://*:17268'
    zsck_status.bind(status_addr)

    recorder = FfmpegRecorder()
    recorder.debug_show_video = args.show_video

    # main loop frequency
    frequency = 30.;

    # last status, to send messages on change only
    recording = recorder.running
    paused = recorder.paused
    has_crashed = recorder.has_crashed

    while True:
        # check if an incoming message is available
        events = zsck_ctrl.poll(1000./frequency)

        if events & zmq.POLLIN:
            zmsg = zsck_ctrl.recv()
            msg.ParseFromString(zmsg)
            logging.debug('recved: %s', msg)

            if msg.opcode == FfmpegControl.RECORD:
                recorder.run()

            elif msg.opcode == FfmpegControl.IDLE:
                recorder.stop()

            elif msg.opcode == FfmpegControl.PAUSE:
                recorder.pause()

            elif msg.opcode == FfmpegControl.UNPAUSE:
                recorder.unpause()

            elif msg.opcode == FfmpegControl.SHUTDOWN:
                recorder.stop()
                break

            # update recording parameters
            if msg.HasField('capture_x'):
                recorder.capture_x = msg.capture_x

            if msg.HasField('capture_y'):
                recorder.capture_y = msg.capture_y

            if msg.HasField('capture_width'):
                recorder.capture_width = msg.capture_width

            if msg.HasField('capture_height'):
                recorder.capture_height = msg.capture_height

            if msg.HasField('capture_fps'):
                recorder.capture_fps = msg.capture_fps

            if msg.HasField('audio_device'):
                recorder.audio_device = msg.audio_device

            if msg.HasField('video_device'):
                recorder.video_device = msg.video_device

            if msg.HasField('scale'):
                recorder.scale = msg.scale

            if msg.HasField('output_file'):
                recorder.output_file = msg.output_file

            if msg.HasField('debug_show_video'):
                recorder.debug_show_video = msg.debug_show_video

        # dispatch kbhit commands
        if not args.nostdin:
            cmd = msvcrt.getch() if msvcrt.kbhit() else 0

            if cmd == 'r':
                recorder.run()

            elif cmd == 't':
                recorder.stop()

            elif cmd == 'p':
                recorder.pause()

            elif cmd == 'P':
                recorder.unpause()

            elif cmd == 'X':
                recorder.stop()
                break

        # publish status
        status = FfmpegStatus()
        dirty = False

        if recorder.running != recording:
            recording = recorder.running
            status.is_recording = recorder.running
            dirty = True

        if recorder.paused != paused:
            paused = recorder.paused
            status.is_paused = recorder.paused
            dirty = True

        if recorder.has_crashed != has_crashed:
            has_crashed = recorder.has_crashed
            status.has_crashed = recorder.has_crashed
            dirty = True

        if dirty:
            zsck_status.send(status.SerializeToString(), zmq.NOBLOCK)

if __name__ == "__main__":
    sys.exit(main())

<|endoftext|>


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

from __future__ import print_function
from __future__ import unicode_literals

import time

from tdclient._model import Model

class Schema(object):
    """Schema of a database table on Treasure Data Service
    """

    class Field(object):
        def __init__(self, name, type):
            self._name = name
            self._type = type

        @property
        def name(self):
            """
            TODO: add docstring
            """
            return self._name

        @property
        def type(self):
            """
            TODO: add docstring
            """
            return self._type

    def __init__(self, fields=[]):
        self._fields = fields

    @property
    def fields(self):
        """
        TODO: add docstring
        """
        return self._fields

    def add_field(self, name, type):
        """
        TODO: add docstring
        """
        self._fields.append(Field(name, type))

class Job(Model):
    """Job on Treasure Data Service
    """

    STATUS_QUEUED = "queued"
    STATUS_BOOTING = "booting"
    STATUS_RUNNING = "running"
    STATUS_SUCCESS = "success"
    STATUS_ERROR = "error"
    STATUS_KILLED = "killed"
    FINISHED_STATUS = [STATUS_SUCCESS, STATUS_ERROR, STATUS_KILLED]

    JOB_PRIORITY = {
        -2: "VERY LOW",
        -1: "LOW",
        0: "NORMAL",
        1: "HIGH",
        2: "VERY HIGH",
    }

    def __init__(self, client, job_id, type, query, **kwargs):
        super(Job, self).__init__(client)
        self._job_id = job_id
        self._type = type
        self._query = query
        self._feed(kwargs)

    def _feed(self, data={}):
        self._url = data.get("url")
        self._status = data.get("status")
        self._debug = data.get("debug")
        self._start_at = data.get("start_at")
        self._end_at = data.get("end_at")
        self._cpu_time = data.get("cpu_time")
        self._result = data.get("result")
        self._result_size = data.get("result_size")
        self._result_url = data.get("result_url")
        self._hive_result_schema = data.get("hive_result_schema")
        self._priority = data.get("priority")
        self._retry_limit = data.get("retry_limit")
        self._org_name = data.get("org_name")
        self._db_name = data.get("db_name")

    def update(self):
        data = self._client.api.show_job(self._job_id)
        self._feed(data)

    def _update_progress(self):
        if self._status not in self.FINISHED_STATUS:
            self._status = self._client.job_status(self._job_id)

    @property
    def id(self):
        """
        Returns: a string represents the identifier of the job
        """
        return self._job_id

    @property
    def job_id(self):
        """
        Returns: a string represents the identifier of the job
        """
        return self._job_id

    @property
    def type(self):
        """
        Returns: a string represents the engine type of the job (e.g. "hive", "presto", etc.)
        """
        return self._type

    @property
    def result_url(self):
        """
        Returns: a string of URL of the result on Treasure Data Service
        """
        return self._result_url

    @property
    def result_schema(self):
        """
        Returns: an array of array represents the type of result columns (Hive specific) (e.g. [["_c1", "string"], ["_c2", "bigint"]])
        """
        return self._hive_result_schema

    @property
    def priority(self):
        """
        Returns: a string represents the priority of the job (e.g. "NORMAL", "HIGH", etc.)
        """
        if self._priority in self.JOB_PRIORITY:
            return self.JOB_PRIORITY[self._priority]
        else:
            # just convert the value in string
            return str(self._priority)

    @property
    def retry_limit(self):
        """
        TODO: add docstring
        """
        return self._retry_limit

    @property
    def org_name(self):
        """
        Returns: organization name
        """
        return self._org_name

    @property
    def db_name(self):
        """
        Returns: a string represents the name of a database that job is running on
        """
        return self._db_name

    def wait(self, timeout=None, wait_interval=1):
        """Sleep until the job has been finished

        Params:
            timeout (int): Timeout in seconds. No timeout by default.
        """
        started_at = time.time()
        while not self.finished():
            if timeout is None or abs(time.time() - started_at) < timeout:
                time.sleep(wait_interval)
            else:
                raise RuntimeError("timeout") # TODO: throw proper error
        self.update()

    def kill(self):
        """Kill the job

        Returns: a string represents the status of killed job ("queued", "running")
        """
        response = self._client.kill(self.job_id)
        self.update()
        return response

    @property
    def query(self):
        """
        Returns: a string represents the query string of the job
        """
        return self._query

    def status(self):
        """
        Returns: a string represents the status of the job ("success", "error", "killed", "queued", "running")
        """
        if self._query is not None and not self.finished():
            self.update()
        return self._status

    @property
    def url(self):
        """
        Returns: a string of URL of the job on Treasure Data Service
        """
        return self._url

    def result(self):
        """
        Returns: an iterator of rows in result set
        """
        if not self.finished():
            raise ValueError("result is not ready")
        else:
            self.update()
            if self._result is None:
                for row in self._client.job_result_each(self._job_id):
                    yield row
            else:
                for row in self._result:
                    yield row

    def result_format(self, format):
        """
        Params:
            format (str): output format of result set

        Returns: an iterator of rows in result set
        """
        if not self.finished():
            raise ValueError("result is not ready")
        else:
            self.update()
            if self._result is None:
                for row in self._client.job_result_format_each(self._job_id, format):
                    yield row
            else:
                for row in self._result:
                    yield row

    def finished(self):
        """
        Returns: `True` if the job has been finished in success, error or killed
        """
        self._update_progress()
        return self._status in self.FINISHED_STATUS

    def success(self):
        """
        Returns: `True` if the job has been finished in success
        """
        self._update_progress()
        return self._status == self.STATUS_SUCCESS

    def error(self):
        """
        Returns: `True` if the job has been finished in error
        """
        self._update_progress()
        return self._status == self.STATUS_ERROR

    def killed(self):
        """
        Returns: `True` if the job has been finished in killed
        """
        self._update_progress()
        return self._status == self.STATUS_KILLED

    def queued(self):
        """
        Returns: `True` if the job is queued
        """
        self._update_progress()
        return self._status == self.STATUS_QUEUED

    def running(self):
        """
        Returns: `True` if the job is running
        """
        self._update_progress()
        return self._status == self.STATUS_RUNNING
<commit_msg>Revert `Job#_update_status` for backward compatibility and add warning<commit_after>#!/usr/bin/env python

from __future__ import print_function
from __future__ import unicode_literals

import time
import warnings

from tdclient._model import Model

class Schema(object):
    """Schema of a database table on Treasure Data Service
    """

    class Field(object):
        def __init__(self, name, type):
            self._name = name
            self._type = type

        @property
        def name(self):
            """
            TODO: add docstring
            """
            return self._name

        @property
        def type(self):
            """
            TODO: add docstring
            """
            return self._type

    def __init__(self, fields=[]):
        self._fields = fields

    @property
    def fields(self):
        """
        TODO: add docstring
        """
        return self._fields

    def add_field(self, name, type):
        """
        TODO: add docstring
        """
        self._fields.append(Field(name, type))

class Job(Model):
    """Job on Treasure Data Service
    """

    STATUS_QUEUED = "queued"
    STATUS_BOOTING = "booting"
    STATUS_RUNNING = "running"
    STATUS_SUCCESS = "success"
    STATUS_ERROR = "error"
    STATUS_KILLED = "killed"
    FINISHED_STATUS = [STATUS_SUCCESS, STATUS_ERROR, STATUS_KILLED]

    JOB_PRIORITY = {
        -2: "VERY LOW",
        -1: "LOW",
        0: "NORMAL",
        1: "HIGH",
        2: "VERY HIGH",
    }

    def __init__(self, client, job_id, type, query, **kwargs):
        super(Job, self).__init__(client)
        self._job_id = job_id
        self._type = type
        self._query = query
        self._feed(kwargs)

    def _feed(self, data={}):
        self._url = data.get("url")
        self._status = data.get("status")
        self._debug = data.get("debug")
        self._start_at = data.get("start_at")
        self._end_at = data.get("end_at")
        self._cpu_time = data.get("cpu_time")
        self._result = data.get("result")
        self._result_size = data.get("result_size")
        self._result_url = data.get("result_url")
        self._hive_result_schema = data.get("hive_result_schema")
        self._priority = data.get("priority")
        self._retry_limit = data.get("retry_limit")
        self._org_name = data.get("org_name")
        self._db_name = data.get("db_name")

    def update(self):
        """Update all fields of the job
        """
        data = self._client.api.show_job(self._job_id)
        self._feed(data)

    def _update_status(self):
        warnings.warn("_update_status() will be removed from future release. Please use update() instaed.)", category=DeprecationWarning)
        self.update()

    def _update_progress(self):
        """Update `_status` field of the job if it's not finished
        """
        if self._status not in self.FINISHED_STATUS:
            self._status = self._client.job_status(self._job_id)

    @property
    def id(self):
        """
        Returns: a string represents the identifier of the job
        """
        return self._job_id

    @property
    def job_id(self):
        """
        Returns: a string represents the identifier of the job
        """
        return self._job_id

    @property
    def type(self):
        """
        Returns: a string represents the engine type of the job (e.g. "hive", "presto", etc.)
        """
        return self._type

    @property
    def result_url(self):
        """
        Returns: a string of URL of the result on Treasure Data Service
        """
        return self._result_url

    @property
    def result_schema(self):
        """
        Returns: an array of array represents the type of result columns (Hive specific) (e.g. [["_c1", "string"], ["_c2", "bigint"]])
        """
        return self._hive_result_schema

    @property
    def priority(self):
        """
        Returns: a string represents the priority of the job (e.g. "NORMAL", "HIGH", etc.)
        """
        if self._priority in self.JOB_PRIORITY:
            return self.JOB_PRIORITY[self._priority]
        else:
            # just convert the value in string
            return str(self._priority)

    @property
    def retry_limit(self):
        """
        TODO: add docstring
        """
        return self._retry_limit

    @property
    def org_name(self):
        """
        Returns: organization name
        """
        return self._org_name

    @property
    def db_name(self):
        """
        Returns: a string represents the name of a database that job is running on
        """
        return self._db_name

    def wait(self, timeout=None, wait_interval=1):
        """Sleep until the job has been finished

        Params:
            timeout (int): Timeout in seconds. No timeout by default.
        """
        started_at = time.time()
        while not self.finished():
            if timeout is None or abs(time.time() - started_at) < timeout:
                time.sleep(wait_interval)
            else:
                raise RuntimeError("timeout") # TODO: throw proper error
        self.update()

    def kill(self):
        """Kill the job

        Returns: a string represents the status of killed job ("queued", "running")
        """
        response = self._client.kill(self.job_id)
        self.update()
        return response

    @property
    def query(self):
        """
        Returns: a string represents the query string of the job
        """
        return self._query

    def status(self):
        """
        Returns: a string represents the status of the job ("success", "error", "killed", "queued", "running")
        """
        if self._query is not None and not self.finished():
            self.update()
        return self._status

    @property
    def url(self):
        """
        Returns: a string of URL of the job on Treasure Data Service
        """
        return self._url

    def result(self):
        """
        Returns: an iterator of rows in result set
        """
        if not self.finished():
            raise ValueError("result is not ready")
        else:
            self.update()
            if self._result is None:
                for row in self._client.job_result_each(self._job_id):
                    yield row
            else:
                for row in self._result:
                    yield row

    def result_format(self, format):
        """
        Params:
            format (str): output format of result set

        Returns: an iterator of rows in result set
        """
        if not self.finished():
            raise ValueError("result is not ready")
        else:
            self.update()
            if self._result is None:
                for row in self._client.job_result_format_each(self._job_id, format):
                    yield row
            else:
                for row in self._result:
                    yield row

    def finished(self):
        """
        Returns: `True` if the job has been finished in success, error or killed
        """
        self._update_progress()
        return self._status in self.FINISHED_STATUS

    def success(self):
        """
        Returns: `True` if the job has been finished in success
        """
        self._update_progress()
        return self._status == self.STATUS_SUCCESS

    def error(self):
        """
        Returns: `True` if the job has been finished in error
        """
        self._update_progress()
        return self._status == self.STATUS_ERROR

    def killed(self):
        """
        Returns: `True` if the job has been finished in killed
        """
        self._update_progress()
        return self._status == self.STATUS_KILLED

    def queued(self):
        """
        Returns: `True` if the job is queued
        """
        self._update_progress()
        return self._status == self.STATUS_QUEUED

    def running(self):
        """
        Returns: `True` if the job is running
        """
        self._update_progress()
        return self._status == self.STATUS_RUNNING
<|endoftext|>