<commit_before>from itertools import chain
from sqlalchemy import case, create_engine, func, text, BigInteger, Boolean, Column, DateTime, Float, Integer, MetaData, String, Table
from sqlalchemy.dialects.postgresql import ARRAY, INT8RANGE
from sqlalchemy.sql.expression import and_, or_, between, join, label
from geoalchemy2 import Geometry
from geoalchemy2.functions import GenericFunction, ST_Intersects, ST_X, ST_Y
import datetime
import itertools
import re
import time
import piecewise.maxmind

def aggregate(config, debug=False):
    from sqlalchemy.schema import CreateTable
    engine = create_engine(config.database_uri)
    metadata = MetaData()
    records = config.make_cache_table(metadata)
    for agg in config.aggregations:
        if not debug:
            agg.build_aggregate_table(engine, metadata, records)
        else:
            query = agg.build_aggregate_query(engine, metadata, records)
            print query.compile(engine)

class ST_MakeBox2D(GenericFunction):
    name = 'ST_MakeBox2D'
    type = Geometry

class ST_Point(GenericFunction):
    name = 'ST_Point'
    type = Geometry

class ST_SnapToGrid(GenericFunction):
    name = 'ST_SnapToGrid'
    type = Geometry

class Aggregation(object):
    def __init__(self, name, statistics_table_name, bins, statistics):
        self.name = name
        self.statistics_table_name = statistics_table_name
        self.bins = bins
        self.statistics = statistics

    def build_aggregate_query(self, engine, metadata, records):
        statistics_table = self.make_table(metadata)
        statistics_table.create(engine, checkfirst = True)

        selection = records.select().with_only_columns([])
        columns = []
        for b in self.bins:
            cols, selection = b.build_query_to_populate(selection, records, statistics_table)
            columns += cols
        for s in self.statistics:
            cols, selection = s.build_query_to_populate(selection, records, statistics_table)
            columns += cols

        return statistics_table.insert().from_select(columns, selection)

    def build_aggregate_table(self, engine, metadata, records):
        query = self.build_aggregate_query(engine, metadata, records)

        with engine.begin() as conn:
            conn.execute(query)

    def make_table(self, metadata):
        bin_columns = chain.from_iterable(b.postgres_columns for b in self.bins)
        stat_columns = chain.from_iterable(b.postgres_columns for b in self.statistics)
        return Table(self.statistics_table_name, metadata,
                Column('id', Integer, primary_key = True),
                *list(chain(bin_columns, stat_columns)),
                keep_existing = True)

    def selection(self, engine, metadata, bins, filters, statistics):
        table = self.make_table(metadata)

        bin_keys = []
        filter_predicates = []

        selection = table.select().with_only_columns([])
        for b in self.bins:
            if b.label in bins:
                selection = b.build_query_to_report(selection, table, bins[b.label])
            if b.label in filters:
                selection = b.filter_query_to_report(selection, table, filters[b.label])
        for s in statistics:
            selection = s.build_query_to_report(selection, table)

        return selection

    def query(self, engine, metadata, bins, filters, statistics):
        selection = self.selection(engine, metadata, bins, filters, statistics)

        with engine.connect() as conn:
            return conn.execute(selection)

    def __repr__(self):
        return "{name}: table={table} bins={bins} stats={stats}".format(name=self.name, table=self.statistics_table_name, bins=self.bins, stats=self.statistics)


class Aggregator(object):
    def __init__(self, database_uri, cache_table_name, filters, aggregations):
        self.database_uri = database_uri
        self.cache_table_name = cache_table_name
        self.filters = filters
        self.aggregations = aggregations

    def make_cache_table(self, metadata):
        return Table(self.cache_table_name, metadata, 
                Column('id', BigInteger, primary_key = True),
                Column('time', DateTime),
                Column('location', Geometry("POINT", srid=4326)),
                Column('client_ip', BigInteger),
                Column('server_ip', BigInteger),
                Column('countrtt', BigInteger),
                Column('sumrtt', BigInteger),
                Column('download_flag', Boolean),
                Column('download_time', BigInteger),
                Column('download_octets', BigInteger),
                Column('upload_time', BigInteger),
                Column('upload_octets', BigInteger),
                Column('bigquery_key', String))

    def ingest_bigquery_query(self):
        select_clause = [
                "web100_log_entry.log_time AS time",
                "connection_spec.client_geolocation.longitude AS longitude",
                "connection_spec.client_geolocation.latitude AS latitude",
                "PARSE_IP(connection_spec.client_ip) AS client_ip",
                "PARSE_IP(connection_spec.server_ip) AS server_ip",
                "web100_log_entry.snap.CountRTT AS countrtt",
                "web100_log_entry.snap.SumRTT AS sumrtt",
                "connection_spec.data_direction AS data_direction",
                "web100_log_entry.snap.SndLimTimeRwin + web100_log_entry.snap.SndLimTimeCwnd + web100_log_entry.snap.SndLimTimeSnd AS download_time",
                "8 * web100_log_entry.snap.HCThruOctetsAcked AS download_octets",
                "web100_log_entry.snap.Duration AS upload_time",
                "8 * web100_log_entry.snap.HCThruOctetsReceived AS upload_octets",
                "CONCAT(String(web100_log_entry.snap.Duration), String(web100_log_entry.snap.CountRTT), String(web100_log_entry.snap.SegsIn) , String(web100_log_entry.snap.SegsOut)) AS bigquery_key"
        ]

        where_clause = list(chain.from_iterable(f.bigquery_filter() for f in self.filters))
        tables = self._tables_for(self.filters)

        return '\n'.join([
            'SELECT',
            ', '.join(select_clause),
            'FROM ',
            ', '.join('[%s]'% t for t in tables),
            'WHERE',
            ' AND '.join(where_clause),
            ';'])

    def _tables_for(self, filters):
        all_tables = []
        current_year_month = time.gmtime()[:2] # first two fields are year and month, handy.
        years = itertools.count(2010)
        months = range(1, 12 + 1)
        possible_year_months = ((year, month) for year in years for month in months)
        valid_year_months = itertools.takewhile(lambda x: x <= current_year_month, possible_year_months)

        for f in filters:
            if isinstance(f, TemporalFilter):
                after = time.gmtime(f.after)[:2]
                before = time.gmtime(f.before)[:2]
                valid_year_months = itertools.ifilter(lambda x: after <= x <= before, valid_year_months)

        return ['plx.google:m_lab.%04d_%02d.all' % ym for ym in valid_year_months]

    def bigquery_row_to_postgres_row(self, bq_row):
        bq_row = [f['v'] for f in bq_row['f']]
        timestamp, longitude, latitude, client_ip, server_ip, countrtt, sumrtt, direction, dl_time, dl_octets, ul_time, ul_octets, bigquery_key = bq_row

        timestamp = timestamp and datetime.datetime.utcfromtimestamp(float(timestamp))
        location = longitude and latitude and "srid=4326;POINT(%f %f)" % (float(longitude), float(latitude))
        client_ip = client_ip and int(client_ip)
        server_ip = server_ip and int(server_ip)
        countrtt = countrtt and int(countrtt)
        sumrtt = sumrtt and int(sumrtt)
        direction = direction and bool(int(direction))
        dl_time = dl_time and float(dl_time)
        dl_octets = dl_octets and int(dl_octets)
        ul_time = ul_time and float(ul_time)
        ul_octets = ul_octets and int(ul_octets)
        bigquery_key = bigquery_key

        pg_row = dict(
                time=timestamp,
                location=location,
                client_ip = client_ip,
                server_ip = server_ip,
                countrtt=countrtt,
                sumrtt = sumrtt,
                download_flag = direction,
                download_time = dl_time,
                download_octets = dl_octets,
                upload_time = ul_time,
                upload_octets = ul_octets,
                bigquery_key = bigquery_key)

        return pg_row

class Bins(object):
    pass

class SpatialGridBins(Bins):
    label = "spatial_grid"

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

    @property<commit_msg>Add test_id field from BigQuery to the results table.<commit_after>from itertools import chain
from sqlalchemy import case, create_engine, func, text, BigInteger, Boolean, Column, DateTime, Float, Integer, MetaData, String, Table
from sqlalchemy.dialects.postgresql import ARRAY, INT8RANGE
from sqlalchemy.sql.expression import and_, or_, between, join, label
from geoalchemy2 import Geometry
from geoalchemy2.functions import GenericFunction, ST_Intersects, ST_X, ST_Y
import datetime
import itertools
import re
import time
import piecewise.maxmind

def aggregate(config, debug=False):
    from sqlalchemy.schema import CreateTable
    engine = create_engine(config.database_uri)
    metadata = MetaData()
    records = config.make_cache_table(metadata)
    for agg in config.aggregations:
        if not debug:
            agg.build_aggregate_table(engine, metadata, records)
        else:
            query = agg.build_aggregate_query(engine, metadata, records)
            print query.compile(engine)

class ST_MakeBox2D(GenericFunction):
    name = 'ST_MakeBox2D'
    type = Geometry

class ST_Point(GenericFunction):
    name = 'ST_Point'
    type = Geometry

class ST_SnapToGrid(GenericFunction):
    name = 'ST_SnapToGrid'
    type = Geometry

class Aggregation(object):
    def __init__(self, name, statistics_table_name, bins, statistics):
        self.name = name
        self.statistics_table_name = statistics_table_name
        self.bins = bins
        self.statistics = statistics

    def build_aggregate_query(self, engine, metadata, records):
        statistics_table = self.make_table(metadata)
        statistics_table.create(engine, checkfirst = True)

        selection = records.select().with_only_columns([])
        columns = []
        for b in self.bins:
            cols, selection = b.build_query_to_populate(selection, records, statistics_table)
            columns += cols
        for s in self.statistics:
            cols, selection = s.build_query_to_populate(selection, records, statistics_table)
            columns += cols

        return statistics_table.insert().from_select(columns, selection)

    def build_aggregate_table(self, engine, metadata, records):
        query = self.build_aggregate_query(engine, metadata, records)

        with engine.begin() as conn:
            conn.execute(query)

    def make_table(self, metadata):
        bin_columns = chain.from_iterable(b.postgres_columns for b in self.bins)
        stat_columns = chain.from_iterable(b.postgres_columns for b in self.statistics)
        return Table(self.statistics_table_name, metadata,
                Column('id', Integer, primary_key = True),
                *list(chain(bin_columns, stat_columns)),
                keep_existing = True)

    def selection(self, engine, metadata, bins, filters, statistics):
        table = self.make_table(metadata)

        bin_keys = []
        filter_predicates = []

        selection = table.select().with_only_columns([])
        for b in self.bins:
            if b.label in bins:
                selection = b.build_query_to_report(selection, table, bins[b.label])
            if b.label in filters:
                selection = b.filter_query_to_report(selection, table, filters[b.label])
        for s in statistics:
            selection = s.build_query_to_report(selection, table)

        return selection

    def query(self, engine, metadata, bins, filters, statistics):
        selection = self.selection(engine, metadata, bins, filters, statistics)

        with engine.connect() as conn:
            return conn.execute(selection)

    def __repr__(self):
        return "{name}: table={table} bins={bins} stats={stats}".format(name=self.name, table=self.statistics_table_name, bins=self.bins, stats=self.statistics)


class Aggregator(object):
    def __init__(self, database_uri, cache_table_name, filters, aggregations):
        self.database_uri = database_uri
        self.cache_table_name = cache_table_name
        self.filters = filters
        self.aggregations = aggregations

    def make_cache_table(self, metadata):
        return Table(self.cache_table_name, metadata, 
                Column('id', BigInteger, primary_key = True),
                Column('time', DateTime),
                Column('location', Geometry("POINT", srid=4326)),
                Column('client_ip', BigInteger),
                Column('server_ip', BigInteger),
                Column('countrtt', BigInteger),
                Column('sumrtt', BigInteger),
                Column('download_flag', Boolean),
                Column('download_time', BigInteger),
                Column('download_octets', BigInteger),
                Column('upload_time', BigInteger),
                Column('upload_octets', BigInteger),
                Column('bigquery_key', String),
                Column('test_id', String))

    def ingest_bigquery_query(self):
        select_clause = [
                "web100_log_entry.log_time AS time",
                "connection_spec.client_geolocation.longitude AS longitude",
                "connection_spec.client_geolocation.latitude AS latitude",
                "PARSE_IP(connection_spec.client_ip) AS client_ip",
                "PARSE_IP(connection_spec.server_ip) AS server_ip",
                "web100_log_entry.snap.CountRTT AS countrtt",
                "web100_log_entry.snap.SumRTT AS sumrtt",
                "connection_spec.data_direction AS data_direction",
                "web100_log_entry.snap.SndLimTimeRwin + web100_log_entry.snap.SndLimTimeCwnd + web100_log_entry.snap.SndLimTimeSnd AS download_time",
                "8 * web100_log_entry.snap.HCThruOctetsAcked AS download_octets",
                "web100_log_entry.snap.Duration AS upload_time",
                "8 * web100_log_entry.snap.HCThruOctetsReceived AS upload_octets",
                "CONCAT(String(web100_log_entry.snap.Duration), String(web100_log_entry.snap.CountRTT), String(web100_log_entry.snap.SegsIn), String(web100_log_entry.snap.SegsOut)) AS bigquery_key",
                "test_id"
        ]

        where_clause = list(chain.from_iterable(f.bigquery_filter() for f in self.filters))
        tables = self._tables_for(self.filters)

        return '\n'.join([
            'SELECT',
            ', '.join(select_clause),
            'FROM ',
            ', '.join('[%s]'% t for t in tables),
            'WHERE',
            ' AND '.join(where_clause),
            ';'])

    def _tables_for(self, filters):
        all_tables = []
        current_year_month = time.gmtime()[:2] # first two fields are year and month, handy.
        years = itertools.count(2010)
        months = range(1, 12 + 1)
        possible_year_months = ((year, month) for year in years for month in months)
        valid_year_months = itertools.takewhile(lambda x: x <= current_year_month, possible_year_months)

        for f in filters:
            if isinstance(f, TemporalFilter):
                after = time.gmtime(f.after)[:2]
                before = time.gmtime(f.before)[:2]
                valid_year_months = itertools.ifilter(lambda x: after <= x <= before, valid_year_months)

        return ['plx.google:m_lab.%04d_%02d.all' % ym for ym in valid_year_months]

    def bigquery_row_to_postgres_row(self, bq_row):
        bq_row = [f['v'] for f in bq_row['f']]
        timestamp, longitude, latitude, client_ip, server_ip, countrtt, sumrtt, direction, dl_time, dl_octets, ul_time, ul_octets, bigquery_key, test_id = bq_row

        timestamp = timestamp and datetime.datetime.utcfromtimestamp(float(timestamp))
        location = longitude and latitude and "srid=4326;POINT(%f %f)" % (float(longitude), float(latitude))
        client_ip = client_ip and int(client_ip)
        server_ip = server_ip and int(server_ip)
        countrtt = countrtt and int(countrtt)
        sumrtt = sumrtt and int(sumrtt)
        direction = direction and bool(int(direction))
        dl_time = dl_time and float(dl_time)
        dl_octets = dl_octets and int(dl_octets)
        ul_time = ul_time and float(ul_time)
        ul_octets = ul_octets and int(ul_octets)
        bigquery_key = bigquery_key
        test_id = test_id

        pg_row = dict(
                time=timestamp,
                location=location,
                client_ip = client_ip,
                server_ip = server_ip,
                countrtt=countrtt,
                sumrtt = sumrtt,
                download_flag = direction,
                download_time = dl_time,
                download_octets = dl_octets,
                upload_time = ul_time,
                upload_octets = ul_octets,
                bigquery_key = bigquery_key,
                test_id = test_id)

        return pg_row

class Bins(object):
    pass

class SpatialGridBins(Bins):
    label = "spatial_grid"

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

    @property
{'fraction_non_alphanumeric': 0.09171815656708804, 'fraction_numerical': 0.007822685788787484, 'mean_word_length': 2.6114636642784035, 'pattern_counts': {'":': 0, '<': 9, '<?xml version=': 0, '>': 3, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 16, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '7464935', 'n_tokens_mistral': 5384, 'n_tokens_neox': 5042, 'n_words': 1141}