<commit_before>#  -*- coding: utf-8 -*-
#  vim: tabstop=4 shiftwidth=4 softtabstop=4

#  Copyright (c) 2013, GEM Foundation

#  OpenQuake is free software: you can redistribute it and/or modify it
#  under the terms of the GNU Affero General Public License as published
#  by the Free Software Foundation, either version 3 of the License, or
#  (at your option) any later version.

#  OpenQuake 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 Affero General Public License
#  along with OpenQuake.  If not, see <http://www.gnu.org/licenses/>.

"""
A script to dump hazard outputs. If you launch it with a given
hazard_calculation_id, it will dump all the hazard outputs relevant for
risk calculations in a tarfile named hc<hazard-calculation-id>.tar.
The tar file can then be moved around and restored in a different
database with the companion script restore_hazards.py.
Internally the dump and restore procedures are based on
COPY TO and COPY FROM commands, so they are quite performant
even for large datasets. They cannot trivially be extended to perform
binary dump/restore since the geography type has no binary form in
PostGIS 1.5.

To restore a hazard computation and all of its outputs into a new database
run ``python restore_hazards.py <tarfile> <host> <dbname> <user> <password>``

The <user> must have sufficient permissions to write on <dbname>.  If
your database already contains a hazard calculation with the same id,
the restore will not override it. If you think that the hazard
calculation on your database is not important and can removed together
with all of its outputs, then remove it by using ``bin/openquake
--delete-hazard-calculation`` (which must be run by a user with
sufficient permissions). Then run again ``restore_hazards.py``.
"""

import os
import shutil
import tarfile
import gzip
import argparse
import psycopg2
import tempfile
import logging

log = logging.getLogger()


# return a string which is a valid SQL tuple
def _tuplestr(tup):
    return '(%s)' % ', '.join(str(x) for x in tup)


class Copier(object):
    """
    Small wrapper around a psycopg2 cursor, which a .copy method
    writing directly to .gz files. It remembers the copied filenames,
    which are stored in the attribute .filenames.
    """
    def __init__(self, psycopg2_cursor):
        self._cursor = psycopg2_cursor
        self.filenames = []

    def tuplestr(self, query, *args):
        """Retrieve tuples of ids a strings"""
        self._cursor.execute(query, args)
        return _tuplestr(row[0] for row in self._cursor)

    def fetchall(self, query, *args):
        """Dispatch to .fetchall"""
        self._cursor.execute(query, args)
        return self._cursor.fetchall()

    def copy(self, query, dest, name, mode):
        """
        Performs a COPY TO/FROM operation. <Works directly with gzipped files.

        :param str query: the COPY query
        :param str dest: the destination directory
        :param str name: the destination file name (no .gz)
        :param chr mode: 'w' (for COPY TO) or 'r' (for COPY FROM)
        """
        fname = os.path.join(dest, name + '.gz')
        log.info('%s\n(-> %s)', query, fname)
        with gzip.open(fname, mode) as f:
            self._cursor.copy_expert(query, f)
            self.filenames.append(fname)


class HazardDumper(object):
    """
    Class to dump a hazard_calculation and related tables.
    The typical usage is

     hd = HazardDumper(conn, '/tmp/somedir')
     hd.dump(42)  # dump the hazard computation #42
     print hd.mktar()  # generate a tarfile
    """

    def __init__(self, conn, outdir=None, format='text'):
        self.conn = conn
        self.curs = Copier(conn.cursor())
        self.format = format
        # there is no binary format for geography in postgis 1.5,
        # this is why we are requiring text format
        assert format == 'text', format
        if outdir:
            os.mkdir(outdir)
        else:
            outdir = tempfile.mkdtemp(prefix='hazard_calculation-')
        self.outdir = outdir

    def hazard_calculation(self, ids):
        """Dump hazard_calculation, lt_realization, hazard_site"""
        self.curs.copy(
            """copy (select * from uiapi.hazard_calculation where id in %s)
                  to stdout with (format '%s')""" % (ids, self.format),
            self.outdir, 'uiapi.hazard_calculation.csv', 'w')
        self.curs.copy(
            """copy (select * from hzrdr.lt_realization
                  where hazard_calculation_id in %s)
                  to stdout with (format '%s')""" % (ids, self.format),
            self.outdir, 'hzrdr.lt_realization.csv', 'w')
        self.curs.copy(
            """copy (select * from hzrdi.hazard_site
                  where hazard_calculation_id in %s)
                  to stdout with (format '%s')""" % (ids, self.format),
            self.outdir, 'hzrdi.hazard_site.csv', 'w')

        # model_content, inputs

        # at this moment, only logic tree content is needed
        self.curs.copy(
            """copy (select mc.* from uiapi.model_content mc
                     join uiapi.input as input
                     on input.model_content_id = mc.id
                     join uiapi.input2hcalc h2c
                     on h2c.input_id = input.id
                     where h2c.hazard_calculation_id = %s
                     and input.input_type in ('source_model_logic_tree',
                                              'gsim_logic_tree'))
               to stdout with (format '%s')""" % (ids, self.format),
            self.outdir, 'uiapi.model_content.csv', 'w')

        input_ids = self.curs.tuplestr(
            """select i.* from uiapi.input i
                     join uiapi.input2hcalc h2c
                     on h2c.input_id = i.id
                     where hazard_calculation_id in %s
                     and input_type in ('source_model_logic_tree',
                                        'gsim_logic_tree')""" % ids)

        self.curs.copy(
            """copy (select * from uiapi.input where id in %s)
               to stdout with (format '%s')""" % (input_ids, self.format),
            self.outdir, 'uiapi.input.csv', 'w')

        self.curs.copy(
            """copy (select * from uiapi.input2hcalc where input_id in %s)
               to stdout with (format '%s')""" % (input_ids, self.format),
            self.outdir, 'uiapi.input2hcalc.csv', 'w')<commit_msg>fix dump_hazards for scenario haz calcs<commit_after>#  -*- coding: utf-8 -*-
#  vim: tabstop=4 shiftwidth=4 softtabstop=4

#  Copyright (c) 2013, GEM Foundation

#  OpenQuake is free software: you can redistribute it and/or modify it
#  under the terms of the GNU Affero General Public License as published
#  by the Free Software Foundation, either version 3 of the License, or
#  (at your option) any later version.

#  OpenQuake 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 Affero General Public License
#  along with OpenQuake.  If not, see <http://www.gnu.org/licenses/>.

"""
A script to dump hazard outputs. If you launch it with a given
hazard_calculation_id, it will dump all the hazard outputs relevant for
risk calculations in a tarfile named hc<hazard-calculation-id>.tar.
The tar file can then be moved around and restored in a different
database with the companion script restore_hazards.py.
Internally the dump and restore procedures are based on
COPY TO and COPY FROM commands, so they are quite performant
even for large datasets. They cannot trivially be extended to perform
binary dump/restore since the geography type has no binary form in
PostGIS 1.5.

To restore a hazard computation and all of its outputs into a new database
run ``python restore_hazards.py <tarfile> <host> <dbname> <user> <password>``

The <user> must have sufficient permissions to write on <dbname>.  If
your database already contains a hazard calculation with the same id,
the restore will not override it. If you think that the hazard
calculation on your database is not important and can removed together
with all of its outputs, then remove it by using ``bin/openquake
--delete-hazard-calculation`` (which must be run by a user with
sufficient permissions). Then run again ``restore_hazards.py``.
"""

import os
import shutil
import tarfile
import gzip
import argparse
import psycopg2
import tempfile
import logging

log = logging.getLogger()


# return a string which is a valid SQL tuple
def _tuplestr(tup):
    return '(%s)' % ', '.join(str(x) for x in tup)


class Copier(object):
    """
    Small wrapper around a psycopg2 cursor, which a .copy method
    writing directly to .gz files. It remembers the copied filenames,
    which are stored in the attribute .filenames.
    """
    def __init__(self, psycopg2_cursor):
        self._cursor = psycopg2_cursor
        self.filenames = []

    def tuplestr(self, query, *args):
        """Retrieve tuples of ids a strings"""
        self._cursor.execute(query, args)
        return _tuplestr(row[0] for row in self._cursor)

    def fetchall(self, query, *args):
        """Dispatch to .fetchall"""
        self._cursor.execute(query, args)
        return self._cursor.fetchall()

    def copy(self, query, dest, name, mode):
        """
        Performs a COPY TO/FROM operation. <Works directly with gzipped files.

        :param str query: the COPY query
        :param str dest: the destination directory
        :param str name: the destination file name (no .gz)
        :param chr mode: 'w' (for COPY TO) or 'r' (for COPY FROM)
        """
        fname = os.path.join(dest, name + '.gz')
        log.info('%s\n(-> %s)', query, fname)
        with gzip.open(fname, mode) as f:
            self._cursor.copy_expert(query, f)
            self.filenames.append(fname)


class HazardDumper(object):
    """
    Class to dump a hazard_calculation and related tables.
    The typical usage is

     hd = HazardDumper(conn, '/tmp/somedir')
     hd.dump(42)  # dump the hazard computation #42
     print hd.mktar()  # generate a tarfile
    """

    def __init__(self, conn, outdir=None, format='text'):
        self.conn = conn
        self.curs = Copier(conn.cursor())
        self.format = format
        # there is no binary format for geography in postgis 1.5,
        # this is why we are requiring text format
        assert format == 'text', format
        if outdir:
            os.mkdir(outdir)
        else:
            outdir = tempfile.mkdtemp(prefix='hazard_calculation-')
        self.outdir = outdir

    def hazard_calculation(self, ids):
        """Dump hazard_calculation, lt_realization, hazard_site"""
        self.curs.copy(
            """copy (select * from uiapi.hazard_calculation where id in %s)
                  to stdout with (format '%s')""" % (ids, self.format),
            self.outdir, 'uiapi.hazard_calculation.csv', 'w')
        self.curs.copy(
            """copy (select * from hzrdr.lt_realization
                  where hazard_calculation_id in %s)
                  to stdout with (format '%s')""" % (ids, self.format),
            self.outdir, 'hzrdr.lt_realization.csv', 'w')
        self.curs.copy(
            """copy (select * from hzrdi.hazard_site
                  where hazard_calculation_id in %s)
                  to stdout with (format '%s')""" % (ids, self.format),
            self.outdir, 'hzrdi.hazard_site.csv', 'w')

        # model_content, inputs

        # at this moment, only logic tree content is needed
        self.curs.copy(
            """copy (select mc.* from uiapi.model_content mc
                     join uiapi.input as input
                     on input.model_content_id = mc.id
                     join uiapi.input2hcalc h2c
                     on h2c.input_id = input.id
                     where h2c.hazard_calculation_id = %s
                     and input.input_type in ('source_model_logic_tree',
                                              'gsim_logic_tree'))
               to stdout with (format '%s')""" % (ids, self.format),
            self.outdir, 'uiapi.model_content.csv', 'w')

        input_ids = self.curs.tuplestr(
            """select i.* from uiapi.input i
                     join uiapi.input2hcalc h2c
                     on h2c.input_id = i.id
                     where hazard_calculation_id in %s
                     and input_type in ('source_model_logic_tree',
                                        'gsim_logic_tree')""" % ids)
        if input_ids != "()":  # scenario calcs do not have logic tree inputs
            self.curs.copy(
                """copy (select * from uiapi.input where id in %s)
                   to stdout with (format '%s')""" % (input_ids, self.format),
                self.outdir, 'uiapi.input.csv', 'w')

            self.curs.copy(
                """copy (select * from uiapi.input2hcalc where input_id in %s)
                   to stdout with (format '%s')""" % (input_ids, self.format),
                self.outdir, 'uiapi.input2hcalc.csv', 'w')
{'fraction_non_alphanumeric': 0.09382679276193313, 'fraction_numerical': 0.00461687392955544, 'mean_word_length': 2.3718302786844085, 'pattern_counts': {'":': 1, '<': 23, '<?xml version=': 0, '>': 23, 'https://': 0, 'lorem ipsum': 0, 'www.': 2, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 6, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '960478', 'n_tokens_mistral': 3995, 'n_tokens_neox': 3815, 'n_words': 1424}