<commit_before>/*
 * Copyright 2019 Spotify AB.
 *
 * 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.
 */

package com.spotify.scio.bigquery

import java.util.UUID

import com.google.api.client.googleapis.json.GoogleJsonResponseException
import com.google.api.services.bigquery.model.{TableFieldSchema, TableSchema}
import com.google.cloud.storage.Storage.BlobListOption
import com.google.cloud.storage.{Blob, StorageOptions}
import com.spotify.scio.bigquery.client.BigQuery
import org.apache.beam.sdk.io.gcp.bigquery.BigQueryHelpers
import org.scalatest.matchers.should.Matchers
import org.scalatest.flatspec.AnyFlatSpec

import scala.collection.JavaConverters._
import scala.util.Success

// scio-test/it:runMain PopulateTestData to re-populate data for integration tests
class BigQueryClientIT extends AnyFlatSpec with Matchers {
  private[this] val bq = BigQuery.defaultInstance()

  val legacyQuery =
    "SELECT word, word_count FROM [bigquery-public-data:samples.shakespeare] LIMIT 10"
  val sqlQuery =
    "SELECT word, word_count FROM `bigquery-public-data.samples.shakespeare` LIMIT 10"

  "QueryService.run" should "run DML queries" in {
    val schema =
      BigQueryUtil.parseSchema("""
      |{
      |  "fields": [
      |    {"mode": "NULLABLE", "name": "word", "type": "STRING"},
      |    {"mode": "NULLABLE", "name": "word_count", "type": "INTEGER"},
      |    {"mode": "NULLABLE", "name": "corpus", "type": "STRING"},
      |    {"mode": "NULLABLE", "name": "corpus_date", "type": "INTEGER"}
      |  ]
      |}
    """.stripMargin)
    val sources = List("gs://data-integration-test-eu/shakespeare-sample-10.json")
    val table = bq.tables.createTemporary(location = "EU")
    val tableRef = bq.load.json(sources, table.asTableSpec, schema = Some(schema))
    tableRef.map { ref =>
      val insertQuery = s"insert into `${ref.asTableSpec}` values (1603, 'alien', 9000, 'alien')"
      bq.query.run(
        insertQuery,
        createDisposition = null,
        writeDisposition = null
      )

      val deleteQuery = s"delete from `${ref.asTableSpec}` where word = 'alien'"
      bq.query.run(
        deleteQuery,
        createDisposition = null,
        writeDisposition = null
      )
    }
  }

  "QueryService.extractLocation" should "work with legacy syntax" in {
    val query = "SELECT word FROM [data-integration-test:samples_%s.shakespeare]"
    bq.query.extractLocation(query.format("eu")) shouldBe Some("EU")
    bq.query.extractLocation(query.format("us")) shouldBe Some("US")
  }

  it should "work with SQL syntax" in {
    val query = "SELECT word FROM `data-integration-test.samples_%s.shakespeare`"
    bq.query.extractLocation(query.format("eu")) shouldBe Some("EU")
    bq.query.extractLocation(query.format("us")) shouldBe Some("US")
  }

  it should "support missing source tables" in {
    bq.query.extractLocation("SELECT 6") shouldBe Some("EU")
  }

  "QueryService.extractTables" should "work with legacy syntax" in {
    val tableSpec = BigQueryHelpers.parseTableSpec("bigquery-public-data:samples.shakespeare")
    bq.query.extractTables(legacyQuery) shouldBe Set(tableSpec)
  }

  it should "work with SQL syntax" in {
    val tableSpec = BigQueryHelpers.parseTableSpec("bigquery-public-data:samples.shakespeare")
    bq.query.extractTables(sqlQuery) shouldBe Set(tableSpec)
  }

  "QueryService.getSchema" should "work with legacy syntax" in {
    val expected = new TableSchema().setFields(
      List(
        new TableFieldSchema().setName("word").setType("STRING").setMode("REQUIRED"),<commit_msg>Fix: temp table location changed back to US (#2905)<commit_after>/*
 * Copyright 2019 Spotify AB.
 *
 * 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.
 */

package com.spotify.scio.bigquery

import java.util.UUID

import com.google.api.client.googleapis.json.GoogleJsonResponseException
import com.google.api.services.bigquery.model.{TableFieldSchema, TableSchema}
import com.google.cloud.storage.Storage.BlobListOption
import com.google.cloud.storage.{Blob, StorageOptions}
import com.spotify.scio.bigquery.client.BigQuery
import org.apache.beam.sdk.io.gcp.bigquery.BigQueryHelpers
import org.scalatest.matchers.should.Matchers
import org.scalatest.flatspec.AnyFlatSpec

import scala.collection.JavaConverters._
import scala.util.Success

// scio-test/it:runMain PopulateTestData to re-populate data for integration tests
class BigQueryClientIT extends AnyFlatSpec with Matchers {
  private[this] val bq = BigQuery.defaultInstance()

  val legacyQuery =
    "SELECT word, word_count FROM [bigquery-public-data:samples.shakespeare] LIMIT 10"
  val sqlQuery =
    "SELECT word, word_count FROM `bigquery-public-data.samples.shakespeare` LIMIT 10"

  "QueryService.run" should "run DML queries" in {
    val schema =
      BigQueryUtil.parseSchema("""
      |{
      |  "fields": [
      |    {"mode": "NULLABLE", "name": "word", "type": "STRING"},
      |    {"mode": "NULLABLE", "name": "word_count", "type": "INTEGER"},
      |    {"mode": "NULLABLE", "name": "corpus", "type": "STRING"},
      |    {"mode": "NULLABLE", "name": "corpus_date", "type": "INTEGER"}
      |  ]
      |}
    """.stripMargin)
    val sources = List("gs://data-integration-test-eu/shakespeare-sample-10.json")
    val table = bq.tables.createTemporary(location = "EU")
    val tableRef = bq.load.json(sources, table.asTableSpec, schema = Some(schema))
    tableRef.map { ref =>
      val insertQuery = s"insert into `${ref.asTableSpec}` values (1603, 'alien', 9000, 'alien')"
      bq.query.run(
        insertQuery,
        createDisposition = null,
        writeDisposition = null
      )

      val deleteQuery = s"delete from `${ref.asTableSpec}` where word = 'alien'"
      bq.query.run(
        deleteQuery,
        createDisposition = null,
        writeDisposition = null
      )
    }
  }

  "QueryService.extractLocation" should "work with legacy syntax" in {
    val query = "SELECT word FROM [data-integration-test:samples_%s.shakespeare]"
    bq.query.extractLocation(query.format("eu")) shouldBe Some("EU")
    bq.query.extractLocation(query.format("us")) shouldBe Some("US")
  }

  it should "work with SQL syntax" in {
    val query = "SELECT word FROM `data-integration-test.samples_%s.shakespeare`"
    bq.query.extractLocation(query.format("eu")) shouldBe Some("EU")
    bq.query.extractLocation(query.format("us")) shouldBe Some("US")
  }

  it should "support missing source tables" in {
    bq.query.extractLocation("SELECT 6") shouldBe Some("US")
  }

  "QueryService.extractTables" should "work with legacy syntax" in {
    val tableSpec = BigQueryHelpers.parseTableSpec("bigquery-public-data:samples.shakespeare")
    bq.query.extractTables(legacyQuery) shouldBe Set(tableSpec)
  }

  it should "work with SQL syntax" in {
    val tableSpec = BigQueryHelpers.parseTableSpec("bigquery-public-data:samples.shakespeare")
    bq.query.extractTables(sqlQuery) shouldBe Set(tableSpec)
  }

  "QueryService.getSchema" should "work with legacy syntax" in {
    val expected = new TableSchema().setFields(
      List(
        new TableFieldSchema().setName("word").setType("STRING").setMode("REQUIRED"),
{'fraction_non_alphanumeric': 0.12331783704428677, 'fraction_numerical': 0.006116956202593589, 'mean_word_length': 3.857397504456328, 'pattern_counts': {'":': 26, '<': 3, '<?xml version=': 0, '>': 5, 'https://': 0, 'lorem ipsum': 0, 'www.': 2, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '3405328', 'n_tokens_mistral': 2523, 'n_tokens_neox': 2471, 'n_words': 635}