<commit_before>#!/usr/bin/env python3
# coding: utf8
import os.path

from . import Utils
from .GooglePhotosMedia import GooglePhotosMedia
from .BaseMedia import MediaType
from .LocalData import LocalData
from .DatabaseMedia import DatabaseMedia
import logging
import requests
import shutil
import tempfile
# apparently this is undocumented and may end up in Threading - could use Pool instead
# but have no need for separate processes since our workers are (probably) IO bound
from multiprocessing.pool import ThreadPool

log = logging.getLogger('gphotos.Photos')

class GooglePhotosSync(object):
    PAGE_SIZE = 100

    def __init__(self, api, root_folder, db):
        """
        :param (RestClient) api
        :param (str) root_folder:
        :param (LocalData) db:
        """
        self._db = db
        self._root_folder = root_folder
        self._api = api
        self._media_folder = 'photos'
        self.download_pool = ThreadPool(10)

        self._latest_download = Utils.minimum_date()
        # properties to be set after init
        # thus in theory one instance could so multiple indexes
        self._startDate = None
        self._endDate = None
        self.includeVideo = True

    @property
    def start_date(self):
        return self._startDate

    @start_date.setter
    def start_date(self, val):
        if val:
            self._startDate = Utils.string_to_date(val)

    @property
    def end_date(self):
        return self._endDate

    @end_date.setter
    def end_date(self, val):
        if val:
            self._endDate = Utils.string_to_date(val)

    @property
    def latest_download(self):
        return self._latest_download

    # todo this will currently do nothing unless using --flush-db
    #  need to look for a 'changes' api if we want to properly support deletes
    def check_for_removed(self):
        # note for partial scans using date filters this is still OK because
        # for a file to exist it must have been indexed in a previous scan
        log.info(u'Finding deleted media ...')
        for (dir_name, _, file_names) in os.walk(self._root_folder):
            for file_name in file_names:
                local_path = os.path.relpath(dir_name, self._root_folder)
                if file_name.startswith('.') or file_name.startswith('gphotos'):
                    continue
                file_row = self._db.get_file_by_path(local_path, file_name)
                if not file_row:
                    name = os.path.join(dir_name, file_name)
                    os.remove(name)
                    log.warning("%s deleted", name)

    def write_media_index(self, media, update=True):
        media.save_to_db(self._db, update)
        if media.modify_date > self._latest_download:
            self._latest_download = media.modify_date

    @classmethod
    def make_search_parameters(cls, page_token=None, start_date=None, end_date=None, do_video=False):
        # todo - instead of this crude code,
        #  should probably work out a nice way to read the REST schema into some useful dynamic classes
        class Y:
            def __init__(self, y, m, d):
                self.year = y
                self.month = m
                self.day = d

            def to_dict(self):
                return {"year": self.year, "month": self.month, "day": self.day}

        start = Y(1800, 1, 1)
        end = Y(3000, 1, 1)
        type_list = ["ALL_MEDIA"]

        if start_date:
            start = Y(start_date.year, start_date.month, start_date.day)
        if end_date:
            end = Y(end_date.year, end_date.month, end_date.day)
        if not do_video:
            type_list = ["PHOTO"]

        body = {
            'pageToken': page_token,
            'pageSize': 100,
            'filters': {
                'dateFilter': {
                    'ranges':
                        [
                            {'startDate': start.to_dict(),
                             'endDate': end.to_dict()
                             }
                        ]
                },
                'mediaTypeFilter': {'mediaTypes': type_list},
            }
        }
        return body

    def index_photos_media(self):
        log.info(u'Indexing Google Photos Files ...')
        count = 0
        try:
            body = self.make_search_parameters(start_date=self.start_date,
                                               end_date=self.end_date,
                                               do_video=self.includeVideo)
            response = self._api.mediaItems.search.execute(body)
            while response:
                items_json = response.json()
                media_json = items_json.get('mediaItems')
                # cope with empty response
                if not media_json:
                    break
                for media_item_json in media_json:
                    media_item = GooglePhotosMedia(media_item_json)
                    media_item.set_path_by_date(self._media_folder)
                    row = media_item.is_indexed(self._db)
                    if not row:
                        count += 1
                        log.info(u"Indexed %d %s", count, media_item.relative_path)
                        self.write_media_index(media_item, False)
                        if count % 1000 == 0:
                            self._db.store()
                    elif media_item.modify_date > row.ModifyDate:
                        log.info(u"Updated %s", media_item.relative_path)
                        self.write_media_index(media_item, True)
                    else:
                        log.debug(u"Skipped %s", media_item.relative_path)
                next_page = items_json.get('nextPageToken')
                if next_page:
                    body = self.make_search_parameters(page_token=next_page,
                                                       start_date=self.start_date,
                                                       end_date=self.end_date,
                                                       do_video=self.includeVideo)
                    response = self._api.mediaItems.search.execute(body)
                else:
                    break
        finally:
            # store latest date for incremental backup only if scanning all
            if not (self.start_date or self.end_date):
                self._db.set_scan_dates(drive_last_date=self._latest_download)

    @classmethod
    def do_download_file(cls, url, local_full_path, media_item):
        # this function runs in a process pool and does the actual downloads
        folder = os.path.dirname(local_full_path)
        with tempfile.NamedTemporaryFile(dir=folder, delete=False) as temp_file:
            r = requests.get(url, stream=True)
            shutil.copyfileobj(r.raw, temp_file)
            os.rename(temp_file.name, local_full_path)
        # set the access date to create date since there is nowhere
        # else to put it on linux (and is useful for debugging)
        os.utime(local_full_path,
                 (Utils.to_timestamp(media_item.modify_date),
                  Utils.to_timestamp(media_item.create_date)))

    def download_file(self, url=None, local_full_path=None, media_item=None):
        # this function farms downloads off to a thread pool
        self.download_pool.apply_async(self.do_download_file, (url, local_full_path, media_item))

    def download_photo_media(self):
        log.info(u'Downloading Photos ...')
        count = 0
        # noinspection PyTypeChecker
        for media_item in DatabaseMedia.get_media_by_search(self._db, media_type=MediaType.PHOTOS,
                                                            start_date=self.start_date, end_date=self.end_date):
            local_folder = os.path.join(self._root_folder, media_item.relative_folder)
            local_full_path = os.path.join(local_folder, media_item.filename)
            if os.path.exists(local_full_path):
                log.debug(u'skipping {} ...'.format(local_full_path))
                # todo is there anyway to detect remote updates with photos API?
                # if Utils.to_timestamp(media.modify_date) > \
                #         os.path.getctime(local_full_path):
                #     log.warning(u'{} was modified'.format(local_full_path))
                # else:
                continue

            if not os.path.isdir(local_folder):
                os.makedirs(local_folder)

            count += 1
            try:
                response = self._api.mediaItems.get.execute(mediaItemId=str(media_item.id))
                r_json = response.json()
                if media_item.is_video():
                    log.info(u'downloading video {} {} ...'.format(count, local_full_path))
                    download_url = '{}=dv'.format(r_json['baseUrl'])
                else:
                    log.info(u'downloading image {} {} ...'.format(count, local_full_path))
                    download_url = '{}=d'.format(r_json['baseUrl'])
                self.download_file(download_url, local_full_path, media_item)

            except Exception as e:
                log.error('failure downloading {}.\n{}{}'.format(local_full_path, type(e), e))
        # allow any remaining background downloads to complete
        self.download_pool.close()
        self.download_pool.join()
<commit_msg>batching media gets for super speedy downloads<commit_after>#!/usr/bin/env python3
# coding: utf8
import os.path

from . import Utils
from .GooglePhotosMedia import GooglePhotosMedia
from .BaseMedia import MediaType
from .LocalData import LocalData
from .DatabaseMedia import DatabaseMedia
from itertools import zip_longest
import logging
import requests
import shutil
import tempfile
# apparently this is undocumented and may end up in Threading - could use Pool instead
# but have no need for separate processes since our workers are (probably) IO bound
from multiprocessing.pool import ThreadPool

log = logging.getLogger('gphotos.Photos')

class GooglePhotosSync(object):
    PAGE_SIZE = 100

    def __init__(self, api, root_folder, db):
        """
        :param (RestClient) api
        :param (str) root_folder:
        :param (LocalData) db:
        """
        self._db = db
        self._root_folder = root_folder
        self._api = api
        self._media_folder = 'photos'
        self.download_pool = ThreadPool(10)

        self._latest_download = Utils.minimum_date()
        # properties to be set after init
        # thus in theory one instance could so multiple indexes
        self._startDate = None
        self._endDate = None
        self.includeVideo = True

    @property
    def start_date(self):
        return self._startDate

    @start_date.setter
    def start_date(self, val):
        if val:
            self._startDate = Utils.string_to_date(val)

    @property
    def end_date(self):
        return self._endDate

    @end_date.setter
    def end_date(self, val):
        if val:
            self._endDate = Utils.string_to_date(val)

    @property
    def latest_download(self):
        return self._latest_download

    # todo this will currently do nothing unless using --flush-db
    #  need to look for a 'changes' api if we want to properly support deletes
    def check_for_removed(self):
        # note for partial scans using date filters this is still OK because
        # for a file to exist it must have been indexed in a previous scan
        log.info(u'Finding deleted media ...')
        for (dir_name, _, file_names) in os.walk(self._root_folder):
            for file_name in file_names:
                local_path = os.path.relpath(dir_name, self._root_folder)
                if file_name.startswith('.') or file_name.startswith('gphotos'):
                    continue
                file_row = self._db.get_file_by_path(local_path, file_name)
                if not file_row:
                    name = os.path.join(dir_name, file_name)
                    os.remove(name)
                    log.warning("%s deleted", name)

    def write_media_index(self, media, update=True):
        media.save_to_db(self._db, update)
        if media.modify_date > self._latest_download:
            self._latest_download = media.modify_date

    @classmethod
    def make_search_parameters(cls, page_token=None, start_date=None, end_date=None, do_video=False):
        # todo - instead of this crude code,
        #  should probably work out a nice way to read the REST schema into some useful dynamic classes
        class Y:
            def __init__(self, y, m, d):
                self.year = y
                self.month = m
                self.day = d

            def to_dict(self):
                return {"year": self.year, "month": self.month, "day": self.day}

        start = Y(1800, 1, 1)
        end = Y(3000, 1, 1)
        type_list = ["ALL_MEDIA"]

        if start_date:
            start = Y(start_date.year, start_date.month, start_date.day)
        if end_date:
            end = Y(end_date.year, end_date.month, end_date.day)
        if not do_video:
            type_list = ["PHOTO"]

        body = {
            'pageToken': page_token,
            'pageSize': 100,
            'filters': {
                'dateFilter': {
                    'ranges':
                        [
                            {'startDate': start.to_dict(),
                             'endDate': end.to_dict()
                             }
                        ]
                },
                'mediaTypeFilter': {'mediaTypes': type_list},
            }
        }
        return body

    def index_photos_media(self):
        log.info(u'Indexing Google Photos Files ...')
        count = 0
        try:
            body = self.make_search_parameters(start_date=self.start_date,
                                               end_date=self.end_date,
                                               do_video=self.includeVideo)
            response = self._api.mediaItems.search.execute(body)
            while response:
                items_json = response.json()
                media_json = items_json.get('mediaItems')
                # cope with empty response
                if not media_json:
                    break
                for media_item_json in media_json:
                    media_item = GooglePhotosMedia(media_item_json)
                    media_item.set_path_by_date(self._media_folder)
                    row = media_item.is_indexed(self._db)
                    if not row:
                        count += 1
                        log.info(u"Indexed %d %s", count, media_item.relative_path)
                        self.write_media_index(media_item, False)
                        if count % 1000 == 0:
                            self._db.store()
                    elif media_item.modify_date > row.ModifyDate:
                        log.info(u"Updated %s", media_item.relative_path)
                        self.write_media_index(media_item, True)
                    else:
                        log.debug(u"Skipped %s", media_item.relative_path)
                next_page = items_json.get('nextPageToken')
                if next_page:
                    body = self.make_search_parameters(page_token=next_page,
                                                       start_date=self.start_date,
                                                       end_date=self.end_date,
                                                       do_video=self.includeVideo)
                    response = self._api.mediaItems.search.execute(body)
                else:
                    break
        finally:
            # store latest date for incremental backup only if scanning all
            if not (self.start_date or self.end_date):
                self._db.set_scan_dates(drive_last_date=self._latest_download)

    @classmethod
    def do_download_file(cls, url, local_full_path, media_item):
        # this function runs in a process pool and does the actual downloads
        folder = os.path.dirname(local_full_path)
        with tempfile.NamedTemporaryFile(dir=folder, delete=False) as temp_file:
            r = requests.get(url, stream=True)
            shutil.copyfileobj(r.raw, temp_file)
            os.rename(temp_file.name, local_full_path)
        # set the access date to create date since there is nowhere
        # else to put it on linux (and is useful for debugging)
        os.utime(local_full_path,
                 (Utils.to_timestamp(media_item.modify_date),
                  Utils.to_timestamp(media_item.create_date)))

    def download_file(self, url=None, local_full_path=None, media_item=None):
        # this function farms downloads off to a thread pool
        # todo add error response function
        self.download_pool.apply_async(self.do_download_file, (url, local_full_path, media_item))

    def download_photo_media(self):
        def grouper(iterable):
            """Collect data into chunks size 20"""
            return zip_longest(*[iter(iterable)] * 20, fillvalue=None)

        log.info(u'Downloading Photos ...')
        count = 0
        for media_items_block in grouper(
                DatabaseMedia.get_media_by_search(self._db, media_type=MediaType.PHOTOS,
                                                  start_date=self.start_date, end_date=self.end_date)):
            batch_ids = []
            for media_item in media_items_block:
                if media_item is None:
                    break

                local_folder = os.path.join(self._root_folder, media_item.relative_folder)
                local_full_path = os.path.join(local_folder, media_item.filename)
                if os.path.exists(local_full_path):
                    log.debug(u'skipping {} ...'.format(local_full_path))
                    # todo is there anyway to detect remote updates with photos API?
                    # if Utils.to_timestamp(media.modify_date) > \
                    #         os.path.getctime(local_full_path):
                    #     log.warning(u'{} was modified'.format(local_full_path))
                    # else:
                    continue

                if not os.path.isdir(local_folder):
                    os.makedirs(local_folder)
                batch_ids.append(media_item.id)

            if len(batch_ids) == 0:
                continue

            try:
                response = self._api.mediaItems.batchGet.execute(mediaItemIds=batch_ids)
                r_json = response.json()
                for media_item_json_status in r_json["mediaItemResults"]:
                    count += 1
                    # todo look at media_item_json_status["status"] for individual errors
                    media_item_json = media_item_json_status["mediaItem"]
                    media_item = GooglePhotosMedia(media_item_json)
                    media_item.set_path_by_date(self._media_folder)
                    try:
                        local_folder = os.path.join(self._root_folder, media_item.relative_folder)
                        local_full_path = os.path.join(local_folder, media_item.filename)
                        if media_item.is_video():
                            log.info(u'downloading video {} {} ...'.format(count, local_full_path))
                            download_url = '{}=dv'.format(media_item_json['baseUrl'])
                        else:
                            log.info(u'downloading image {} {} ...'.format(count, local_full_path))
                            download_url = '{}=d'.format(media_item_json['baseUrl'])
                        self.download_file(download_url, local_full_path, media_item)
                    except Exception as e:
                        log.error('failure downloading of {}.\n{}{}'.format(media_item.filename, type(e), e))
            except Exception as e:
                log.error('failure in batch get of {}.\n{}{}'.format(batch_ids, type(e), e))

        # allow any remaining background downloads to complete
        self.download_pool.close()
        self.download_pool.join()
<|endoftext|>


<commit_before>"""Cinder tests."""
# Copyright 2014 - 2015 Solinea, Inc.
#
# Licensed under the Solinea Software License Agreement (goldstone),
# Version 1.0 (the "License"); you may not use this file except in compliance
# with the License. You may obtain a copy of the License at:
#
#     http://www.solinea.com/goldstone/LICENSE.pdf
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either expressed or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
from django.test import SimpleTestCase
from mock import patch

from goldstone.apps.cinder.tasks import time_service_list
from goldstone.test_utils import create_and_login, AUTHORIZATION_PAYLOAD

class TaskTests(SimpleTestCase):
    """Test cinder tasks."""

    @patch('goldstone.apps.cinder.tasks.time_api_call')
    @patch('goldstone.apps.cinder.tasks.stack_api_request_base')
    def test_time_image_list_api(self, m_base, m_time_api_call):

        import requests
        from requests import Response

        response = Response()
        response._content = '{"services": [{"id": 1}]}'
        response.status_code = requests.codes.ok
        m_base.return_value = {'url': 'http://url', 'headers': {}}
        m_time_api_call.return_value = {'created': True,
                                        'response': response}
        result = time_service_list()
        self.assertEqual(m_time_api_call.call_count, 1)
        self.assertEqual(result, m_time_api_call.return_value)

class ViewTests(SimpleTestCase):
    """Test api_perf."""

    def test_report_view(self):
        uri = '/cinder/report'
        response = self.client.get(uri)
        self.assertEqual(response.status_code, 200)
        self.assertTemplateUsed(response, 'cinder_report.html')

class DataViewTests(SimpleTestCase):
    """Test a grabbag of cinder API endpoints."""

    def setUp(self):
        """Run before every test."""

        get_user_model().objects.all().delete()
        self.token = create_and_login()

    def _evaluate(self, response):
        """Check a test's results."""
        import json
        from django.http import HttpResponse

        self.assertIsInstance(response, HttpResponse)
        self.assertNotEqual(response.content, None)

        try:
            results = json.loads(response.content)
        except Exception:          # pylint: disable=W0703
            self.fail("Could not convert content to JSON, content was %s" %
                      response.content)
        else:
            self.assertIsInstance(results, list)
            self.assertGreaterEqual(len(results), 1)
            self.assertIsInstance(results[0], list)

    def test_get_volumes(self):
        self._evaluate(self.client.get("/cinder/volumes",
            HTTP_AUTHORIZATION=AUTHORIZATION_PAYLOAD % self.token)

    def test_get_backups(self):
        self._evaluate(self.client.get(
            "/cinder/backups",
            HTTP_AUTHORIZATION=AUTHORIZATION_PAYLOAD % self.token))

    def test_get_snapshots(self):
        self._evaluate(self.client.get(
            "/cinder/snapshots",
            HTTP_AUTHORIZATION=AUTHORIZATION_PAYLOAD % self.token))

    def test_get_services(self):
        self._evaluate(self.client.get(
            "/cinder/services",
            HTTP_AUTHORIZATION=AUTHORIZATION_PAYLOAD % self.token))

    def test_get_volume_types(self):
        self._evaluate(self.client.get(
            "/cinder/volume_types",
            HTTP_AUTHORIZATION=AUTHORIZATION_PAYLOAD % self.token))

    def test_get_transfers(self):
        self._evaluate(self.client.get(
            "/cinder/transfers",
            HTTP_AUTHORIZATION=AUTHORIZATION_PAYLOAD % self.token))
<commit_msg>GOLD-639: Converted syncdb, migrate, and runserver to fabric commands. Split the custom User model into its own application to solve a circular import bug. Bug, AUTH_USER_MODEL had bad value. API change, /accounts/me is now /user/me.<commit_after>"""Cinder tests."""
# Copyright 2014 - 2015 Solinea, Inc.
#
# Licensed under the Solinea Software License Agreement (goldstone),
# Version 1.0 (the "License"); you may not use this file except in compliance
# with the License. You may obtain a copy of the License at:
#
#     http://www.solinea.com/goldstone/LICENSE.pdf
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either expressed or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
from django.test import SimpleTestCase
from mock import patch

from goldstone.apps.cinder.tasks import time_service_list
from goldstone.test_utils import create_and_login, AUTHORIZATION_PAYLOAD

class TaskTests(SimpleTestCase):
    """Test cinder tasks."""

    @patch('goldstone.apps.cinder.tasks.time_api_call')
    @patch('goldstone.apps.cinder.tasks.stack_api_request_base')
    def test_time_image_list_api(self, m_base, m_time_api_call):

        import requests
        from requests import Response

        response = Response()
        response._content = '{"services": [{"id": 1}]}'
        response.status_code = requests.codes.ok
        m_base.return_value = {'url': 'http://url', 'headers': {}}
        m_time_api_call.return_value = {'created': True,
                                        'response': response}
        result = time_service_list()
        self.assertEqual(m_time_api_call.call_count, 1)
        self.assertEqual(result, m_time_api_call.return_value)

class ViewTests(SimpleTestCase):
    """Test api_perf."""

    def test_report_view(self):
        uri = '/cinder/report'
        response = self.client.get(uri)
        self.assertEqual(response.status_code, 200)
        self.assertTemplateUsed(response, 'cinder_report.html')

class DataViewTests(SimpleTestCase):
    """Test a grabbag of cinder API endpoints."""

    def setUp(self):
        """Run before every test."""

        get_user_model().objects.all().delete()
        self.token = create_and_login()

    def _evaluate(self, response):
        """Check a test's results."""
        import json
        from django.http import HttpResponse

        self.assertIsInstance(response, HttpResponse)
        self.assertNotEqual(response.content, None)

        try:
            results = json.loads(response.content)
        except Exception:          # pylint: disable=W0703
            self.fail("Could not convert content to JSON, content was %s" %
                      response.content)
        else:
            self.assertIsInstance(results, list)
            self.assertGreaterEqual(len(results), 1)
            self.assertIsInstance(results[0], list)

    def test_get_volumes(self):
        self._evaluate(self.client.get("/cinder/volumes",
            HTTP_AUTHORIZATION=AUTHORIZATION_PAYLOAD % self.token))

    def test_get_backups(self):
        self._evaluate(self.client.get(
            "/cinder/backups",
            HTTP_AUTHORIZATION=AUTHORIZATION_PAYLOAD % self.token))

    def test_get_snapshots(self):
        self._evaluate(self.client.get(
            "/cinder/snapshots",
            HTTP_AUTHORIZATION=AUTHORIZATION_PAYLOAD % self.token))

    def test_get_services(self):
        self._evaluate(self.client.get(
            "/cinder/services",
            HTTP_AUTHORIZATION=AUTHORIZATION_PAYLOAD % self.token))

    def test_get_volume_types(self):
        self._evaluate(self.client.get(
            "/cinder/volume_types",
            HTTP_AUTHORIZATION=AUTHORIZATION_PAYLOAD % self.token))

    def test_get_transfers(self):
        self._evaluate(self.client.get(
            "/cinder/transfers",
            HTTP_AUTHORIZATION=AUTHORIZATION_PAYLOAD % self.token))
<|endoftext|>


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

import warnings
import torch
from ..distributions import MultivariateNormal, MultitaskMultivariateNormal
from ..likelihoods import _GaussianLikelihoodBase
from .. import settings
from .gp import GP
from .exact_prediction_strategies import prediction_strategy

class ExactGP(GP):
    def __init__(self, train_inputs, train_targets, likelihood):
        if train_inputs is not None and torch.is_tensor(train_inputs):
            train_inputs = (train_inputs,)
        if train_inputs is not None and not all(torch.is_tensor(train_input) for train_input in train_inputs):
            raise RuntimeError("Train inputs must be a tensor, or a list/tuple of tensors")
        if not isinstance(likelihood, _GaussianLikelihoodBase):
            raise RuntimeError("ExactGP can only handle Gaussian likelihoods")

        super(ExactGP, self).__init__()
        if train_inputs is not None:
            self.train_inputs = tuple(tri.unsqueeze(-1) if tri.ndimension() == 1 else tri for tri in train_inputs)
            self.train_targets = train_targets
        else:
            self.train_inputs = None
            self.train_targets = None
        self.likelihood = likelihood

        self.prediction_strategy = None

    def _apply(self, fn):
        if self.train_inputs is not None:
            self.train_inputs = tuple(fn(train_input) for train_input in self.train_inputs)
            self.train_targets = fn(self.train_targets)
        return super(ExactGP, self)._apply(fn)

    def set_train_data(self, inputs=None, targets=None, strict=True):
        """Set training data (does not re-fit model hyper-parameters)"""
        if inputs is not None:
            if torch.is_tensor(inputs):
                inputs = (inputs,)
            inputs = tuple(input_.unsqueeze(-1) if input_.ndimension() == 1 else input_ for input_ in inputs)
            for input, t_input in zip(inputs, self.train_inputs):
                for attr in {"shape", "dtype", "device"}:
                    if strict and getattr(input, attr) != getattr(t_input, attr):
                        raise RuntimeError("Cannot modify {attr} of inputs".format(attr=attr))
            self.train_inputs = inputs
        if targets is not None:
            for attr in {"shape", "dtype", "device"}:
                if strict and getattr(targets, attr) != getattr(self.train_targets, attr):
                    raise RuntimeError("Cannot modify {attr} of targets".format(attr=attr))
            self.train_targets = targets
        self.prediction_strategy = None

    def train(self, mode=True):
        if mode:
            self.prediction_strategy = None
        return super(ExactGP, self).train(mode)

    def __call__(self, *args, **kwargs):
        train_inputs = list(self.train_inputs) if self.train_inputs is not None else []
        inputs = list(i.unsqueeze(-1) if i.ndimension() == 1 else i for i in args)

        # Training mode: optimizing
        if self.training:
            if self.train_inputs is None:
                raise RuntimeError(
                    "train_inputs, train_targets cannot be None in training mode. "
                    "Call .eval() for prior predictions, or call .set_train_data() to add training data."
                )
            if settings.check_training_data.on():
                if not all(torch.equal(train_input, input) for train_input, input in zip(train_inputs, inputs)):
                    raise RuntimeError("You must train on the training inputs!")
            res = super(ExactGP, self).__call__(*inputs, **kwargs)
            return res

        # Prior mode
        elif self.train_inputs is None or self.train_targets is None:
            full_inputs = args
            full_output = super(ExactGP, self).__call__(*full_inputs, **kwargs)
            if settings.debug().on():
                if not isinstance(full_output, MultivariateNormal):
                    raise RuntimeError("ExactGP.forward must return a MultivariateNormal")
            return full_output
        # Posterior mode
        else:
            inputs = list(i.squeeze(0) if i.ndimension() == 3 else i for i in inputs)
            if settings.debug.on():
                if all(torch.equal(train_input, input) for train_input, input in zip(train_inputs, inputs)):
                    warnings.warn(
                        "The input matches the stored training data. Did you forget to call model.train()?", UserWarning
                    )

            # Exact inference
            non_batch_train = False

            if all(tin.dim() == 2 for tin in self.train_inputs):
                # Train inputs were non-batch
                non_batch_train = True
            # Expand train_input to match input's batch size, or vice versa if they don't match.
            for i, (train_input, input) in enumerate(zip(train_inputs, inputs)):
                if train_input.dim() < input.dim():
                    batch_shape = inputs[0].shape[:-2]
                    train_inputs[i] = train_input.expand(*batch_shape, *train_input.shape[-2:])
                elif input.dim() < train_input.dim():
                    batch_shape = train_inputs[0].shape[:-2]
                    inputs[i] = input.expand(*batch_shape, *input.shape[-2:])
                else:
                    batch_shape = train_inputs[0].shape[:-2]

            full_inputs = tuple(
                torch.cat([train_input, input], dim=-2) for train_input, input in zip(train_inputs, inputs)
            )

            full_output = super(ExactGP, self).__call__(*full_inputs, **kwargs)
            if settings.debug().on():
                if not isinstance(full_output, MultivariateNormal):
                    raise RuntimeError("ExactGP.forward must return a MultivariateNormal")
            full_mean, full_covar = full_output.mean, full_output.lazy_covariance_matrix

            train_targets = self.train_targets
            if any(
                orig_train_input.dim() < train_input.dim()
                for orig_train_input, train_input in zip(self.train_inputs, train_inputs)
            ):
                train_targets = train_targets.unsqueeze(0).expand(*batch_shape, *train_targets.size())

            num_tasks = 1
            if isinstance(full_output, MultitaskMultivariateNormal):
                num_tasks = full_output.num_tasks

            full_mean = full_mean.view(*batch_shape, -1)
            num_train = train_targets.size(len(batch_shape))
            train_targets = train_targets.view(*batch_shape, -1)

            train_mean = full_mean.narrow(-1, 0, train_targets.size(-1))

            if self.prediction_strategy is None:
                train_train_covar = full_covar[..., :num_train, :num_train].evaluate_kernel()
                self.prediction_strategy = prediction_strategy(
                    num_train,
                    train_inputs,
                    train_mean,
                    train_train_covar,
                    train_targets,
                    self.likelihood,
                    non_batch_train,
                )

            if train_inputs[0].shape != self.prediction_strategy.train_inputs[0].shape:
                # The test batch shape has changed, update prediction strategy with expanded objects
                batch_shape = train_inputs[0].shape[:-2]
                train_train_covar = self.prediction_strategy.train_train_covar
                if len(batch_shape) > len(train_train_covar.batch_shape):
                    # Expanding to add more batches
                    expanded_covar = train_train_covar.expand(*batch_shape, *train_train_covar.matrix_shape)
                elif (
                    len(batch_shape) == len(train_train_covar.batch_shape)
                    and batch_shape.numel() != train_train_covar.batch_shape.numel()
                ):
                    # The test batch size has changed, we need to repeat it to a new batch size.
                    expanded_covar = train_train_covar[0].expand(*batch_shape, *train_train_covar.matrix_shape)
                else:
                    # We are leaving batch mode, not entering it.
                    expanded_covar = train_train_covar[0]
                self.prediction_strategy.train_train_covar = expanded_covar
                self.prediction_strategy.num_train = num_train
                self.prediction_strategy.train_targets = train_targets
                self.prediction_strategy.train_inputs = train_inputs
                self.prediction_strategy.train_mean = train_mean
                self.prediction_strategy.non_batch_train = non_batch_train

            test_mean = full_mean.narrow(-1, train_targets.size(-1), full_mean.size(-1) - train_targets.size(-1))
            test_test_covar = full_covar[..., num_train:, num_train:]
            test_train_covar = full_covar[..., num_train:, :num_train]

            predictive_mean = self.prediction_strategy.exact_predictive_mean(test_mean, test_train_covar)
            predictive_covar = self.prediction_strategy.exact_predictive_covar(test_test_covar, test_train_covar)

            if num_tasks > 1:
                if train_targets.ndimension() == 2:
                    # Batch multitask
                    predictive_mean = predictive_mean.view(train_targets.size(0), -1, num_tasks).contiguous()
                else:
                    # Standard multitask
                    predictive_mean = predictive_mean.view(-1, num_tasks).contiguous()

            return full_output.__class__(predictive_mean, predictive_covar)
<commit_msg>Explicitly support singleton batch dimensions<commit_after>#!/usr/bin/env python3

import warnings
import torch
from ..distributions import MultivariateNormal, MultitaskMultivariateNormal
from ..likelihoods import _GaussianLikelihoodBase
from .. import settings
from .gp import GP
from .exact_prediction_strategies import prediction_strategy

class ExactGP(GP):
    def __init__(self, train_inputs, train_targets, likelihood):
        if train_inputs is not None and torch.is_tensor(train_inputs):
            train_inputs = (train_inputs,)
        if train_inputs is not None and not all(torch.is_tensor(train_input) for train_input in train_inputs):
            raise RuntimeError("Train inputs must be a tensor, or a list/tuple of tensors")
        if not isinstance(likelihood, _GaussianLikelihoodBase):
            raise RuntimeError("ExactGP can only handle Gaussian likelihoods")

        super(ExactGP, self).__init__()
        if train_inputs is not None:
            self.train_inputs = tuple(tri.unsqueeze(-1) if tri.ndimension() == 1 else tri for tri in train_inputs)
            self.train_targets = train_targets
        else:
            self.train_inputs = None
            self.train_targets = None
        self.likelihood = likelihood

        self.prediction_strategy = None

    def _apply(self, fn):
        if self.train_inputs is not None:
            self.train_inputs = tuple(fn(train_input) for train_input in self.train_inputs)
            self.train_targets = fn(self.train_targets)
        return super(ExactGP, self)._apply(fn)

    def set_train_data(self, inputs=None, targets=None, strict=True):
        """Set training data (does not re-fit model hyper-parameters)"""
        if inputs is not None:
            if torch.is_tensor(inputs):
                inputs = (inputs,)
            inputs = tuple(input_.unsqueeze(-1) if input_.ndimension() == 1 else input_ for input_ in inputs)
            for input, t_input in zip(inputs, self.train_inputs):
                for attr in {"shape", "dtype", "device"}:
                    if strict and getattr(input, attr) != getattr(t_input, attr):
                        raise RuntimeError("Cannot modify {attr} of inputs".format(attr=attr))
            self.train_inputs = inputs
        if targets is not None:
            for attr in {"shape", "dtype", "device"}:
                if strict and getattr(targets, attr) != getattr(self.train_targets, attr):
                    raise RuntimeError("Cannot modify {attr} of targets".format(attr=attr))
            self.train_targets = targets
        self.prediction_strategy = None

    def train(self, mode=True):
        if mode:
            self.prediction_strategy = None
        return super(ExactGP, self).train(mode)

    def __call__(self, *args, **kwargs):
        train_inputs = list(self.train_inputs) if self.train_inputs is not None else []
        inputs = list(i.unsqueeze(-1) if i.ndimension() == 1 else i for i in args)

        # Training mode: optimizing
        if self.training:
            if self.train_inputs is None:
                raise RuntimeError(
                    "train_inputs, train_targets cannot be None in training mode. "
                    "Call .eval() for prior predictions, or call .set_train_data() to add training data."
                )
            if settings.check_training_data.on():
                if not all(torch.equal(train_input, input) for train_input, input in zip(train_inputs, inputs)):
                    raise RuntimeError("You must train on the training inputs!")
            res = super(ExactGP, self).__call__(*inputs, **kwargs)
            return res

        # Prior mode
        elif self.train_inputs is None or self.train_targets is None:
            full_inputs = args
            full_output = super(ExactGP, self).__call__(*full_inputs, **kwargs)
            if settings.debug().on():
                if not isinstance(full_output, MultivariateNormal):
                    raise RuntimeError("ExactGP.forward must return a MultivariateNormal")
            return full_output
        # Posterior mode
        else:
            if settings.debug.on():
                if all(torch.equal(train_input, input) for train_input, input in zip(train_inputs, inputs)):
                    warnings.warn(
                        "The input matches the stored training data. Did you forget to call model.train()?", UserWarning
                    )

            # Exact inference
            non_batch_train = False

            if all(tin.dim() == 2 for tin in self.train_inputs):
                # Train inputs were non-batch
                non_batch_train = True
            # Expand train_input to match input's batch size, or vice versa if they don't match.
            for i, (train_input, input) in enumerate(zip(train_inputs, inputs)):
                if train_input.dim() < input.dim():
                    batch_shape = inputs[0].shape[:-2]
                    train_inputs[i] = train_input.expand(*batch_shape, *train_input.shape[-2:])
                elif input.dim() < train_input.dim():
                    batch_shape = train_inputs[0].shape[:-2]
                    inputs[i] = input.expand(*batch_shape, *input.shape[-2:])
                else:
                    batch_shape = train_inputs[0].shape[:-2]

            full_inputs = tuple(
                torch.cat([train_input, input], dim=-2) for train_input, input in zip(train_inputs, inputs)
            )

            full_output = super(ExactGP, self).__call__(*full_inputs, **kwargs)
            if settings.debug().on():
                if not isinstance(full_output, MultivariateNormal):
                    raise RuntimeError("ExactGP.forward must return a MultivariateNormal")
            full_mean, full_covar = full_output.mean, full_output.lazy_covariance_matrix

            train_targets = self.train_targets
            if any(
                orig_train_input.dim() < train_input.dim()
                for orig_train_input, train_input in zip(self.train_inputs, train_inputs)
            ):
                train_targets = train_targets.unsqueeze(0).expand(*batch_shape, *train_targets.size())

            num_tasks = 1
            if isinstance(full_output, MultitaskMultivariateNormal):
                num_tasks = full_output.num_tasks

            full_mean = full_mean.view(*batch_shape, -1)
            num_train = train_targets.size(len(batch_shape))
            train_targets = train_targets.view(*batch_shape, -1)

            train_mean = full_mean.narrow(-1, 0, train_targets.size(-1))

            if self.prediction_strategy is None:
                train_train_covar = full_covar[..., :num_train, :num_train].evaluate_kernel()
                self.prediction_strategy = prediction_strategy(
                    num_train,
                    train_inputs,
                    train_mean,
                    train_train_covar,
                    train_targets,
                    self.likelihood,
                    non_batch_train,
                )

            if train_inputs[0].shape != self.prediction_strategy.train_inputs[0].shape:
                # The test batch shape has changed, update prediction strategy with expanded objects
                batch_shape = train_inputs[0].shape[:-2]
                train_train_covar = self.prediction_strategy.train_train_covar
                if len(batch_shape) > len(train_train_covar.batch_shape):
                    # Expanding to add more batches
                    expanded_covar = train_train_covar.expand(*batch_shape, *train_train_covar.matrix_shape)
                elif (
                    len(batch_shape) == len(train_train_covar.batch_shape)
                    and batch_shape.numel() != train_train_covar.batch_shape.numel()
                ):
                    # The test batch size has changed, we need to repeat it to a new batch size.
                    expanded_covar = train_train_covar[0].expand(*batch_shape, *train_train_covar.matrix_shape)
                else:
                    # We are leaving batch mode, not entering it.
                    expanded_covar = train_train_covar[0]
                self.prediction_strategy.train_train_covar = expanded_covar
                self.prediction_strategy.num_train = num_train
                self.prediction_strategy.train_targets = train_targets
                self.prediction_strategy.train_inputs = train_inputs
                self.prediction_strategy.train_mean = train_mean
                self.prediction_strategy.non_batch_train = non_batch_train

            test_mean = full_mean.narrow(-1, train_targets.size(-1), full_mean.size(-1) - train_targets.size(-1))
            test_test_covar = full_covar[..., num_train:, num_train:]
            test_train_covar = full_covar[..., num_train:, :num_train]

            predictive_mean = self.prediction_strategy.exact_predictive_mean(test_mean, test_train_covar)
            predictive_covar = self.prediction_strategy.exact_predictive_covar(test_test_covar, test_train_covar)

            if num_tasks > 1:
                if train_targets.ndimension() == 2:
                    # Batch multitask
                    predictive_mean = predictive_mean.view(train_targets.size(0), -1, num_tasks).contiguous()
                else:
                    # Standard multitask
                    predictive_mean = predictive_mean.view(-1, num_tasks).contiguous()

            return full_output.__class__(predictive_mean, predictive_covar)
<|endoftext|>


<commit_before># Copyright 2015 Google Inc.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""Define API Queries."""

import six

from google.cloud.bigquery._helpers import _TypedProperty
from google.cloud.bigquery._helpers import _rows_from_json
from google.cloud.bigquery.dataset import Dataset
from google.cloud.bigquery.job import QueryJob
from google.cloud.bigquery.table import _parse_schema_resource
from google.cloud.bigquery._helpers import QueryParametersProperty
from google.cloud.bigquery._helpers import UDFResourcesProperty

class _SyncQueryConfiguration(object):
    """User-settable configuration options for synchronous query jobs.

    Values which are ``None`` -> server defaults.
    """
    _default_dataset = None
    _dry_run = None
    _max_results = None
    _timeout_ms = None
    _preserve_nulls = None
    _use_query_cache = None
    _use_legacy_sql = None

class QueryResults(object):
    """Synchronous job: query tables.

    :type query: str
    :param query: SQL query string

    :type client: :class:`google.cloud.bigquery.client.Client`
    :param client: A client which holds credentials and project configuration
                   for the dataset (which requires a project).

    :type udf_resources: tuple
    :param udf_resources: An iterable of
                        :class:`google.cloud.bigquery.job.UDFResource`
                        (empty by default)

    :type query_parameters: tuple
    :param query_parameters:
        An iterable of
        :class:`google.cloud.bigquery._helpers.AbstractQueryParameter`
        (empty by default)
    """

    _UDF_KEY = 'userDefinedFunctionResources'
    _QUERY_PARAMETERS_KEY = 'queryParameters'

    def __init__(self, query, client, udf_resources=(), query_parameters=()):
        self._client = client
        self._properties = {}
        self.query = query
        self._configuration = _SyncQueryConfiguration()
        self.udf_resources = udf_resources
        self.query_parameters = query_parameters
        self._job = None

    @classmethod
    def from_query_job(cls, job):
        """Factory: construct from an existing job.

        :type job: :class:`~google.cloud.bigquery.job.QueryJob`
        :param job: existing job

        :rtype: :class:`QueryResults`
        :returns: the instance, bound to the job
        """
        instance = cls(job.query, job._client, job.udf_resources)
        instance._job = job
        job_ref = instance._properties.setdefault('jobReference', {})
        job_ref['jobId'] = job.name
        if job.default_dataset is not None:
            instance.default_dataset = job.default_dataset
        if job.use_query_cache is not None:
            instance.use_query_cache = job.use_query_cache
        if job.use_legacy_sql is not None:
            instance.use_legacy_sql = job.use_legacy_sql
        return instance

    @property
    def project(self):
        """Project bound to the job.

        :rtype: str
        :returns: the project (derived from the client).
        """
        return self._client.project

    def _require_client(self, client):
        """Check client or verify over-ride.

        :type client: :class:`~google.cloud.bigquery.client.Client` or
                      ``NoneType``
        :param client: the client to use.  If not passed, falls back to the
                       ``client`` stored on the current dataset.

        :rtype: :class:`google.cloud.bigquery.client.Client`
        :returns: The client passed in or the currently bound client.
        """
        if client is None:
            client = self._client
        return client

    @property
    def cache_hit(self):
        """Query results served from cache.

        See:
        https://cloud.google.com/bigquery/docs/reference/rest/v2/jobs/query#cacheHit

        :rtype: bool or ``NoneType``
        :returns: True if the query results were served from cache (None
                  until set by the server).
        """
        return self._properties.get('cacheHit')

    @property
    def complete(self):
        """Server completed query.

        See:
        https://cloud.google.com/bigquery/docs/reference/rest/v2/jobs/query#jobComplete

        :rtype: bool or ``NoneType``
        :returns: True if the query completed on the server (None
                  until set by the server).
        """
        return self._properties.get('jobComplete')

    @property
    def errors(self):
        """Errors generated by the query.

        See:
        https://cloud.google.com/bigquery/docs/reference/rest/v2/jobs/query#errors

        :rtype: list of mapping, or ``NoneType``
        :returns: Mappings describing errors generated on the server (None
                  until set by the server).
        """
        return self._properties.get('errors')

    @property
    def name(self):
        """Job name, generated by the back-end.

        See:
        https://cloud.google.com/bigquery/docs/reference/rest/v2/jobs/query#jobReference

        :rtype: list of mapping, or ``NoneType``
        :returns: Mappings describing errors generated on the server (None
                  until set by the server).
        """
        return self._properties.get('jobReference', {}).get('jobId')

    @property
    def job(self):
        """Job instance used to run the query.

        :rtype: :class:`google.cloud.bigquery.job.QueryJob`, or ``NoneType``
        :returns: Job instance used to run the query (None until
                  ``jobReference`` property is set by the server).
        """
        if self._job is None:
            job_ref = self._properties.get('jobReference')
            if job_ref is not None:
                self._job = QueryJob(job_ref['jobId'], self.query,
                                     self._client)
        return self._job

    @property
    def page_token(self):
        """Token for fetching next bach of results.

        See:
        https://cloud.google.com/bigquery/docs/reference/rest/v2/jobs/query#pageToken

        :rtype: str, or ``NoneType``
        :returns: Token generated on the server (None until set by the server).
        """
        return self._properties.get('pageToken')

    @property
    def total_rows(self):
        """Total number of rows returned by the query.

        See:
        https://cloud.google.com/bigquery/docs/reference/rest/v2/jobs/query#totalRows

        :rtype: int, or ``NoneType``
        :returns: Count generated on the server (None until set by the server).
        """
        total_rows = self._properties.get('totalRows')
        if total_rows is not None:
            return int(total_rows)

    @property
    def total_bytes_processed(self):
        """Total number of bytes processed by the query.

        See:
        https://cloud.google.com/bigquery/docs/reference/rest/v2/jobs/query#totalBytesProcessed

        :rtype: int, or ``NoneType``
        :returns: Count generated on the server (None until set by the server).
        """
        total_bytes_processed = self._properties.get('totalBytesProcessed')
        if total_bytes_processed is not None:
            return int(total_bytes_processed)

    @property
    def rows(self):
        """Query results.

        See:
        https://cloud.google.com/bigquery/docs/reference/rest/v2/jobs/query#rows

        :rtype: list of tuples of row values, or ``NoneType``
        :returns: fields describing the schema (None until set by the server).
        """
        return _rows_from_json(self._properties.get('rows', ()), self.schema)

    @property
    def schema(self):
        """Schema for query results.

        See:
        https://cloud.google.com/bigquery/docs/reference/rest/v2/jobs/query#schema

        :rtype: list of :class:`SchemaField`, or ``NoneType``
        :returns: fields describing the schema (None until set by the server).
        """
        return _parse_schema_resource(self._properties.get('schema', {}))

    default_dataset = _TypedProperty('default_dataset', Dataset)
    """See:
    https://cloud.google.com/bigquery/docs/reference/rest/v2/jobs/query#defaultDataset
    """

    dry_run = _TypedProperty('dry_run', bool)
    """See:
    https://cloud.google.com/bigquery/docs/reference/rest/v2/jobs/query#dryRun
    """

    max_results = _TypedProperty('max_results', six.integer_types)
    """See:
    https://cloud.google.com/bigquery/docs/reference/rest/v2/jobs/query#maxResults
    """

    preserve_nulls = _TypedProperty('preserve_nulls', bool)
    """See:
    https://cloud.google.com/bigquery/docs/reference/rest/v2/jobs/query#preserveNulls
    """

    query_parameters = QueryParametersProperty()

    timeout_ms = _TypedProperty('timeout_ms', six.integer_types)
    """See:
    https://cloud.google.com/bigquery/docs/reference/rest/v2/jobs/query#timeoutMs
    """

    udf_resources = UDFResourcesProperty()

    use_query_cache = _TypedProperty('use_query_cache', bool)
    """See:
    https://cloud.google.com/bigquery/docs/reference/rest/v2/jobs/query#useQueryCache
    """

    use_legacy_sql = _TypedProperty('use_legacy_sql', bool)
    """See:
    https://cloud.google.com/bigquery/docs/\
    reference/v2/jobs/query#useLegacySql
    """

    def _set_properties(self, api_response):
        """Update properties from resource in body of ``api_response``

        :type api_response: httplib2.Response
        :param api_response: response returned from an API call
        """
        self._properties.clear()
        self._properties.update(api_response)

    def _build_resource(self):
        """Generate a resource for :meth:`begin`."""
        resource = {'query': self.query}

        if self.default_dataset is not None:
            resource['defaultDataset'] = {
                'projectId': self.project,
                'datasetId': self.default_dataset.name,
            }

        if self.max_results is not None:
            resource['maxResults'] = self.max_results

        if self.preserve_nulls is not None:
            resource['preserveNulls'] = self.preserve_nulls

        if self.timeout_ms is not None:
            resource['timeoutMs'] = self.timeout_ms

        if self.use_query_cache is not None:
            resource['useQueryCache'] = self.use_query_cache

        if self.use_legacy_sql is not None:
            resource['useLegacySql'] = self.use_legacy_sql

        if self.dry_run is not None:
            resource['dryRun'] = self.dry_run

        if len(self._udf_resources) > 0:
            resource[self._UDF_KEY] = [
                {udf_resource.udf_type: udf_resource.value}
                for udf_resource in self._udf_resources
            ]
        if len(self._query_parameters) > 0:
            resource[self._QUERY_PARAMETERS_KEY] = [
                query_parameter.to_api_repr()
                for query_parameter in self._query_parameters
            ]
            if self._query_parameters[0].name is None:
                resource['parameterMode'] = 'POSITIONAL'
            else:
                resource['parameterMode'] = 'NAMED'

        return resource

    def run(self, client=None):
        """API call:  run the query via a POST request

        See:
        https://cloud.google.com/bigquery/docs/reference/rest/v2/jobs/query

        :type client: :class:`~google.cloud.bigquery.client.Client` or
                      ``NoneType``
        :param client: the client to use.  If not passed, falls back to the
                       ``client`` stored on the current dataset.
        """
        if self.job is not None:
            raise ValueError("Query job is already running.")

        client = self._require_client(client)
        path = '/projects/%s/queries' % (self.project,)
        api_response = client._connection.api_request(
            method='POST', path=path, data=self._build_resource())
        self._set_properties(api_response)

    def fetch_data(self, max_results=None, page_token=None, start_index=None,
                   timeout_ms=None, client=None):
        """API call:  fetch a page of query result data via a GET request

        See:
        https://cloud.google.com/bigquery/docs/reference/rest/v2/jobs/getQueryResults

        :type max_results: int
        :param max_results: (Optional) maximum number of rows to return.

        :type page_token: str
        :param page_token:
            (Optional) token representing a cursor into the table's rows.

        :type start_index: int
        :param start_index: (Optional) zero-based index of starting row

        :type timeout_ms: int
        :param timeout_ms:
            (Optional) timeout, in milliseconds, to wait for http request
            to complete.

        :type client: :class:`~google.cloud.bigquery.client.Client` or
                      ``NoneType``
        :param client: the client to use.  If not passed, falls back to the
                       ``client`` stored on the current dataset.

        :rtype: tuple
        :returns: ``(row_data, total_rows, page_token)``, where ``row_data``
                  is a list of tuples, one per result row, containing only
                  the values;  ``total_rows`` is a count of the total number
                  of rows in the table;  and ``page_token`` is an opaque
                  string which can be used to fetch the next batch of rows
                  (``None`` if no further batches can be fetched).
        :raises: ValueError if the query has not yet been executed.
        """
        if self.name is None:
            raise ValueError("Query not yet executed:  call 'run()'")

        client = self._require_client(client)
        params = {}

        if max_results is not None:
            params['maxResults'] = max_results

        if page_token is not None:
            params['pageToken'] = page_token

        if start_index is not None:
            params['startIndex'] = start_index

        if timeout_ms is not None:
            params['timeoutMs'] = timeout_ms

        path = '/projects/%s/queries/%s' % (self.project, self.name)
        response = client._connection.api_request(method='GET',
                                                  path=path,
                                                  query_params=params)
        self._set_properties(response)

        total_rows = response.get('totalRows')
        if total_rows is not None:
            total_rows = int(total_rows)
        page_token = response.get('pageToken')
        rows_data = _rows_from_json(response.get('rows', ()), self.schema)

        return rows_data, total_rows, page_token
<commit_msg>Include full API reference<commit_after># Copyright 2015 Google Inc.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""Define API Queries."""

import six

from google.cloud.bigquery._helpers import _TypedProperty
from google.cloud.bigquery._helpers import _rows_from_json
from google.cloud.bigquery.dataset import Dataset
from google.cloud.bigquery.job import QueryJob
from google.cloud.bigquery.table import _parse_schema_resource
from google.cloud.bigquery._helpers import QueryParametersProperty
from google.cloud.bigquery._helpers import UDFResourcesProperty

class _SyncQueryConfiguration(object):
    """User-settable configuration options for synchronous query jobs.

    Values which are ``None`` -> server defaults.
    """
    _default_dataset = None
    _dry_run = None
    _max_results = None
    _timeout_ms = None
    _preserve_nulls = None
    _use_query_cache = None
    _use_legacy_sql = None

class QueryResults(object):
    """Synchronous job: query tables.

    :type query: str
    :param query: SQL query string

    :type client: :class:`google.cloud.bigquery.client.Client`
    :param client: A client which holds credentials and project configuration
                   for the dataset (which requires a project).

    :type udf_resources: tuple
    :param udf_resources: An iterable of
                        :class:`google.cloud.bigquery.job.UDFResource`
                        (empty by default)

    :type query_parameters: tuple
    :param query_parameters:
        An iterable of
        :class:`google.cloud.bigquery._helpers.AbstractQueryParameter`
        (empty by default)
    """

    _UDF_KEY = 'userDefinedFunctionResources'
    _QUERY_PARAMETERS_KEY = 'queryParameters'

    def __init__(self, query, client, udf_resources=(), query_parameters=()):
        self._client = client
        self._properties = {}
        self.query = query
        self._configuration = _SyncQueryConfiguration()
        self.udf_resources = udf_resources
        self.query_parameters = query_parameters
        self._job = None

    @classmethod
    def from_query_job(cls, job):
        """Factory: construct from an existing job.

        :type job: :class:`~google.cloud.bigquery.job.QueryJob`
        :param job: existing job

        :rtype: :class:`QueryResults`
        :returns: the instance, bound to the job
        """
        instance = cls(job.query, job._client, job.udf_resources)
        instance._job = job
        job_ref = instance._properties.setdefault('jobReference', {})
        job_ref['jobId'] = job.name
        if job.default_dataset is not None:
            instance.default_dataset = job.default_dataset
        if job.use_query_cache is not None:
            instance.use_query_cache = job.use_query_cache
        if job.use_legacy_sql is not None:
            instance.use_legacy_sql = job.use_legacy_sql
        return instance

    @property
    def project(self):
        """Project bound to the job.

        :rtype: str
        :returns: the project (derived from the client).
        """
        return self._client.project

    def _require_client(self, client):
        """Check client or verify over-ride.

        :type client: :class:`~google.cloud.bigquery.client.Client` or
                      ``NoneType``
        :param client: the client to use.  If not passed, falls back to the
                       ``client`` stored on the current dataset.

        :rtype: :class:`google.cloud.bigquery.client.Client`
        :returns: The client passed in or the currently bound client.
        """
        if client is None:
            client = self._client
        return client

    @property
    def cache_hit(self):
        """Query results served from cache.

        See:
        https://cloud.google.com/bigquery/docs/reference/rest/v2/jobs/query#cacheHit

        :rtype: bool or ``NoneType``
        :returns: True if the query results were served from cache (None
                  until set by the server).
        """
        return self._properties.get('cacheHit')

    @property
    def complete(self):
        """Server completed query.

        See:
        https://cloud.google.com/bigquery/docs/reference/rest/v2/jobs/query#jobComplete

        :rtype: bool or ``NoneType``
        :returns: True if the query completed on the server (None
                  until set by the server).
        """
        return self._properties.get('jobComplete')

    @property
    def errors(self):
        """Errors generated by the query.

        See:
        https://cloud.google.com/bigquery/docs/reference/rest/v2/jobs/query#errors

        :rtype: list of mapping, or ``NoneType``
        :returns: Mappings describing errors generated on the server (None
                  until set by the server).
        """
        return self._properties.get('errors')

    @property
    def name(self):
        """Job name, generated by the back-end.

        See:
        https://cloud.google.com/bigquery/docs/reference/rest/v2/jobs/query#jobReference

        :rtype: list of mapping, or ``NoneType``
        :returns: Mappings describing errors generated on the server (None
                  until set by the server).
        """
        return self._properties.get('jobReference', {}).get('jobId')

    @property
    def job(self):
        """Job instance used to run the query.

        :rtype: :class:`google.cloud.bigquery.job.QueryJob`, or ``NoneType``
        :returns: Job instance used to run the query (None until
                  ``jobReference`` property is set by the server).
        """
        if self._job is None:
            job_ref = self._properties.get('jobReference')
            if job_ref is not None:
                self._job = QueryJob(job_ref['jobId'], self.query,
                                     self._client)
        return self._job

    @property
    def page_token(self):
        """Token for fetching next bach of results.

        See:
        https://cloud.google.com/bigquery/docs/reference/rest/v2/jobs/query#pageToken

        :rtype: str, or ``NoneType``
        :returns: Token generated on the server (None until set by the server).
        """
        return self._properties.get('pageToken')

    @property
    def total_rows(self):
        """Total number of rows returned by the query.

        See:
        https://cloud.google.com/bigquery/docs/reference/rest/v2/jobs/query#totalRows

        :rtype: int, or ``NoneType``
        :returns: Count generated on the server (None until set by the server).
        """
        total_rows = self._properties.get('totalRows')
        if total_rows is not None:
            return int(total_rows)

    @property
    def total_bytes_processed(self):
        """Total number of bytes processed by the query.

        See:
        https://cloud.google.com/bigquery/docs/reference/rest/v2/jobs/query#totalBytesProcessed

        :rtype: int, or ``NoneType``
        :returns: Count generated on the server (None until set by the server).
        """
        total_bytes_processed = self._properties.get('totalBytesProcessed')
        if total_bytes_processed is not None:
            return int(total_bytes_processed)

    @property
    def rows(self):
        """Query results.

        See:
        https://cloud.google.com/bigquery/docs/reference/rest/v2/jobs/query#rows

        :rtype: list of tuples of row values, or ``NoneType``
        :returns: fields describing the schema (None until set by the server).
        """
        return _rows_from_json(self._properties.get('rows', ()), self.schema)

    @property
    def schema(self):
        """Schema for query results.

        See:
        https://cloud.google.com/bigquery/docs/reference/rest/v2/jobs/query#schema

        :rtype: list of :class:`SchemaField`, or ``NoneType``
        :returns: fields describing the schema (None until set by the server).
        """
        return _parse_schema_resource(self._properties.get('schema', {}))

    default_dataset = _TypedProperty('default_dataset', Dataset)
    """See:
    https://cloud.google.com/bigquery/docs/reference/rest/v2/jobs/query#defaultDataset
    """

    dry_run = _TypedProperty('dry_run', bool)
    """See:
    https://cloud.google.com/bigquery/docs/reference/rest/v2/jobs/query#dryRun
    """

    max_results = _TypedProperty('max_results', six.integer_types)
    """See:
    https://cloud.google.com/bigquery/docs/reference/rest/v2/jobs/query#maxResults
    """

    preserve_nulls = _TypedProperty('preserve_nulls', bool)
    """See:
    https://cloud.google.com/bigquery/docs/reference/rest/v2/jobs/query#preserveNulls
    """

    query_parameters = QueryParametersProperty()

    timeout_ms = _TypedProperty('timeout_ms', six.integer_types)
    """See:
    https://cloud.google.com/bigquery/docs/reference/rest/v2/jobs/query#timeoutMs
    """

    udf_resources = UDFResourcesProperty()

    use_query_cache = _TypedProperty('use_query_cache', bool)
    """See:
    https://cloud.google.com/bigquery/docs/reference/rest/v2/jobs/query#useQueryCache
    """

    use_legacy_sql = _TypedProperty('use_legacy_sql', bool)
    """See:
    https://cloud.google.com/bigquery/docs/\
    reference/v2/jobs/query#useLegacySql
    """

    def _set_properties(self, api_response):
        """Update properties from resource in body of ``api_response``

        :type api_response: httplib2.Response
        :param api_response: response returned from an API call
        """
        self._properties.clear()
        self._properties.update(api_response)

    def _build_resource(self):
        """Generate a resource for :meth:`begin`."""
        resource = {'query': self.query}

        if self.default_dataset is not None:
            resource['defaultDataset'] = {
                'projectId': self.project,
                'datasetId': self.default_dataset.name,
            }

        if self.max_results is not None:
            resource['maxResults'] = self.max_results

        if self.preserve_nulls is not None:
            resource['preserveNulls'] = self.preserve_nulls

        if self.timeout_ms is not None:
            resource['timeoutMs'] = self.timeout_ms

        if self.use_query_cache is not None:
            resource['useQueryCache'] = self.use_query_cache

        if self.use_legacy_sql is not None:
            resource['useLegacySql'] = self.use_legacy_sql

        if self.dry_run is not None:
            resource['dryRun'] = self.dry_run

        if len(self._udf_resources) > 0:
            resource[self._UDF_KEY] = [
                {udf_resource.udf_type: udf_resource.value}
                for udf_resource in self._udf_resources
            ]
        if len(self._query_parameters) > 0:
            resource[self._QUERY_PARAMETERS_KEY] = [
                query_parameter.to_api_repr()
                for query_parameter in self._query_parameters
            ]
            if self._query_parameters[0].name is None:
                resource['parameterMode'] = 'POSITIONAL'
            else:
                resource['parameterMode'] = 'NAMED'

        return resource

    def run(self, client=None):
        """API call:  run the query via a POST request

        See:
        https://cloud.google.com/bigquery/docs/reference/rest/v2/jobs/query

        :type client: :class:`~google.cloud.bigquery.client.Client` or
                      ``NoneType``
        :param client: the client to use.  If not passed, falls back to the
                       ``client`` stored on the current dataset.
        """
        if self.job is not None:
            raise ValueError("Query job is already running.")

        client = self._require_client(client)
        path = '/projects/%s/queries' % (self.project,)
        api_response = client._connection.api_request(
            method='POST', path=path, data=self._build_resource())
        self._set_properties(api_response)

    def fetch_data(self, max_results=None, page_token=None, start_index=None,
                   timeout_ms=None, client=None):
        """API call:  fetch a page of query result data via a GET request

        See:
        https://cloud.google.com/bigquery/docs/reference/rest/v2/jobs/getQueryResults

        :type max_results: int
        :param max_results: (Optional) maximum number of rows to return.

        :type page_token: str
        :param page_token:
            (Optional) token representing a cursor into the table's rows.

        :type start_index: int
        :param start_index: (Optional) zero-based index of starting row

        :type timeout_ms: int
        :param timeout_ms:
            (Optional) How long to wait for the query to complete, in
            milliseconds, before the request times out and returns. Note that
            this is only a timeout for the request, not the query. If the query
            takes longer to run than the timeout value, the call returns without
            any results and with the 'jobComplete' flag set to false. You can
            call GetQueryResults() to wait for the query to complete and read
            the results. The default value is 10000 milliseconds (10 seconds).

        :type client: :class:`~google.cloud.bigquery.client.Client` or
                      ``NoneType``
        :param client: the client to use.  If not passed, falls back to the
                       ``client`` stored on the current dataset.

        :rtype: tuple
        :returns: ``(row_data, total_rows, page_token)``, where ``row_data``
                  is a list of tuples, one per result row, containing only
                  the values;  ``total_rows`` is a count of the total number
                  of rows in the table;  and ``page_token`` is an opaque
                  string which can be used to fetch the next batch of rows
                  (``None`` if no further batches can be fetched).
        :raises: ValueError if the query has not yet been executed.
        """
        if self.name is None:
            raise ValueError("Query not yet executed:  call 'run()'")

        client = self._require_client(client)
        params = {}

        if max_results is not None:
            params['maxResults'] = max_results

        if page_token is not None:
            params['pageToken'] = page_token

        if start_index is not None:
            params['startIndex'] = start_index

        if timeout_ms is not None:
            params['timeoutMs'] = timeout_ms

        path = '/projects/%s/queries/%s' % (self.project, self.name)
        response = client._connection.api_request(method='GET',
                                                  path=path,
                                                  query_params=params)
        self._set_properties(response)

        total_rows = response.get('totalRows')
        if total_rows is not None:
            total_rows = int(total_rows)
        page_token = response.get('pageToken')
        rows_data = _rows_from_json(response.get('rows', ()), self.schema)

        return rows_data, total_rows, page_token
<|endoftext|>