<commit_before>#!/usr/bin/env python
# -*- mode: python; coding: utf-8 -*-
# Copyright 2022 the HERA Collaboration
# Licensed under the 2-clause BSD license.
"""Accumulate SNAP spectra from redis and write to a UVH5 file."""
import time
import json
import logging
import argparse
import traceback
from pathlib import Path

import numpy as np
from pyuvdata import UVData
from pyuvdata import utils as uvutils

from astropy.time import Time, TimeDelta

from hera_corr_cm import HeraCorrCM
from hera_corr_cm.redis_cm import read_maps_from_redis, read_cminfo_from_redis

formatter = "%(asctime)s.%(msecs)03d %(levelname)s - %(module)s: %(message)s"
logging.basicConfig(
    level=logging.INFO,
    format=formatter,
    datefmt="%Y-%m-%d %H:%M:%S",
)
logger = logging.getLogger(__file__)
logger.ident = "hera_mini_snap_catcher"

syslog_handler = logging.handlers.SysLogHandler(address="/dev/log")
syslog_handler.setFormatter(
    logging.Formatter(
        "%(module)s[%(process)s]: %(message)s",
        "%FT%H:%m:%s%z",
    )
)
logger.addHandler(syslog_handler)

def get_blt_index(antnum, time, ant_array, time_array):
    indices = np.nonzero(np.asarray(ant_array) == antnum)[0]

    if indices.size != 0:
        sub_inds = np.nonzero(np.asarray(time_array)[indices] == time)[0]
        if sub_inds.size != 0:
            return indices[sub_inds]
        else:
            return None
    else:
        return None

parser = argparse.ArgumentParser(
    description="Create UVH5 file from snap autocorrelations in Redis.",
    formatter_class=argparse.ArgumentDefaultsHelpFormatter,
)
parser.add_argument(
    "-r",
    dest="redishost",
    type=str,
    default="redishost",
    help="Host servicing redis requests",
)
parser.add_argument(
    "--max_downtime",
    dest="max_downtime",
    type=float,
    default=60.0,
    help=(
        "Seconds between to wait between redis updates "
        "before restarting with new file."
    ),
)

parser.add_argument(
    "--max_file_len",
    dest="max_file_len",
    type=float,
    default=60.0,
    help=(
        "Maximum length of a UVH5 file in seconds. "
        "This is the total time between the first and last entry in the time_array."
    ),
)

parser.add_argument(
    "--outdir",
    "-o",
    dest="outdir",
    type=str,
    default=".",
    help="The directory to write files to.",
)

args = parser.parse_args()

logger.info("Starting HERA snap mini data catcher.")
while True:
    try:
        corr_cm = HeraCorrCM(redishost=args.redishost)
        logger.info("HERA CorrCM Redis connection established.")

        mapping = read_maps_from_redis()
        snap_to_ant_mapping = mapping["snap_to_ant"]
        if mapping is None:
            raise ValueError(
                "None value encountered for antenna/snap mappings in redis."
            )

        last_time_mapping = {}
        last_loop_completion = Time.now()

        logger.info("Beginning data catching.")
        while True:
            try:

                # use future_array_shapes
                # Nblts, Nfreqs, Npols
                # we don't necessarily know how many freqs we'll have.
                # it is most likely 1024 but this will guard against X-Engine changes.
                data_array = None
                time_array = []
                ant_array = []

                downtime = 0
                file_len = 0

                while downtime < args.max_downtime and file_len < args.max_file_len:
                    time.sleep(0.01)
                    # get integrations from HeraCorrCM.get_snaprf_status()
                    # compare to the last time for each snap

                    for (
                        snap_hostname,
                        status,
                    ) in corr_cm.get_snaprf_status().items():

                        if status["timestamp"] is None or status["timestamp"] == "None":
                            continue

                        snap, stream = snap_hostname.split(":")
                        stream = int(stream)

                        hera_ant_num = snap_to_ant_mapping[snap][stream]
                        if hera_ant_num is None:
                            continue

                        # arbitrarily assign east as polarization 0
                        pol = hera_ant_num[-1].lower()
                        pol_ind = 0 if pol == "e" else 1

                        # trim off the HH and the Polarization
                        hera_ant_num = int(hera_ant_num[2:-1])

                        timestamp = Time(status["timestamp"]).jd

                        if snap_hostname not in last_time_mapping:
                            last_time_mapping[snap_hostname] = timestamp
                        else:
                            if timestamp <= last_time_mapping[snap_hostname]:
                                # Time has not progressed yet.
                                # Just sit tight
                                continue
                            last_time_mapping[snap_hostname] = timestamp

                        auto_corr = status["autocorrelation"]

                        if data_array is None:
                            data_array = np.zeros(
                                (0, auto_corr.shape[0], 2), dtype=np.float32
                            )
                        data_index = get_blt_index(
                            hera_ant_num, timestamp, ant_array, time_array
                        )
                        if data_index is not None:
                            # this baseline is already in the time array
                            # insert into the correct polarization
                            # at the given index
                            data_array[data_index, :, pol_ind] = auto_corr

                        else:
                            # otherwise append the data array and insert the meta-data
                            data_insert = np.zeros(
                                (1, auto_corr.shape[0], 2), dtype=np.float32
                            )
                            data_insert[..., pol_ind] = auto_corr

                            data_array = np.append(data_array, data_insert, axis=0)
                            time_array.append(timestamp)
                            ant_array.append(hera_ant_num)

                    last_time = Time(max(last_time_mapping.values()), format="jd")
                    # Take the least time between the last snap update
                    # and the last time a file was written
                    downtime = min(
                        (Time.now() - last_time).to_value("s"),
                        (Time.now() - last_loop_completion).to_value("s"),
                    )
                    if len(time_array) > 0:
                        file_len = TimeDelta(
                            time_array[-1] - time_array[0], format="jd"
                        ).to_value("s")

                time_array = np.asarray(time_array)
                if time_array.size == 0:
                    # No Data was taken
                    logger.info(
                        f"Downtime timeout. No new data recieved since {last_time.jd} "
                        f"({last_time.iso} UTC)."
                    )
                    continue

                ant_array = np.asarray(ant_array)

                uvd = UVData()

                uvd.data_array = data_array.astype(np.complex64)
                uvd.vis_units = "UNCALIB"
                uvd.time_array = time_array

                uvd.Nblts = time_array.shape[0]
                uvd.Ntimes = np.unique(uvd.time_array).size

                uvd.Nfreqs = data_array.shape[1]
                uvd.freq_array = np.linspace(0, 256, uvd.Nfreqs) * 1e6

                uvd.channel_width = np.full_like(
                    uvd.freq_array, np.diff(uvd.freq_array)[0]
                )
                uvd.Nants_data = np.unique(ant_array).size
                uvd.Nbls = uvd.Nants_data
                uvd.Nspws = 1
                uvd.Npols = 2

                uvd.ant_1_array = ant_array
                uvd.ant_2_array = ant_array
                uvd.baseline_array = uvd.antnums_to_baseline(
                    uvd.ant_1_array, uvd.ant_2_array
                )

                uvd.nsample_array = np.ones_like(uvd.data_array, np.float32)
                uvd.integration_time = np.ones((uvd.Nblts), np.float32)
                uvd.flag_array = np.zeros_like(uvd.data_array, dtype=bool)
                uvd.x_orientation = "north"
                uvd.Nants_data = len(set(ant_array))
                uvd.polarization_array = uvutils.polstr2num(
                    ["ee", "nn"], x_orientation=uvd.x_orientation
                )
                uvd.spw_array = np.array([0])

                cminfo = read_cminfo_from_redis(return_as="dict")

                uvd.antenna_positions = np.asarray(cminfo["antenna_positions"])
                uvd.antenna_names = np.asarray(cminfo["antenna_names"])
                uvd.antenna_numbers = np.asarray(cminfo["antenna_numbers"])
                uvd.Nants_telescope = len(uvd.antenna_numbers)
                lat, lon, alt = (
                    cminfo["cofa_lat"],
                    cminfo["cofa_lon"],
                    cminfo["cofa_alt"],
                )
                uvd.telescope_location_lat_lon_alt_degrees = (lat, lon, alt)
                uvd.set_uvws_from_antenna_positions()
                uvd.set_lsts_from_time_array()

                uvd.object_name = ""
                uvd.multi_phase_center = False
                uvd.telescope_name = "HERA"
                uvd.instrument = "HERA"
                uvd.history = f"Created by {__file__}."
                uvd.phase_type = "drift"

                # lets us not have to spectral windows
                uvd._set_future_array_shapes()

                uvd.reorder_blts("time", "baseline")
                uvd.extra_keywords["snap_to_ant_mapping"] = json.dumps(
                    snap_to_ant_mapping
                )
                fname = f"zen.{Time.now().jd:.6f}.snap_autos.uvh5"
                logger.info(f"Writing output file {fname}.")
                uvd.write_uvh5(Path(args.outdir) / fname)

                last_loop_completion = Time.now()

            except Exception:
                logger.error(
                    "{t} -- error while accumulating datafile".format(t=time.asctime()),
                )
                logger.error(traceback.print_exc())
                continue

    except Exception:
        logger.error(traceback.print_exc())
        continue
<commit_msg>write mini_catcher files to a separate folder by JD<commit_after>#!/usr/bin/env python
# -*- mode: python; coding: utf-8 -*-
# Copyright 2022 the HERA Collaboration
# Licensed under the 2-clause BSD license.
"""Accumulate SNAP spectra from redis and write to a UVH5 file."""
import time
import json
import logging
import argparse
import traceback
from pathlib import Path

import numpy as np
from pyuvdata import UVData
from pyuvdata import utils as uvutils

from astropy.time import Time, TimeDelta

from hera_corr_cm import HeraCorrCM
from hera_corr_cm.redis_cm import read_maps_from_redis, read_cminfo_from_redis

formatter = "%(asctime)s.%(msecs)03d %(levelname)s - %(module)s: %(message)s"
logging.basicConfig(
    level=logging.INFO,
    format=formatter,
    datefmt="%Y-%m-%d %H:%M:%S",
)
logger = logging.getLogger(__file__)
logger.ident = "hera_mini_snap_catcher"

syslog_handler = logging.handlers.SysLogHandler(address="/dev/log")
syslog_handler.setFormatter(
    logging.Formatter(
        "%(module)s[%(process)s]: %(message)s",
        "%FT%H:%m:%s%z",
    )
)
logger.addHandler(syslog_handler)

def get_blt_index(antnum, time, ant_array, time_array):
    indices = np.nonzero(np.asarray(ant_array) == antnum)[0]

    if indices.size != 0:
        sub_inds = np.nonzero(np.asarray(time_array)[indices] == time)[0]
        if sub_inds.size != 0:
            return indices[sub_inds]
        else:
            return None
    else:
        return None

parser = argparse.ArgumentParser(
    description="Create UVH5 file from snap autocorrelations in Redis.",
    formatter_class=argparse.ArgumentDefaultsHelpFormatter,
)
parser.add_argument(
    "-r",
    dest="redishost",
    type=str,
    default="redishost",
    help="Host servicing redis requests",
)
parser.add_argument(
    "--max_downtime",
    dest="max_downtime",
    type=float,
    default=60.0,
    help=(
        "Seconds between to wait between redis updates "
        "before restarting with new file."
    ),
)

parser.add_argument(
    "--max_file_len",
    dest="max_file_len",
    type=float,
    default=60.0,
    help=(
        "Maximum length of a UVH5 file in seconds. "
        "This is the total time between the first and last entry in the time_array."
    ),
)

parser.add_argument(
    "--outdir",
    "-o",
    dest="outdir",
    type=str,
    default=".",
    help="The directory to write files to.",
)

args = parser.parse_args()

logger.info("Starting HERA snap mini data catcher.")
while True:
    try:
        corr_cm = HeraCorrCM(redishost=args.redishost)
        logger.info("HERA CorrCM Redis connection established.")

        mapping = read_maps_from_redis()
        snap_to_ant_mapping = mapping["snap_to_ant"]
        if mapping is None:
            raise ValueError(
                "None value encountered for antenna/snap mappings in redis."
            )

        last_time_mapping = {}
        last_loop_completion = Time.now()

        logger.info("Beginning data catching.")
        while True:
            try:

                # use future_array_shapes
                # Nblts, Nfreqs, Npols
                # we don't necessarily know how many freqs we'll have.
                # it is most likely 1024 but this will guard against X-Engine changes.
                data_array = None
                time_array = []
                ant_array = []

                downtime = 0
                file_len = 0

                while downtime < args.max_downtime and file_len < args.max_file_len:
                    time.sleep(0.01)
                    # get integrations from HeraCorrCM.get_snaprf_status()
                    # compare to the last time for each snap

                    for (
                        snap_hostname,
                        status,
                    ) in corr_cm.get_snaprf_status().items():

                        if status["timestamp"] is None or status["timestamp"] == "None":
                            continue

                        snap, stream = snap_hostname.split(":")
                        stream = int(stream)

                        hera_ant_num = snap_to_ant_mapping[snap][stream]
                        if hera_ant_num is None:
                            continue

                        # arbitrarily assign east as polarization 0
                        pol = hera_ant_num[-1].lower()
                        pol_ind = 0 if pol == "e" else 1

                        # trim off the HH and the Polarization
                        hera_ant_num = int(hera_ant_num[2:-1])

                        timestamp = Time(status["timestamp"]).jd

                        if snap_hostname not in last_time_mapping:
                            last_time_mapping[snap_hostname] = timestamp
                        else:
                            if timestamp <= last_time_mapping[snap_hostname]:
                                # Time has not progressed yet.
                                # Just sit tight
                                continue
                            last_time_mapping[snap_hostname] = timestamp

                        auto_corr = status["autocorrelation"]

                        if data_array is None:
                            data_array = np.zeros(
                                (0, auto_corr.shape[0], 2), dtype=np.float32
                            )
                        data_index = get_blt_index(
                            hera_ant_num, timestamp, ant_array, time_array
                        )
                        if data_index is not None:
                            # this baseline is already in the time array
                            # insert into the correct polarization
                            # at the given index
                            data_array[data_index, :, pol_ind] = auto_corr

                        else:
                            # otherwise append the data array and insert the meta-data
                            data_insert = np.zeros(
                                (1, auto_corr.shape[0], 2), dtype=np.float32
                            )
                            data_insert[..., pol_ind] = auto_corr

                            data_array = np.append(data_array, data_insert, axis=0)
                            time_array.append(timestamp)
                            ant_array.append(hera_ant_num)

                    last_time = Time(max(last_time_mapping.values()), format="jd")
                    # Take the least time between the last snap update
                    # and the last time a file was written
                    downtime = min(
                        (Time.now() - last_time).to_value("s"),
                        (Time.now() - last_loop_completion).to_value("s"),
                    )
                    if len(time_array) > 0:
                        file_len = TimeDelta(
                            time_array[-1] - time_array[0], format="jd"
                        ).to_value("s")

                time_array = np.asarray(time_array)
                if time_array.size == 0:
                    # No Data was taken
                    logger.info(
                        f"Downtime timeout. No new data recieved since {last_time.jd} "
                        f"({last_time.iso} UTC)."
                    )
                    continue

                ant_array = np.asarray(ant_array)

                uvd = UVData()

                uvd.data_array = data_array.astype(np.complex64)
                uvd.vis_units = "UNCALIB"
                uvd.time_array = time_array

                uvd.Nblts = time_array.shape[0]
                uvd.Ntimes = np.unique(uvd.time_array).size

                uvd.Nfreqs = data_array.shape[1]
                uvd.freq_array = np.linspace(0, 256, uvd.Nfreqs) * 1e6

                uvd.channel_width = np.full_like(
                    uvd.freq_array, np.diff(uvd.freq_array)[0]
                )
                uvd.Nants_data = np.unique(ant_array).size
                uvd.Nbls = uvd.Nants_data
                uvd.Nspws = 1
                uvd.Npols = 2

                uvd.ant_1_array = ant_array
                uvd.ant_2_array = ant_array
                uvd.baseline_array = uvd.antnums_to_baseline(
                    uvd.ant_1_array, uvd.ant_2_array
                )

                uvd.nsample_array = np.ones_like(uvd.data_array, np.float32)
                uvd.integration_time = np.ones((uvd.Nblts), np.float32)
                uvd.flag_array = np.zeros_like(uvd.data_array, dtype=bool)
                uvd.x_orientation = "north"
                uvd.Nants_data = len(set(ant_array))
                uvd.polarization_array = uvutils.polstr2num(
                    ["ee", "nn"], x_orientation=uvd.x_orientation
                )
                uvd.spw_array = np.array([0])

                cminfo = read_cminfo_from_redis(return_as="dict")

                uvd.antenna_positions = np.asarray(cminfo["antenna_positions"])
                uvd.antenna_names = np.asarray(cminfo["antenna_names"])
                uvd.antenna_numbers = np.asarray(cminfo["antenna_numbers"])
                uvd.Nants_telescope = len(uvd.antenna_numbers)
                lat, lon, alt = (
                    cminfo["cofa_lat"],
                    cminfo["cofa_lon"],
                    cminfo["cofa_alt"],
                )
                uvd.telescope_location_lat_lon_alt_degrees = (lat, lon, alt)
                uvd.set_uvws_from_antenna_positions()
                uvd.set_lsts_from_time_array()

                uvd.object_name = ""
                uvd.multi_phase_center = False
                uvd.telescope_name = "HERA"
                uvd.instrument = "HERA"
                uvd.history = f"Created by {__file__}."
                uvd.phase_type = "drift"

                # lets us not have to spectral windows
                uvd._set_future_array_shapes()

                uvd.reorder_blts("time", "baseline")
                uvd.extra_keywords["snap_to_ant_mapping"] = json.dumps(
                    snap_to_ant_mapping
                )

                date = Time.now().jd

                folder_name = f"{int(date.floor()):d}"

                # make the fold for the JD if it does not exist
                folder = Path(args.outdir) / folder_name
                folder.mkdir(exists_ok=True)

                fname = folder / f"zen.{date:.6f}.snap_autos.uvh5"
                logger.info(f"Writing output file {fname}.")
                uvd.write_uvh5(fname)

                last_loop_completion = Time.now()

            except Exception:
                logger.error(
                    "{t} -- error while accumulating datafile".format(t=time.asctime()),
                )
                logger.error(traceback.print_exc())
                continue

    except Exception:
        logger.error(traceback.print_exc())
        continue
<|endoftext|>


<commit_before>from itertools import chain
from torch.nn import DataParallel
import pdb

class LightningDataParallel(DataParallel):
    """
    Override the forward call in lightning so it goes to training and validation step respectively
    """

    def forward(self, *inputs, **kwargs):
        if not self.device_ids:
            # -------------
            # MAIN CHANGE
            if self.module.training:
                return self.module.training_step(*inputs, **kwargs)
            else:
                return self.module.validation_step(*inputs, **kwargs)
            # -------------

        for t in chain(self.module.parameters(), self.module.buffers()):
            if t.device != self.src_device_obj:
                raise RuntimeError("module must have its parameters and buffers "
                                   "on device {} (device_ids[0]) but found one of "
                                   "them on device: {}".format(self.src_device_obj, t.device))

        inputs, kwargs = self.scatter(inputs, kwargs, self.device_ids)
        if len(self.device_ids) == 1:
            if self.module.training:
                return self.module.training_step(*inputs[0], **kwargs[0])
            else:
                return self.module.validation_step(*inputs[0], **kwargs[0])
        replicas = self.replicate(self.module, self.device_ids[:len(inputs)])
        outputs = self.parallel_apply(replicas, inputs, kwargs)
        return self.gather(outputs, self.output_device)
<commit_msg>updated args<commit_after>from itertools import chain
from torch.nn import DataParallel
import pdb

import threading
import torch
from torch.cuda._utils import _get_device_index

def get_a_var(obj):
    if isinstance(obj, torch.Tensor):
        return obj

    if isinstance(obj, list) or isinstance(obj, tuple):
        for result in map(get_a_var, obj):
            if isinstance(result, torch.Tensor):
                return result
    if isinstance(obj, dict):
        for result in map(get_a_var, obj.items()):
            if isinstance(result, torch.Tensor):
                return result
    return None

class LightningDataParallel(DataParallel):
    """
    Override the forward call in lightning so it goes to training and validation step respectively
    """

    def parallel_apply(self, replicas, inputs, kwargs):
        return parallel_apply(replicas, inputs, kwargs, self.device_ids[:len(replicas)])

def parallel_apply(modules, inputs, kwargs_tup=None, devices=None):
    r"""Applies each `module` in :attr:`modules` in parallel on arguments
    contained in :attr:`inputs` (positional) and :attr:`kwargs_tup` (keyword)
    on each of :attr:`devices`.

    Args:
        modules (Module): modules to be parallelized
        inputs (tensor): inputs to the modules
        devices (list of int or torch.device): CUDA devices

    :attr:`modules`, :attr:`inputs`, :attr:`kwargs_tup` (if given), and
    :attr:`devices` (if given) should all have same length. Moreover, each
    element of :attr:`inputs` can either be a single object as the only argument
    to a module, or a collection of positional arguments.
    """
    assert len(modules) == len(inputs)
    if kwargs_tup is not None:
        assert len(modules) == len(kwargs_tup)
    else:
        kwargs_tup = ({},) * len(modules)
    if devices is not None:
        assert len(modules) == len(devices)
    else:
        devices = [None] * len(modules)
    devices = list(map(lambda x: _get_device_index(x, True), devices))
    lock = threading.Lock()
    results = {}
    grad_enabled = torch.is_grad_enabled()

    def _worker(i, module, input, kwargs, device=None):
        torch.set_grad_enabled(grad_enabled)
        if device is None:
            device = get_a_var(input).get_device()
        try:
            with torch.cuda.device(device):
                # this also avoids accidental slicing of `input` if it is a Tensor
                if not isinstance(input, (list, tuple)):
                    input = (input,)

                if module.training:
                    return module.training_step(*input, **kwargs)
                else:
                    return module.validation_step(*input, **kwargs)
            with lock:
                results[i] = output
        except Exception as e:
            with lock:
                results[i] = e

    if len(modules) > 1:
        threads = [threading.Thread(target=_worker,
                                    args=(i, module, input, kwargs, device))
                   for i, (module, input, kwargs, device) in
                   enumerate(zip(modules, inputs, kwargs_tup, devices))]

        for thread in threads:
            thread.start()
        for thread in threads:
            thread.join()
    else:
        _worker(0, modules[0], inputs[0], kwargs_tup[0], devices[0])

    outputs = []
    for i in range(len(inputs)):
        output = results[i]
        if isinstance(output, Exception):
            raise output
        outputs.append(output)
    return outputs<|endoftext|>


<commit_before>import mailbox

from six import StringIO

from fulltext.backends.__eml import _get_file

def _get_path(path, **kwargs):
    text, mb = StringIO(), mailbox.mbox(path, create=False)

    for k in mb.keys():
        text.write(_get_file(mb.get_file(k)))
        text.write(u'\n\n')

    return text.getvalue()<commit_msg>Code quality.<commit_after>import mailbox

from six import StringIO

from fulltext.backends.__eml import _get_file

def _get_path(path, **kwargs):
    text, mb = StringIO(), mailbox.mbox(path, create=False)

    for k in mb.keys():
        text.write(_get_file(mb.get_file(k)))
        text.write(u'\n\n')

    return text.getvalue()
<|endoftext|>


<commit_before>#!/usr/bin/python
# -*- coding: utf-8 -*-

# Hive Netius System
# Copyright (C) 2008-2012 Hive Solutions Lda.
#
# This file is part of Hive Netius System.
#
# Hive Netius System is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# Hive Netius System is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with Hive Netius System. If not, see <http://www.gnu.org/licenses/>.

__author__ = "João Magalhães joamag@hive.pt>"
""" The author(s) of the module """

__version__ = "1.0.0"
""" The version of the module """

__revision__ = "$LastChangedRevision$"
""" The revision number of the module """

__date__ = "$LastChangedDate$"
""" The last change date of the module """

__copyright__ = "Copyright (c) 2008-2012 Hive Solutions Lda."
""" The copyright for the module """

__license__ = "GNU General Public License (GPL), Version 3"
""" The license for the module """

import http

import netius.common
import netius.clients

class ProxyServer(http.HTTPServer):

    def __init__(self, *args, **kwargs):
        http.HTTPServer.__init__(self, *args, **kwargs)
        self.conn_map = {}

        self.http_client = netius.clients.HTTPClient(
            thread = False,
            *args,
            **kwargs
        )
        self.http_client.bind("headers", self._on_prx_headers)
        self.http_client.bind("message", self._on_prx_message)
        self.http_client.bind("partial", self._on_prx_partial)
        self.http_client.bind("chunk", self._on_prx_chunk)
        self.http_client.bind("acquire", self._on_prx_acquire)
        self.http_client.bind("close", self._on_prx_close)
        self.http_client.bind("error", self._on_prx_error)

        self.raw_client = netius.clients.RawClient(
            thread = False,
            *args,
            **kwargs
        )
        self.raw_client.bind("connect", self._on_raw_connect)
        self.raw_client.bind("data", self._on_raw_data)
        self.raw_client.bind("close", self._on_raw_close)

        self.container = netius.Container(*args, **kwargs)
        self.container.add_base(self)
        self.container.add_base(self.http_client)
        self.container.add_base(self.raw_client)

    def start(self):
        self.http_client.load()
        self.raw_client.load()
        self.container.start()

    def stop(self):
        self.container.stop()

    def on_data(self, connection, data):
        netius.StreamServer.on_data(self, connection, data)

        tunnel_c = hasattr(connection, "tunnel_c") and connection.tunnel_c

        if tunnel_c: tunnel_c.send(data)
        else: connection.parse(data)

    def on_connection_d(self, connection):
        http.HTTPServer.on_connection_d(self, connection)

        tunnel_c = hasattr(connection, "tunnel_c") and connection.tunnel_c
        proxy_c = hasattr(connection, "proxy_c") and connection.proxy_c

        if tunnel_c: tunnel_c.close()
        if proxy_c: proxy_c.close()

    def _prx_close(self, connection):
        connection.close()

    def _prx_keep(self, connection):
        pass

    def _on_prx_headers(self, client, parser, headers):
        # retrieves the owner of the parser as the client connection
        # and then retrieves all the other http specific values
        _connection = parser.owner
        code_s = parser.code_s
        status_s = parser.status_s
        version_s = parser.version_s

        # resolves the client connection into the proper proxy connection
        # to be used to send the headers (and status line) to the client
        connection = self.conn_map[_connection]

        # creates a buffer list that will hold the complete set of
        # lines that compose both the status lines and the headers
        # then appends the start line and the various header lines
        # to it so that it contains all of them
        buffer = []
        buffer.append("%s %s %s\r\n" % (version_s, code_s, status_s))
        for key, value in headers.iteritems():
            key = netius.common.header_up(key)
            buffer.append("%s: %s\r\n" % (key, value))
        buffer.append("\r\n")

        # joins the header strings list as the data string that contains
        # the headers and then sends the value through the connection
        data = "".join(buffer)
        connection.send(data)

    def _on_prx_message(self, client, parser, message):
        # retrieves the back-end connection from the provided parser and
        # then evaluates if that connection is of type chunked
        _connection = parser.owner
        is_chunked = parser.chunked

        # sets the current client connection as unused and then retrieves
        # the requester connection associated with the client (back-end)
        # connection in order to be used in the current processing
        _connection.used = False
        connection = self.conn_map[_connection]

        # creates the clojure function that will be used to close both
        # the client and the server connections on a callback action
        def close(connection): connection.close(); _connection.close()

        # verifies that the connection is meant to be kept alive, the
        # connection is meant to be kept alive when both the client and
        # the final (back-end) server respond with the keep alive flag
        keep_alive = parser.keep_alive and connection.parser.keep_alive

        # defines the proper callback function to be called at the end
        # of the flushing of the connection according to the result of
        # the keep alive evaluation (as defined in specification)
        if keep_alive: callback = None
        else: callback = close

        # verifies if the current connection is of type chunked an in case
        # it is must first send the final (close) chunk and then call the
        # proper callback otherwise in case it's a plain connection the
        # callback is immediately called in case it's defined
        if is_chunked: connection.send("0\r\n\r\n", callback = callback)
        elif callback: connection.send("", callback = callback)

    def _on_prx_partial(self, client, parser, data):
        _connection = parser.owner
        is_chunked = parser.chunked

        if is_chunked: return

        connection = self.conn_map[_connection]
        connection.send(data)

    def _on_prx_chunk(self, client, parser, range):
        _connection = parser.owner
        connection = self.conn_map[_connection]

        start, end = range
        data = parser.message[start:end]
        data_s = "".join(data)
        data_l = len(data_s)
        header = "%x\r\n" % data_l
        chunk = header + data_s + "\r\n"
        connection.send(chunk)

    def _on_prx_acquire(self, client, _connection):
        _connection.used = True

    def _on_prx_close(self, client, _connection):
        connection = self.conn_map[_connection]
        if _connection.used: connection.close(flush = True)
        del self.conn_map[_connection]

    def _on_prx_error(self, client, _connection, error):
        error_m = str(error) or "Unknown proxy relay error"
        connection = self.conn_map[_connection]
        connection.send_response(
            data = error_m,
            code = 500,
            code_s = "Internal Error",
            callback = self._prx_close
        )

    def _on_raw_connect(self, client, _connection):
        connection = self.conn_map[_connection]
        connection.send_response(
            code = 200,
            code_s = "Connection established"
        )

    def _on_raw_data(self, client, _connection, data):
        connection = self.conn_map[_connection]
        connection.send(data)

    def _on_raw_close(self, client, _connection):
        connection = self.conn_map[_connection]
        connection.close(flush = True)
        del self.conn_map[_connection]
<commit_msg>new support for connection closing<commit_after>#!/usr/bin/python
# -*- coding: utf-8 -*-

# Hive Netius System
# Copyright (C) 2008-2012 Hive Solutions Lda.
#
# This file is part of Hive Netius System.
#
# Hive Netius System is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# Hive Netius System is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with Hive Netius System. If not, see <http://www.gnu.org/licenses/>.

__author__ = "João Magalhães joamag@hive.pt>"
""" The author(s) of the module """

__version__ = "1.0.0"
""" The version of the module """

__revision__ = "$LastChangedRevision$"
""" The revision number of the module """

__date__ = "$LastChangedDate$"
""" The last change date of the module """

__copyright__ = "Copyright (c) 2008-2012 Hive Solutions Lda."
""" The copyright for the module """

__license__ = "GNU General Public License (GPL), Version 3"
""" The license for the module """

import http

import netius.common
import netius.clients

class ProxyServer(http.HTTPServer):

    def __init__(self, *args, **kwargs):
        http.HTTPServer.__init__(self, *args, **kwargs)
        self.conn_map = {}

        self.http_client = netius.clients.HTTPClient(
            thread = False,
            *args,
            **kwargs
        )
        self.http_client.bind("headers", self._on_prx_headers)
        self.http_client.bind("message", self._on_prx_message)
        self.http_client.bind("partial", self._on_prx_partial)
        self.http_client.bind("chunk", self._on_prx_chunk)
        self.http_client.bind("acquire", self._on_prx_acquire)
        self.http_client.bind("close", self._on_prx_close)
        self.http_client.bind("error", self._on_prx_error)

        self.raw_client = netius.clients.RawClient(
            thread = False,
            *args,
            **kwargs
        )
        self.raw_client.bind("connect", self._on_raw_connect)
        self.raw_client.bind("data", self._on_raw_data)
        self.raw_client.bind("close", self._on_raw_close)

        self.container = netius.Container(*args, **kwargs)
        self.container.add_base(self)
        self.container.add_base(self.http_client)
        self.container.add_base(self.raw_client)

    def start(self):
        self.http_client.load()
        self.raw_client.load()
        self.container.start()

    def stop(self):
        self.container.stop()

    def on_data(self, connection, data):
        netius.StreamServer.on_data(self, connection, data)

        tunnel_c = hasattr(connection, "tunnel_c") and connection.tunnel_c

        if tunnel_c: tunnel_c.send(data)
        else: connection.parse(data)

    def on_connection_d(self, connection):
        http.HTTPServer.on_connection_d(self, connection)

        tunnel_c = hasattr(connection, "tunnel_c") and connection.tunnel_c
        proxy_c = hasattr(connection, "proxy_c") and connection.proxy_c

        if tunnel_c: tunnel_c.close()
        if proxy_c: proxy_c.close()

    def _prx_close(self, connection):
        connection.close()

    def _prx_keep(self, connection):
        pass

    def _on_prx_headers(self, client, parser, headers):
        # retrieves the owner of the parser as the client connection
        # and then retrieves all the other http specific values
        _connection = parser.owner
        code_s = parser.code_s
        status_s = parser.status_s
        version_s = parser.version_s

        # resolves the client connection into the proper proxy connection
        # to be used to send the headers (and status line) to the client
        connection = self.conn_map[_connection]

        # creates a buffer list that will hold the complete set of
        # lines that compose both the status lines and the headers
        # then appends the start line and the various header lines
        # to it so that it contains all of them
        buffer = []
        buffer.append("%s %s %s\r\n" % (version_s, code_s, status_s))
        for key, value in headers.iteritems():
            key = netius.common.header_up(key)
            buffer.append("%s: %s\r\n" % (key, value))
        buffer.append("\r\n")

        # joins the header strings list as the data string that contains
        # the headers and then sends the value through the connection
        data = "".join(buffer)
        connection.send(data)

    def _on_prx_message(self, client, parser, message):
        # retrieves the back-end connection from the provided parser and
        # then evaluates if that connection is of type chunked
        _connection = parser.owner
        is_chunked = parser.chunked

        # sets the current client connection as unused and then retrieves
        # the requester connection associated with the client (back-end)
        # connection in order to be used in the current processing
        _connection.used = False
        connection = self.conn_map[_connection]

        # invalidates both the tunnel connection and the proxy connection
        # as they are no longer associated with the current connection
        # not attached to the client connection
        connection.tunnel_c = None
        connection.proxy_c = None

        # creates the clojure function that will be used to close the
        # current client connection and that may or may not close the
        # corresponding back-end connection (as define in specification)
        def close(connection): connection.close()

        # verifies that the connection is meant to be kept alive, the
        # connection is meant to be kept alive when both the client and
        # the final (back-end) server respond with the keep alive flag
        keep_alive = parser.keep_alive and connection.parser.keep_alive

        # defines the proper callback function to be called at the end
        # of the flushing of the connection according to the result of
        # the keep alive evaluation (as defined in specification)
        if keep_alive: callback = None
        else: callback = close

        # verifies if the current connection is of type chunked an in case
        # it is must first send the final (close) chunk and then call the
        # proper callback otherwise in case it's a plain connection the
        # callback is immediately called in case it's defined
        if is_chunked: connection.send("0\r\n\r\n", callback = callback)
        elif callback: connection.send("", callback = callback)

    def _on_prx_partial(self, client, parser, data):
        _connection = parser.owner
        is_chunked = parser.chunked

        if is_chunked: return

        connection = self.conn_map[_connection]
        connection.send(data)

    def _on_prx_chunk(self, client, parser, range):
        _connection = parser.owner
        connection = self.conn_map[_connection]

        start, end = range
        data = parser.message[start:end]
        data_s = "".join(data)
        data_l = len(data_s)
        header = "%x\r\n" % data_l
        chunk = header + data_s + "\r\n"
        connection.send(chunk)

    def _on_prx_acquire(self, client, _connection):
        _connection.used = True

    def _on_prx_close(self, client, _connection):
        connection = self.conn_map[_connection]
        if _connection.used: connection.close(flush = True)
        del self.conn_map[_connection]

    def _on_prx_error(self, client, _connection, error):
        error_m = str(error) or "Unknown proxy relay error"
        connection = self.conn_map[_connection]
        connection.send_response(
            data = error_m,
            code = 500,
            code_s = "Internal Error",
            callback = self._prx_close
        )

    def _on_raw_connect(self, client, _connection):
        connection = self.conn_map[_connection]
        connection.send_response(
            code = 200,
            code_s = "Connection established"
        )

    def _on_raw_data(self, client, _connection, data):
        connection = self.conn_map[_connection]
        connection.send(data)

    def _on_raw_close(self, client, _connection):
        connection = self.conn_map[_connection]
        connection.close(flush = True)
        del self.conn_map[_connection]
<|endoftext|>


<commit_before>"""
Misc database methods
"""

import os
import logging
from datetime import datetime,timedelta
from functools import partial,reduce
import operator
from collections import OrderedDict, Iterable
import math
import uuid
import shutil
from io import BytesIO

import tornado.gen

# load mysql Error
try:
    import MySQLdb
except ImportError:
    try:
        import pymysql as MySQLdb
    except:
        class MySQLdb:
            class Error(Exception):
                pass

import iceprod.core.functions
from iceprod.core.dataclasses import Number,String
from iceprod.core import serialization
from iceprod.core.jsonUtil import json_encode,json_decode,json_compressor

from iceprod.server.dbmethods import _Methods_Base,filtered_input,datetime2str,str2datetime

logger = logging.getLogger('dbmethods.misc')

class misc(_Methods_Base):
    """
    misc DB methods.

    Takes a handle to a subclass of iceprod.server.modules.db.DBAPI
    as an argument.
    """

    ### site to site ###

    def misc_site_to_site_upload(self,src,dest):
        # TODO: actually write this method
        raise NotImplementedError()

    @tornado.gen.coroutine
    def misc_get_tables_for_task(self,task_ids):
        """
        Get all tables necessary to run task(s).

        Args:
            task_ids (iterable): An iterable of task_ids

        Returns:
            dict: table entries
        """
        if isinstance(task_ids,str):
            task_ids = [task_ids]
        elif not isinstance(task_ids,Iterable):
            raise Exception('task_ids not Iterable')
        task_ids = set(task_ids)

        tables = {}

        sql = 'select depends from task where task_id in ('
        sql += ','.join('?' for _ in task_ids)
        sql += ')'
        bindings = tuple(task_ids)
        ret = yield self.parent.db.query(sql, bindings)
        if not ret:
            raise tornado.gen.Return({})
        for row in ret:
            for d in row[0].split(','):
                d = d.strip()
                if d:
                    task_ids.add(d)

        sql = 'select * from search where task_id in ('
        sql += ','.join('?' for _ in task_ids)
        sql += ')'
        bindings = tuple(task_ids)
        ret = yield self.parent.db.query(sql, bindings)
        if not ret:
            raise tornado.gen.Return({})

        search_table = {}
        keys = []
        for row in ret:
            search_table[row[0]] = self._list_to_dict('search',row)
            if not keys:
                keys = search_table[row[0]].keys()
        tables['search'] = {'keys':keys,'values':ret}

        sql = 'select * from task where task_id in ('
        sql += ','.join('?' for _ in task_ids)
        sql += ')'
        bindings = tuple(task_ids)
        ret = yield self.parent.db.query(sql, bindings)
        task_rel_ids = set()
        if ret:
            keys = []
            for row in ret:
                row2 = self._list_to_dict('task',row)
                task_rel_ids.add(row2['task_rel_id'])
                if not keys:
                    keys = row2.keys()
            tables['task'] = {'keys':keys,'values':ret}

        job_ids = set(search_table[id]['job_id'] for id in task_ids)
        sql = 'select * from job where job_id in ('
        sql += ','.join('?' for _ in job_ids)
        sql += ')'
        bindings = tuple(job_ids)
        ret = yield self.parent.db.query(sql, bindings)
        if ret:
            keys = []
            for row in ret:
                if not keys:
                    keys = self._list_to_dict('job',row).keys()
                    break
            tables['job'] = {'keys':keys,'values':ret}

        dataset_ids = set(search_table[id]['dataset_id'] for id in task_ids)
        sql = 'select * from dataset where dataset_id in ('
        sql += ','.join('?' for _ in dataset_ids)
        sql += ')'
        bindings = tuple(dataset_ids)
        ret = yield self.parent.db.query(sql, bindings)
        categoryvalue_ids = set()
        group_ids = set()
        if ret:
            keys = []
            for row in ret:
                row2 = self._list_to_dict('dataset',row)
                if row2['categoryvalue_ids']:
                    for cv_id in row2['categoryvalue_ids'].split(','):
                        categoryvalue_ids.add(cv_id)
                if 'groupd_id' in row2 and row2['group_id']:
                    group_ids.add(row2['group_id'])
                if not keys:
                    keys = row2.keys()
            tables['dataset'] = {'keys':keys,'values':ret}

        if group_ids:
            sql = 'select * from groups where group_ids in ('
            sql += ','.join('?' for _ in group_ids) + ')'
            bindings = tuple(group_ids)
            ret = yield self.parent.db.query(sql, bindings)
            if ret:
                keys = []
                for row in ret:
                    if not keys:
                        keys = self._list_to_dict('groups',row).keys()
                        break
                tables['groups'] = {'keys':keys,'values':ret}

        sql = 'select * from task_rel where dataset_id in ('
        sql += ','.join('?' for _ in dataset_ids)
        sql += ')'
        bindings = tuple(dataset_ids)
        ret = yield self.parent.db.query(sql, bindings)
        if ret:
            keys = []
            for row in ret:
                if not keys:
                    keys = self._list_to_dict('task_rel',row).keys()
                    break
            tables['task_rel'] = {'keys':keys,'values':ret}

        sql = 'select * from config where dataset_id in ('
        sql += ','.join('?' for _ in dataset_ids)
        sql += ')'
        bindings = tuple(dataset_ids)
        ret = yield self.parent.db.query(sql, bindings)
        if ret:
            keys = []
            for row in ret:
                if not keys:
                    keys = self._list_to_dict('config',row).keys()
                    break
            tables['config'] = {'keys':keys,'values':ret}

        categorydef_ids = set()
        if categoryvalue_ids:
            sql = 'select * from categoryvalue where categoryvalue_id in ('
            sql += ','.join('?' for _ in categoryvalue_ids)
            sql += ')'
            bindings = tuple(categoryvalue_ids)
            ret = yield self.parent.db.query(sql, bindings)
            if ret:
                keys = []
                for row in ret:
                    row2 = self._list_to_dict('categoryvalue',row)
                    categorydef_ids.add(row2['categorydef_id'])
                    if not keys:
                        keys = row2.keys()
                tables['categoryvalue'] = {'keys':keys,'values':ret}

        if categorydef_ids:
            sql = 'select * from categorydef where categorydef_id in ('
            sql += ','.join('?' for _ in categorydef_ids)
            sql += ')'
            bindings = tuple(categorydef_ids)
            ret = yield self.parent.db.query(sql, bindings)
            if ret:
                keys = []
                for row in ret:
                    if not keys:
                        keys = self._list_to_dict('categorydef',row).keys()
                        break
                tables['categorydef'] = {'keys':keys,'values':ret}

        raise tornado.gen.Return(tables)

    @tornado.gen.coroutine
    def misc_update_tables(self, tables):
        """
        Update the DB tables with the incoming information.

        Args:
            tables (dict): {table_name:{keys:[],values:[[]]}}

        Returns:
            bool: success or failure
        """
        if not tables:
            return
        if not isinstance(tables,dict):
            raise Exception('tables not a dict')
        try:
            for name in tables:
                sql = 'replace into %s ('%filtered_input(name)
                sql += ','.join(str(filtered_input(k)) for k in tables[name]['keys'])
                sql += ') values ('
                sql += ','.join('?' for _ in tables[name]['keys'])
                sql += ')'
                for values in tables[name]['values']:
                    bindings = tuple(values)
                    yield self.parent.db.query(sql, bindings)
        except:
            logger.warn('error updating tables', exc_info=True)
            raise

    @tornado.gen.coroutine
    def misc_update_master_db(self, table, index, timestamp, sql, bindings):
        """
        Update the DB with incoming information (query provided).

        Args:
            table (str): The table affected
            index (str): That table's index id
            timestamp (str): An ISO 8601 UTC timestamp
            sql (str): An sql statement
            bindings (tuple): Bindings for the sql statement
        """
        with (yield self.parent.db.acquire_lock('update_master')):
            try:
                sql2 = 'select timestamp from master_update_history '
                sql2 += 'where table_name = ? and update_index = ?'
                bindings2 = (table,index)
                ret = yield self.parent.db.query(sql2, bindings2)
                prev_timestamp = None
                for row in ret:
                    prev_timestamp = row[0]
                if prev_timestamp and prev_timestamp >= timestamp:
                    logger.info('newer data already present for %s %s %s',
                                table, index, timestamp)
                    return
                yield self.parent.db.query(sql, bindings)
            except MySQLdb.Error:
                logger.warn('dropping history for %r', sql, exc_info=True)
            except:
                logger.warn('error updating master', exc_info=True)
                raise
            else:
                sql2 = 'replace into master_update_history (table_name,update_index,timestamp) values (?,?,?)'
                bindings2 = (table,index,timestamp)
                try:
                    yield self.parent.db.query(sql2, bindings2)
                except MySQLdb.Error:
                    logger.warn('mysql error updating update_history',
                                exc_info=True)
                except:
                    logger.warn('error updating update_history', exc_info=True)
                    raise
<commit_msg>fix get_tables_for_task so we actually send all the job info<commit_after>"""
Misc database methods
"""

import os
import logging
from datetime import datetime,timedelta
from functools import partial,reduce
import operator
from collections import OrderedDict, Iterable
import math
import uuid
import shutil
from io import BytesIO

import tornado.gen

# load mysql Error
try:
    import MySQLdb
except ImportError:
    try:
        import pymysql as MySQLdb
    except:
        class MySQLdb:
            class Error(Exception):
                pass

import iceprod.core.functions
from iceprod.core.dataclasses import Number,String
from iceprod.core import serialization
from iceprod.core.jsonUtil import json_encode,json_decode,json_compressor

from iceprod.server.dbmethods import _Methods_Base,filtered_input,datetime2str,str2datetime

logger = logging.getLogger('dbmethods.misc')

class misc(_Methods_Base):
    """
    misc DB methods.

    Takes a handle to a subclass of iceprod.server.modules.db.DBAPI
    as an argument.
    """

    ### site to site ###

    def misc_site_to_site_upload(self,src,dest):
        # TODO: actually write this method
        raise NotImplementedError()

    @tornado.gen.coroutine
    def misc_get_tables_for_task(self,task_ids):
        """
        Get all tables necessary to run task(s).

        Args:
            task_ids (iterable): An iterable of task_ids

        Returns:
            dict: table entries
        """
        if isinstance(task_ids,str):
            task_ids = [task_ids]
        elif not isinstance(task_ids,Iterable):
            raise Exception('task_ids not Iterable')
        task_ids = set(task_ids)

        tables = {}

        sql = 'select depends from task where task_id in (%s)'
        keys = []
        values = []
        for f in self._bulk_select(sql, task_ids):
            ret = yield f
            for row in ret:
                for d in row[0].split(','):
                    d = d.strip()
                    if d:
                        task_ids.add(d)

        job_ids = set()
        dataset_ids = set()
        sql = 'select * from search where task_id in (%s)'
        keys = []
        values = []
        for f in self._bulk_select(sql, task_ids):
            ret = yield f
            for row in ret:
                tmp = self._list_to_dict('search',row)
                if not keys:
                    keys = list(tmp.keys())
                job_ids.add(tmp['job_id'])
                dataset_ids.add(tmp['dataset_id'])
            values.extend(ret)
        if keys and values:
            tables['search'] = {'keys':keys,'values':values}
        else:
            raise tornado.gen.Return({})

        sql = 'select * from task where task_id in (%s)'
        keys = []
        values = []
        for f in self._bulk_select(sql, task_ids):
            ret = yield f
            for row in ret:
                row2 = self._list_to_dict('task',row)
                if not keys:
                    keys = list(row2.keys())
            values.extend(ret)
        if keys and values:
            tables['task'] = {'keys':keys,'values':values}

        sql = 'select * from job where job_id in (%s)'
        keys = []
        values = []
        for f in self._bulk_select(sql, job_ids):
            ret = yield f
            for row in ret:
                if not keys:
                    keys = list(self._list_to_dict('job',row).keys())
                    break
            values.extend(ret)
        if keys and values:
            tables['job'] = {'keys':keys,'values':values}

        categoryvalue_ids = set()
        group_ids = set()
        sql = 'select * from dataset where dataset_id in (%s)'
        keys = []
        values = []
        for f in self._bulk_select(sql, dataset_ids):
            ret = yield f
            for row in ret:
                row2 = self._list_to_dict('dataset',row)
                if row2['categoryvalue_ids']:
                    for cv_id in row2['categoryvalue_ids'].split(','):
                        categoryvalue_ids.add(cv_id)
                if 'groupd_id' in row2 and row2['group_id']:
                    group_ids.add(row2['group_id'])
                if not keys:
                    keys = list(row2.keys())
            values.extend(ret)
        if keys and values:
            tables['dataset'] = {'keys':keys,'values':values}

        if group_ids:
            sql = 'select * from groups where group_ids in (%s)'
            keys = []
            values = []
            for f in self._bulk_select(sql, group_ids):
                ret = yield f
                for row in ret:
                    if not keys:
                        keys = list(self._list_to_dict('groups',row).keys())
                        break
                values.extend(ret)
            if keys and values:
                tables['groups'] = {'keys':keys,'values':values}

        sql = 'select * from task_rel where dataset_id in (%s)'
        keys = []
        values = []
        for f in self._bulk_select(sql, dataset_ids):
            ret = yield f
            for row in ret:
                if not keys:
                    keys = self._list_to_dict('task_rel',row).keys()
                    break
            values.extend(ret)
        if keys and values:
            tables['task_rel'] = {'keys':keys,'values':values}

        sql = 'select * from config where dataset_id in (%s)'
        keys = []
        values = []
        for f in self._bulk_select(sql, dataset_ids):
            ret = yield f
            for row in ret:
                if not keys:
                    keys = list(self._list_to_dict('config',row).keys())
                    break
            values.extend(ret)
        if keys and values:
            tables['config'] = {'keys':keys,'values':values}

        categorydef_ids = set()
        if categoryvalue_ids:
            sql = 'select * from categoryvalue where categoryvalue_id in (%s)'
            keys = []
            values = []
            for f in self._bulk_select(sql, categoryvalue_ids):
                ret = yield f
                for row in ret:
                    row2 = self._list_to_dict('categoryvalue',row)
                    categorydef_ids.add(row2['categorydef_id'])
                    if not keys:
                        keys = list(row2.keys())
                values.extend(ret)
            if keys and values:
                tables['categoryvalue'] = {'keys':keys,'values':values}

        if categorydef_ids:
            sql = 'select * from categorydef where categorydef_id in (%s)'
            keys = []
            values = []
            for f in self._bulk_select(sql, categorydef_ids):
                ret = yield f
                for row in ret:
                    if not keys:
                        keys = list(self._list_to_dict('categorydef',row).keys())
                        break
                values.extend(ret)
            if keys and values:
                tables['categorydef'] = {'keys':keys,'values':values}

        raise tornado.gen.Return(tables)

    @tornado.gen.coroutine
    def misc_update_tables(self, tables):
        """
        Update the DB tables with the incoming information.

        Args:
            tables (dict): {table_name:{keys:[],values:[[]]}}

        Returns:
            bool: success or failure
        """
        if not tables:
            return
        if not isinstance(tables,dict):
            raise Exception('tables not a dict')
        try:
            for name in tables:
                sql = 'replace into %s ('%filtered_input(name)
                sql += ','.join(str(filtered_input(k)) for k in tables[name]['keys'])
                sql += ') values ('
                sql += ','.join('?' for _ in tables[name]['keys'])
                sql += ')'
                for values in tables[name]['values']:
                    bindings = tuple(values)
                    yield self.parent.db.query(sql, bindings)
        except:
            logger.warn('error updating tables', exc_info=True)
            raise

    @tornado.gen.coroutine
    def misc_update_master_db(self, table, index, timestamp, sql, bindings):
        """
        Update the DB with incoming information (query provided).

        Args:
            table (str): The table affected
            index (str): That table's index id
            timestamp (str): An ISO 8601 UTC timestamp
            sql (str): An sql statement
            bindings (tuple): Bindings for the sql statement
        """
        with (yield self.parent.db.acquire_lock('update_master')):
            try:
                sql2 = 'select timestamp from master_update_history '
                sql2 += 'where table_name = ? and update_index = ?'
                bindings2 = (table,index)
                ret = yield self.parent.db.query(sql2, bindings2)
                prev_timestamp = None
                for row in ret:
                    prev_timestamp = row[0]
                if prev_timestamp and prev_timestamp >= timestamp:
                    logger.info('newer data already present for %s %s %s',
                                table, index, timestamp)
                    return
                yield self.parent.db.query(sql, bindings)
            except MySQLdb.Error:
                logger.warn('dropping history for %r', sql, exc_info=True)
            except:
                logger.warn('error updating master', exc_info=True)
                raise
            else:
                sql2 = 'replace into master_update_history (table_name,update_index,timestamp) values (?,?,?)'
                bindings2 = (table,index,timestamp)
                try:
                    yield self.parent.db.query(sql2, bindings2)
                except MySQLdb.Error:
                    logger.warn('mysql error updating update_history',
                                exc_info=True)
                except:
                    logger.warn('error updating update_history', exc_info=True)
                    raise
<|endoftext|>


<commit_before>#!/usr/bin/env python
# -*- mode: python; coding: utf-8 -*-
# Copyright 2022 the HERA Collaboration
# Licensed under the 2-clause BSD license.
"""Accumulate SNAP spectra from redis and write to a UVH5 file."""
import argparse
import json
import logging
import time
import traceback
from pathlib import Path

import numpy as np
from astropy.time import Time, TimeDelta
from hera_corr_cm import HeraCorrCM
from hera_corr_cm.redis_cm import read_cminfo_from_redis, read_maps_from_redis
from pyuvdata import UVData
from pyuvdata import utils as uvutils

formatter = "%(asctime)s.%(msecs)03d %(levelname)s - %(module)s: %(message)s"
logging.basicConfig(
    level=logging.INFO,
    format=formatter,
    datefmt="%Y-%m-%d %H:%M:%S",
)
logger = logging.getLogger(__file__)
logger.ident = "hera_mini_snap_catcher"

syslog_handler = logging.handlers.SysLogHandler(address="/dev/log")
syslog_handler.setFormatter(
    logging.Formatter(
        "%(module)s[%(process)s]: %(message)s",
        "%FT%H:%m:%s%z",
    )
)
logger.addHandler(syslog_handler)

def get_blt_index(antnum, time, ant_array, time_array):
    indices = np.nonzero(np.asarray(ant_array) == antnum)[0]

    if indices.size != 0:
        sub_inds = np.nonzero(np.asarray(time_array)[indices] == time)[0]
        if sub_inds.size != 0:
            return indices[sub_inds]
        else:
            return None
    else:
        return None

parser = argparse.ArgumentParser(
    description="Create UVH5 file from snap autocorrelations in Redis.",
    formatter_class=argparse.ArgumentDefaultsHelpFormatter,
)
parser.add_argument(
    "-r",
    dest="redishost",
    type=str,
    default="redishost",
    help="Host servicing redis requests",
)
parser.add_argument(
    "--max_downtime",
    dest="max_downtime",
    type=float,
    default=60.0,
    help=(
        "Seconds between to wait between redis updates "
        "before restarting with new file."
    ),
)

parser.add_argument(
    "--max_file_len",
    dest="max_file_len",
    type=float,
    default=60.0,
    help=(
        "Maximum length of a UVH5 file in seconds. "
        "This is the total time between the first and last entry in the time_array."
    ),
)

parser.add_argument(
    "--outdir",
    "-o",
    dest="outdir",
    type=str,
    default=".",
    help="The directory to write files to.",
)

args = parser.parse_args()

logger.info("Starting HERA snap mini data catcher.")
while True:
    try:
        corr_cm = HeraCorrCM(redishost=args.redishost)
        logger.info("HERA CorrCM Redis connection established.")

        mapping = read_maps_from_redis()
        snap_to_ant_mapping = mapping["snap_to_ant"]
        if mapping is None:
            raise ValueError(
                "None value encountered for antenna/snap mappings in redis."
            )

        last_time_mapping = {}
        last_loop_completion = Time.now()

        logger.info("Beginning data catching.")
        while True:
            try:

                # use future_array_shapes
                # Nblts, Nfreqs, Npols
                # we don't necessarily know how many freqs we'll have.
                # it is most likely 1024 but this will guard against X-Engine changes.
                data_array = None
                time_array = []
                ant_array = []

                downtime = 0
                file_len = 0

                inner_loop_restart = Time.now()

                while (downtime < args.max_downtime) and (file_len < args.max_file_len):
                    time.sleep(0.01)
                    # get integrations from HeraCorrCM.get_snaprf_status()
                    # compare to the last time for each snap

                    for (
                        snap_hostname,
                        status,
                    ) in corr_cm.get_snaprf_status().items():

                        if status["timestamp"] is None or status["timestamp"] == "None":
                            continue

                        snap, stream = snap_hostname.split(":")
                        stream = int(stream)

                        hera_ant_num = snap_to_ant_mapping[snap][stream]
                        if hera_ant_num is None:
                            continue

                        # arbitrarily assign east as polarization 0
                        pol = hera_ant_num[-1].lower()
                        pol_ind = 0 if pol == "e" else 1

                        # trim off the HH and the Polarization
                        hera_ant_num = int(hera_ant_num[2:-1])

                        timestamp = Time(status["timestamp"]).jd

                        if snap_hostname not in last_time_mapping:
                            last_time_mapping[snap_hostname] = timestamp

                            # Data has changed, restart the inner loop timer
                            inner_loop_restart = Time.now()
                        else:
                            if timestamp <= last_time_mapping[snap_hostname]:
                                # Time has not progressed yet.
                                # Just sit tight
                                continue
                            last_time_mapping[snap_hostname] = timestamp

                            # Data has changed, restart the inner loop timer
                            inner_loop_restart = Time.now()

                        auto_corr = status["autocorrelation"]

                        if data_array is None:
                            data_array = np.zeros(
                                (0, auto_corr.shape[0], 2), dtype=np.float32
                            )
                        data_index = get_blt_index(
                            hera_ant_num, timestamp, ant_array, time_array
                        )
                        if data_index is not None:
                            # this baseline is already in the time array
                            # insert into the correct polarization
                            # at the given index
                            data_array[data_index, :, pol_ind] = auto_corr

                        else:
                            # otherwise append the data array and insert the meta-data
                            data_insert = np.zeros(
                                (1, auto_corr.shape[0], 2), dtype=np.float32
                            )
                            data_insert[..., pol_ind] = auto_corr

                            data_array = np.append(data_array, data_insert, axis=0)
                            time_array.append(timestamp)
                            ant_array.append(hera_ant_num)

                    last_time = Time(max(last_time_mapping.values()), format="jd")
                    # Take the least time between the last snap update
                    # and the last time a file was written
                    # and when we started this inner loop

                    downtime = (
                        (
                            Time.now()
                            - np.asarray(
                                [inner_loop_restart, last_time, last_loop_completion]
                            )
                        )
                        .min()
                        .to_value("s")
                    )
                    if len(time_array) > 0:
                        file_len = TimeDelta(
                            time_array[-1] - time_array[0], format="jd"
                        ).to_value("s")

                time_array = np.asarray(time_array)
                if time_array.size == 0:
                    # No Data was taken
                    logger.info(
                        f"Downtime timeout. No new data recieved since {last_time.jd} "
                        f"({last_time.iso} UTC)."
                    )
                    continue

                ant_array = np.asarray(ant_array)

                uvd = UVData()

                uvd.data_array = data_array.astype(np.complex64)
                uvd.vis_units = "UNCALIB"
                uvd.time_array = time_array

                uvd.Nblts = time_array.shape[0]
                uvd.Ntimes = np.unique(uvd.time_array).size

                uvd.Nfreqs = data_array.shape[1]
                uvd.freq_array = np.linspace(0, 256, uvd.Nfreqs) * 1e6

                uvd.channel_width = np.full_like(
                    uvd.freq_array, np.diff(uvd.freq_array)[0]
                )
                uvd.Nants_data = np.unique(ant_array).size
                uvd.Nbls = uvd.Nants_data
                uvd.Nspws = 1
                uvd.Npols = 2

                uvd.ant_1_array = ant_array
                uvd.ant_2_array = ant_array
                uvd.baseline_array = uvd.antnums_to_baseline(
                    uvd.ant_1_array, uvd.ant_2_array
                )

                uvd.nsample_array = np.ones_like(uvd.data_array, np.float32)
                uvd.integration_time = np.ones((uvd.Nblts), np.float32)
                uvd.flag_array = np.zeros_like(uvd.data_array, dtype=bool)
                uvd.x_orientation = "north"
                uvd.Nants_data = len(set(ant_array))
                uvd.polarization_array = uvutils.polstr2num(
                    ["ee", "nn"], x_orientation=uvd.x_orientation
                )
                uvd.spw_array = np.array([0])

                cminfo = read_cminfo_from_redis(return_as="dict")

                uvd.antenna_positions = np.asarray(cminfo["antenna_positions"])
                uvd.antenna_names = np.asarray(cminfo["antenna_names"])
                uvd.antenna_numbers = np.asarray(cminfo["antenna_numbers"])
                uvd.Nants_telescope = len(uvd.antenna_numbers)
                lat, lon, alt = (
                    cminfo["cofa_lat"],
                    cminfo["cofa_lon"],
                    cminfo["cofa_alt"],
                )
                uvd.telescope_location_lat_lon_alt_degrees = (lat, lon, alt)
                uvd.set_uvws_from_antenna_positions()
                uvd.set_lsts_from_time_array()

                uvd.object_name = ""
                uvd.multi_phase_center = False
                uvd.telescope_name = "HERA"
                uvd.instrument = "HERA"
                uvd.history = f"Created by {__file__}."
                uvd.phase_type = "drift"

                # lets us not have to spectral windows
                uvd._set_future_array_shapes()

                uvd.reorder_blts("time", "baseline")
                uvd.extra_keywords["snap_to_ant_mapping"] = json.dumps(
                    snap_to_ant_mapping
                )
                uvd.extra_keywords["tag"] = "engineering"

                date = Time.now().jd

                folder_name = f"{int(np.floor(date)):d}"

                # make the folder for the JD if it does not exist
                folder = Path(args.outdir) / folder_name
                folder.mkdir(exist_ok=True)

                fname = folder / f"zen.{date:.6f}.snap_autos.uvh5"
                logger.info(f"Writing output file {fname}.")
                uvd.write_uvh5(fname)

                last_loop_completion = Time.now()

            except Exception:
                logger.error(
                    "{t} -- error while accumulating datafile".format(t=time.asctime()),
                )
                logger.error(traceback.print_exc())
                continue

    except Exception:
        logger.error(traceback.print_exc())
        continue
<commit_msg>add 30s sleep to mini catcher, autos only update every once and a while<commit_after>#!/usr/bin/env python
# -*- mode: python; coding: utf-8 -*-
# Copyright 2022 the HERA Collaboration
# Licensed under the 2-clause BSD license.
"""Accumulate SNAP spectra from redis and write to a UVH5 file."""
import argparse
import json
import logging
import time
import traceback
from pathlib import Path

import numpy as np
from astropy.time import Time, TimeDelta
from hera_corr_cm import HeraCorrCM
from hera_corr_cm.redis_cm import read_cminfo_from_redis, read_maps_from_redis
from pyuvdata import UVData
from pyuvdata import utils as uvutils

formatter = "%(asctime)s.%(msecs)03d %(levelname)s - %(module)s: %(message)s"
logging.basicConfig(
    level=logging.INFO,
    format=formatter,
    datefmt="%Y-%m-%d %H:%M:%S",
)
logger = logging.getLogger(__file__)
logger.ident = "hera_mini_snap_catcher"

syslog_handler = logging.handlers.SysLogHandler(address="/dev/log")
syslog_handler.setFormatter(
    logging.Formatter(
        "%(module)s[%(process)s]: %(message)s",
        "%FT%H:%m:%s%z",
    )
)
logger.addHandler(syslog_handler)

def get_blt_index(antnum, time, ant_array, time_array):
    indices = np.nonzero(np.asarray(ant_array) == antnum)[0]

    if indices.size != 0:
        sub_inds = np.nonzero(np.asarray(time_array)[indices] == time)[0]
        if sub_inds.size != 0:
            return indices[sub_inds]
        else:
            return None
    else:
        return None

parser = argparse.ArgumentParser(
    description="Create UVH5 file from snap autocorrelations in Redis.",
    formatter_class=argparse.ArgumentDefaultsHelpFormatter,
)
parser.add_argument(
    "-r",
    dest="redishost",
    type=str,
    default="redishost",
    help="Host servicing redis requests",
)
parser.add_argument(
    "--max_downtime",
    dest="max_downtime",
    type=float,
    default=60.0,
    help=(
        "Seconds between to wait between redis updates "
        "before restarting with new file."
    ),
)

parser.add_argument(
    "--max_file_len",
    dest="max_file_len",
    type=float,
    default=60.0,
    help=(
        "Maximum length of a UVH5 file in seconds. "
        "This is the total time between the first and last entry in the time_array."
    ),
)

parser.add_argument(
    "--outdir",
    "-o",
    dest="outdir",
    type=str,
    default=".",
    help="The directory to write files to.",
)

args = parser.parse_args()

logger.info("Starting HERA snap mini data catcher.")
while True:
    try:
        corr_cm = HeraCorrCM(redishost=args.redishost)
        logger.info("HERA CorrCM Redis connection established.")

        mapping = read_maps_from_redis()
        snap_to_ant_mapping = mapping["snap_to_ant"]
        if mapping is None:
            raise ValueError(
                "None value encountered for antenna/snap mappings in redis."
            )

        last_time_mapping = {}
        last_loop_completion = Time.now()

        logger.info("Beginning data catching.")
        while True:
            try:

                # use future_array_shapes
                # Nblts, Nfreqs, Npols
                # we don't necessarily know how many freqs we'll have.
                # it is most likely 1024 but this will guard against X-Engine changes.
                data_array = None
                time_array = []
                ant_array = []

                downtime = 0
                file_len = 0

                inner_loop_restart = Time.now()

                while (downtime < args.max_downtime) and (file_len < args.max_file_len):
                    time.sleep(30)
                    # get integrations from HeraCorrCM.get_snaprf_status()
                    # compare to the last time for each snap

                    for (
                        snap_hostname,
                        status,
                    ) in corr_cm.get_snaprf_status().items():

                        if status["timestamp"] is None or status["timestamp"] == "None":
                            continue

                        snap, stream = snap_hostname.split(":")
                        stream = int(stream)

                        hera_ant_num = snap_to_ant_mapping[snap][stream]
                        if hera_ant_num is None:
                            continue

                        # arbitrarily assign east as polarization 0
                        pol = hera_ant_num[-1].lower()
                        pol_ind = 0 if pol == "e" else 1

                        # trim off the HH and the Polarization
                        hera_ant_num = int(hera_ant_num[2:-1])

                        timestamp = Time(status["timestamp"]).jd

                        if snap_hostname not in last_time_mapping:
                            last_time_mapping[snap_hostname] = timestamp

                            # Data has changed, restart the inner loop timer
                            inner_loop_restart = Time.now()
                        else:
                            if timestamp <= last_time_mapping[snap_hostname]:
                                # Time has not progressed yet.
                                # Just sit tight
                                continue
                            last_time_mapping[snap_hostname] = timestamp

                            # Data has changed, restart the inner loop timer
                            inner_loop_restart = Time.now()

                        auto_corr = status["autocorrelation"]

                        if data_array is None:
                            data_array = np.zeros(
                                (0, auto_corr.shape[0], 2), dtype=np.float32
                            )
                        data_index = get_blt_index(
                            hera_ant_num, timestamp, ant_array, time_array
                        )
                        if data_index is not None:
                            # this baseline is already in the time array
                            # insert into the correct polarization
                            # at the given index
                            data_array[data_index, :, pol_ind] = auto_corr

                        else:
                            # otherwise append the data array and insert the meta-data
                            data_insert = np.zeros(
                                (1, auto_corr.shape[0], 2), dtype=np.float32
                            )
                            data_insert[..., pol_ind] = auto_corr

                            data_array = np.append(data_array, data_insert, axis=0)
                            time_array.append(timestamp)
                            ant_array.append(hera_ant_num)

                    last_time = Time(max(last_time_mapping.values()), format="jd")
                    # Take the least time between the last snap update
                    # and the last time a file was written
                    # and when we started this inner loop

                    downtime = (
                        (
                            Time.now()
                            - np.asarray(
                                [inner_loop_restart, last_time, last_loop_completion]
                            )
                        )
                        .min()
                        .to_value("s")
                    )
                    if len(time_array) > 0:
                        file_len = TimeDelta(
                            time_array[-1] - time_array[0], format="jd"
                        ).to_value("s")

                time_array = np.asarray(time_array)
                if time_array.size == 0:
                    # No Data was taken
                    logger.info(
                        f"Downtime timeout. No new data recieved since {last_time.jd} "
                        f"({last_time.iso} UTC)."
                    )
                    continue

                ant_array = np.asarray(ant_array)

                uvd = UVData()

                uvd.data_array = data_array.astype(np.complex64)
                uvd.vis_units = "UNCALIB"
                uvd.time_array = time_array

                uvd.Nblts = time_array.shape[0]
                uvd.Ntimes = np.unique(uvd.time_array).size

                uvd.Nfreqs = data_array.shape[1]
                uvd.freq_array = np.linspace(0, 256, uvd.Nfreqs) * 1e6

                uvd.channel_width = np.full_like(
                    uvd.freq_array, np.diff(uvd.freq_array)[0]
                )
                uvd.Nants_data = np.unique(ant_array).size
                uvd.Nbls = uvd.Nants_data
                uvd.Nspws = 1
                uvd.Npols = 2

                uvd.ant_1_array = ant_array
                uvd.ant_2_array = ant_array
                uvd.baseline_array = uvd.antnums_to_baseline(
                    uvd.ant_1_array, uvd.ant_2_array
                )

                uvd.nsample_array = np.ones_like(uvd.data_array, np.float32)
                uvd.integration_time = np.ones((uvd.Nblts), np.float32)
                uvd.flag_array = np.zeros_like(uvd.data_array, dtype=bool)
                uvd.x_orientation = "north"
                uvd.Nants_data = len(set(ant_array))
                uvd.polarization_array = uvutils.polstr2num(
                    ["ee", "nn"], x_orientation=uvd.x_orientation
                )
                uvd.spw_array = np.array([0])

                cminfo = read_cminfo_from_redis(return_as="dict")

                uvd.antenna_positions = np.asarray(cminfo["antenna_positions"])
                uvd.antenna_names = np.asarray(cminfo["antenna_names"])
                uvd.antenna_numbers = np.asarray(cminfo["antenna_numbers"])
                uvd.Nants_telescope = len(uvd.antenna_numbers)
                lat, lon, alt = (
                    cminfo["cofa_lat"],
                    cminfo["cofa_lon"],
                    cminfo["cofa_alt"],
                )
                uvd.telescope_location_lat_lon_alt_degrees = (lat, lon, alt)
                uvd.set_uvws_from_antenna_positions()
                uvd.set_lsts_from_time_array()

                uvd.object_name = ""
                uvd.multi_phase_center = False
                uvd.telescope_name = "HERA"
                uvd.instrument = "HERA"
                uvd.history = f"Created by {__file__}."
                uvd.phase_type = "drift"

                # lets us not have to spectral windows
                uvd._set_future_array_shapes()

                uvd.reorder_blts("time", "baseline")
                uvd.extra_keywords["snap_to_ant_mapping"] = json.dumps(
                    snap_to_ant_mapping
                )
                uvd.extra_keywords["tag"] = "engineering"

                date = Time.now().jd

                folder_name = f"{int(np.floor(date)):d}"

                # make the folder for the JD if it does not exist
                folder = Path(args.outdir) / folder_name
                folder.mkdir(exist_ok=True)

                fname = folder / f"zen.{date:.6f}.snap_autos.uvh5"
                logger.info(f"Writing output file {fname}.")
                uvd.write_uvh5(fname)

                last_loop_completion = Time.now()

            except Exception:
                logger.error(
                    "{t} -- error while accumulating datafile".format(t=time.asctime()),
                )
                logger.error(traceback.print_exc())
                continue

    except Exception:
        logger.error(traceback.print_exc())
        continue
<|endoftext|>