<commit_before>import base64
import connexion
import six
import json
import os
import random
import string
import sys
import time
import urllib

from collections import OrderedDict
from elasticsearch import Elasticsearch
from elasticsearch_dsl import Search
from flask import request
from flask import current_app
from werkzeug.exceptions import BadRequest
from google.cloud import storage
from oauth2client.service_account import ServiceAccountCredentials

from data_explorer.models.export_url_response import ExportUrlResponse  # noqa: E501
from data_explorer.util import elasticsearch_util
from data_explorer.util.dataset_faceted_search import DatasetFacetedSearch

# Export to Saturn flow
# - User clicks export button on bottom right of Data Explorer
# - A dialog pops up, asking user to name cohort. User types "African females"
#   and clicks Export button
# - UI server calls API server /export_url (this file):
#   - Constructs a JSON array of JSON entity objects
#   - Writes JSON array to a GCS file
#   - Concat that with an existing samples JSON entity file, which contains
#     all the samples in the dataset and is created on indexing.
#   - Creates a signed url for GCS file. Returns signed url to UI server
# - UI server redirects to Saturn add-import?url=SIGNED_URL
# - On add-import page, user selects Workspace. User clicks Import button.
# - User is redirected to selected workspace Data tab, showing newly imported
#   entities.


def _check_preconditions():
    config_path = os.path.join(current_app.config['DATASET_CONFIG_DIR'],
                               'deploy.json')
    if not os.path.isfile(config_path):
        error_msg = (
            'deploy.json not found. Export to Saturn feature will not work. '
            'See https://github.com/DataBiosphere/data-explorer#one-time-setup-for-export-to-saturn-feature'
        )
        current_app.logger.error(error_msg)
        raise BadRequest(error_msg)

    if not current_app.config['DEPLOY_PROJECT_ID']:
        error_msg = (
            'Project not set in deploy.json. Export to Saturn feature will not work. '
            'See https://github.com/DataBiosphere/data-explorer#one-time-setup-for-export-to-saturn-feature-for-export-to-saturn-feature'
        )
        current_app.logger.error(error_msg)
        raise BadRequest(error_msg)

    if not current_app.config['EXPORT_URL_GCS_BUCKET']:
        error_msg = (
            'Project not set in deploy.json or export URL GCS bucket not '
            'found. Export to Saturn feature will not work. '
            'See https://github.com/DataBiosphere/data-explorer#one-time-setup-for-export-to-saturn-feature'
        )
        current_app.logger.error(error_msg)
        raise BadRequest(error_msg)

    private_key_path = os.path.join(current_app.config['DATASET_CONFIG_DIR'],
                                    'private-key.json')
    if not os.path.isfile(private_key_path):
        error_msg = (
            'Private key not found. Export to Saturn feature will not work. '
            'See https://github.com/DataBiosphere/data-explorer#one-time-setup-for-export-to-saturn-feature'
        )
        current_app.logger.error(error_msg)
        raise BadRequest(error_msg)


def _get_doc_generator(filter_arr):
    if len(filter_arr) == 0:
        return

    es = Elasticsearch(current_app.config['ELASTICSEARCH_URL'])
    es_facets = OrderedDict(current_app.config['ELASTICSEARCH_FACETS'].items())
    filters = elasticsearch_util.get_facet_value_dict(filter_arr, es_facets)
    search_dict = DatasetFacetedSearch(filters,
                                       es_facets).build_search().to_dict().get(
                                           'post_filter', {})
    search = Search(using=es, index=current_app.config['INDEX_NAME'])
    search.update_from_dict({'post_filter': search_dict})
    for result in search.scan():
        yield result.to_dict()


def _get_entities_dict(cohort_name, query, filter_arr):
    """Returns a dict representing the JSON list of entities."""
    # Saturn add-import expects a JSON list of entities, where each entity is
    # the entity JSON passed into
    # https://rawls.dsde-prod.broadinstitute.org/#!/entities/create_entity
    entities = []
    for table_name in current_app.config['TABLE_NAMES']:
        splits = table_name.split('.')
        if len(splits) != 3:
            raise BadRequest(
                'Unknown format for table name %s. Expected BigQuery project_id.dataset_id.table_name'
            )
        entities.append({
            # FireCloud doesn't allow spaces, so use underscore.
            'entityType': 'BigQuery_table',
            # This is the entity ID. Ideally this would be
            # project_id.dataset_id.table_name, and we wouldn't need the
            # table_name attribute. Unfortunately RAWLS doesn't allow
            # periods here. RAWLS does allow periods in attributes. So use
            # underscores here and periods in table_name attribute.
            'name': table_name.replace('.', '_'),
            'attributes': {
                'table_name': table_name
            }
        })

    # If a cohort was selected, create a query entity and get Elasticsearch documents
    # for the cohort so we can create sample_set entity.
    if cohort_name:
        entities.append({
            'entityType': 'cohort',
            'name': cohort_name,
            'attributes': {
                'query': query
            }
        })

        sample_items = []
        for doc in _get_doc_generator(filter_arr):
            sample_items.extend([{
                'entityType': 'sample',
                'entityName': s['sample_id']
            } for s in doc.get('samples', [])])
        entities.append({
            'entityType': 'sample_set',
            'name': cohort_name,
            'attributes': {
                'samples': {
                    'itemsType': 'EntityReference',
                    'items': sample_items
                }
            }
        })

    return entities


def _random_str():
    # Random 10 character string
    return ''.join(
        random.choice(string.ascii_letters + string.digits) for _ in range(10))


def _write_gcs_file(entities):
    """Returns GCS file path of the format /bucket/object."""
    client = storage.Client(project=current_app.config['DEPLOY_PROJECT_ID'])
    export_bucket = client.get_bucket(
        current_app.config['EXPORT_URL_GCS_BUCKET'])
    samples_bucket = client.get_bucket(
        current_app.config['EXPORT_URL_SAMPLES_GCS_BUCKET'])

    # Copy the samples blob to the export bucket in order to compose with the other
    # object containing the rest of the entities JSON.
    samples_blob = samples_bucket.get_blob('samples')
    samples_blob = samples_bucket.copy_blob(samples_blob, export_bucket)

    blob = export_bucket.blob(_random_str())
    entities_json = json.dumps(entities)
    # Remove the leading '[' character since this is being concatenated with the
    # sample entities JSON, which has the trailing ']' stripped in the indexer.
    entities_json = ',%s' % entities_json[1:]
    blob.upload_from_string(entities_json)

    merged = export_bucket.blob(_random_str())
    merged.upload_from_string('')
    merged.compose([samples_blob, blob])

    current_app.logger.info(
        'Wrote gs://%s/%s' % (current_app.config['EXPORT_URL_GCS_BUCKET'],
                              merged.name))<commit_msg>Use rewrite API to copy samples JSON file (#181)<commit_after>import base64
import connexion
import six
import json
import os
import random
import string
import sys
import time
import urllib

from collections import OrderedDict
from elasticsearch import Elasticsearch
from elasticsearch_dsl import Search
from flask import request
from flask import current_app
from werkzeug.exceptions import BadRequest
from google.cloud import storage
from oauth2client.service_account import ServiceAccountCredentials

from data_explorer.models.export_url_response import ExportUrlResponse  # noqa: E501
from data_explorer.util import elasticsearch_util
from data_explorer.util.dataset_faceted_search import DatasetFacetedSearch

# Export to Saturn flow
# - User clicks export button on bottom right of Data Explorer
# - A dialog pops up, asking user to name cohort. User types "African females"
#   and clicks Export button
# - UI server calls API server /export_url (this file):
#   - Constructs a JSON array of JSON entity objects
#   - Writes JSON array to a GCS file
#   - Concat that with an existing samples JSON entity file, which contains
#     all the samples in the dataset and is created on indexing.
#   - Creates a signed url for GCS file. Returns signed url to UI server
# - UI server redirects to Saturn add-import?url=SIGNED_URL
# - On add-import page, user selects Workspace. User clicks Import button.
# - User is redirected to selected workspace Data tab, showing newly imported
#   entities.


def _check_preconditions():
    config_path = os.path.join(current_app.config['DATASET_CONFIG_DIR'],
                               'deploy.json')
    if not os.path.isfile(config_path):
        error_msg = (
            'deploy.json not found. Export to Saturn feature will not work. '
            'See https://github.com/DataBiosphere/data-explorer#one-time-setup-for-export-to-saturn-feature'
        )
        current_app.logger.error(error_msg)
        raise BadRequest(error_msg)

    if not current_app.config['DEPLOY_PROJECT_ID']:
        error_msg = (
            'Project not set in deploy.json. Export to Saturn feature will not work. '
            'See https://github.com/DataBiosphere/data-explorer#one-time-setup-for-export-to-saturn-feature-for-export-to-saturn-feature'
        )
        current_app.logger.error(error_msg)
        raise BadRequest(error_msg)

    if not current_app.config['EXPORT_URL_GCS_BUCKET']:
        error_msg = (
            'Project not set in deploy.json or export URL GCS bucket not '
            'found. Export to Saturn feature will not work. '
            'See https://github.com/DataBiosphere/data-explorer#one-time-setup-for-export-to-saturn-feature'
        )
        current_app.logger.error(error_msg)
        raise BadRequest(error_msg)

    private_key_path = os.path.join(current_app.config['DATASET_CONFIG_DIR'],
                                    'private-key.json')
    if not os.path.isfile(private_key_path):
        error_msg = (
            'Private key not found. Export to Saturn feature will not work. '
            'See https://github.com/DataBiosphere/data-explorer#one-time-setup-for-export-to-saturn-feature'
        )
        current_app.logger.error(error_msg)
        raise BadRequest(error_msg)


def _get_doc_generator(filter_arr):
    if len(filter_arr) == 0:
        return

    es = Elasticsearch(current_app.config['ELASTICSEARCH_URL'])
    es_facets = OrderedDict(current_app.config['ELASTICSEARCH_FACETS'].items())
    filters = elasticsearch_util.get_facet_value_dict(filter_arr, es_facets)
    search_dict = DatasetFacetedSearch(filters,
                                       es_facets).build_search().to_dict().get(
                                           'post_filter', {})
    search = Search(using=es, index=current_app.config['INDEX_NAME'])
    search.update_from_dict({'post_filter': search_dict})
    for result in search.scan():
        yield result.to_dict()


def _get_entities_dict(cohort_name, query, filter_arr):
    """Returns a dict representing the JSON list of entities."""
    # Saturn add-import expects a JSON list of entities, where each entity is
    # the entity JSON passed into
    # https://rawls.dsde-prod.broadinstitute.org/#!/entities/create_entity
    entities = []
    for table_name in current_app.config['TABLE_NAMES']:
        splits = table_name.split('.')
        if len(splits) != 3:
            raise BadRequest(
                'Unknown format for table name %s. Expected BigQuery project_id.dataset_id.table_name'
            )
        entities.append({
            # FireCloud doesn't allow spaces, so use underscore.
            'entityType': 'BigQuery_table',
            # This is the entity ID. Ideally this would be
            # project_id.dataset_id.table_name, and we wouldn't need the
            # table_name attribute. Unfortunately RAWLS doesn't allow
            # periods here. RAWLS does allow periods in attributes. So use
            # underscores here and periods in table_name attribute.
            'name': table_name.replace('.', '_'),
            'attributes': {
                'table_name': table_name
            }
        })

    # If a cohort was selected, create a query entity and get Elasticsearch documents
    # for the cohort so we can create sample_set entity.
    if cohort_name:
        entities.append({
            'entityType': 'cohort',
            'name': cohort_name,
            'attributes': {
                'query': query
            }
        })

        sample_items = []
        for doc in _get_doc_generator(filter_arr):
            sample_items.extend([{
                'entityType': 'sample',
                'entityName': s['sample_id']
            } for s in doc.get('samples', [])])
        entities.append({
            'entityType': 'sample_set',
            'name': cohort_name,
            'attributes': {
                'samples': {
                    'itemsType': 'EntityReference',
                    'items': sample_items
                }
            }
        })

    return entities


def _random_str():
    # Random 10 character string
    return ''.join(
        random.choice(string.ascii_letters + string.digits) for _ in range(10))


def _write_gcs_file(entities):
    """Returns GCS file path of the format /bucket/object."""
    client = storage.Client(project=current_app.config['DEPLOY_PROJECT_ID'])
    export_bucket = client.get_bucket(
        current_app.config['EXPORT_URL_GCS_BUCKET'])
    samples_bucket = client.get_bucket(
        current_app.config['EXPORT_URL_SAMPLES_GCS_BUCKET'])

    # Copy the samples blob to the export bucket in order to compose with the other
    # object containing the rest of the entities JSON.
    samples_blob = samples_bucket.get_blob('samples')
    copied_samples_blob = export_bucket.blob('samples')
    # Use the rewrite rather than the copy API because the copy can timeout.
    copied_samples_blob.rewrite(samples_blob)

    blob = export_bucket.blob(_random_str())
    entities_json = json.dumps(entities)
    # Remove the leading '[' character since this is being concatenated with the
    # sample entities JSON, which has the trailing ']' stripped in the indexer.
    entities_json = ',%s' % entities_json[1:]
    blob.upload_from_string(entities_json)

    merged = export_bucket.blob(_random_str())
    merged.upload_from_string('')
    merged.compose([copied_samples_blob, blob])

    current_app.logger.info(
        'Wrote gs://%s/%s' % (current_app.config['EXPORT_URL_GCS_BUCKET'],
                              merged.name))
{'fraction_non_alphanumeric': 0.08893175468937076, 'fraction_numerical': 0.0019289610216841825, 'mean_word_length': 2.655482616095308, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 3, 'https://': 10, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 8, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '1082500', 'n_tokens_mistral': 4345, 'n_tokens_neox': 4200, 'n_words': 1160}