prompt: We are currently solving the following issue within our repository. Here is the issue text:
  --- BEGIN ISSUE ---
  Bug: Dataset not found in US
  I'm getting an error while running the BigQuery kfp component to export some data from BQ to GCS. It says: `google.api_core.exceptions.NotFound: 404 Not found: xxx was not found in location US` which is correct, because the dataset is in EU, but I couldn't find any way to change the default behaviour of looking for datasets in the US.
  
  On our regular codebase we specify the location on the BigQuery client initialisation, but the code here just uses the default: https://github.com/danicat/pipelines/blob/master/component_sdk/python/kfp_component/google/bigquery/_query.py
  
  It is a one line change in the code, but I don't know how to rebuild the ml-pipeline container. The developer guide doesn't mention how to do it (but it does explain how to build containers for other components of the solution).
  
  Please help! :)
  
  --- END ISSUE ---
  
  Below are some code segments, each from a relevant file. One or more of these files may contain bugs.
  --- BEGIN FILES ---
  Path: `component_sdk/python/kfp_component/google/bigquery/_query.py`
  Content:
  ```
  1 # Copyright 2018 Google LLC
  2 #
  3 # Licensed under the Apache License, Version 2.0 (the "License");
  4 # you may not use this file except in compliance with the License.
  5 # You may obtain a copy of the License at
  6 #
  7 #      http://www.apache.org/licenses/LICENSE-2.0
  8 #
  9 # Unless required by applicable law or agreed to in writing, software
  10 # distributed under the License is distributed on an "AS IS" BASIS,
  11 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  12 # See the License for the specific language governing permissions and
  13 # limitations under the License.
  14 
  15 import json
  16 import logging
  17 
  18 from google.cloud import bigquery
  19 from google.api_core import exceptions
  20 
  21 from kfp_component.core import KfpExecutionContext, display
  22 from .. import common as gcp_common
  23 
  24 # TODO(hongyes): make this path configurable as a environment variable
  25 KFP_OUTPUT_PATH = '/tmp/kfp/output/'
  26 
  27 def query(query, project_id, dataset_id=None, table_id=None, 
  28     output_gcs_path=None, dataset_location='US', job_config=None):
  29     """Submit a query to Bigquery service and dump outputs to Bigquery table or 
  30     a GCS blob.
  31     
  32     Args:
  33         query (str): The query used by Bigquery service to fetch the results.
  34         project_id (str): The project to execute the query job.
  35         dataset_id (str): The ID of the persistent dataset to keep the results
  36             of the query. If the dataset does not exist, the operation will 
  37             create a new one.
  38         table_id (str): The ID of the table to keep the results of the query. If
  39             absent, the operation will generate a random id for the table.
  40         output_gcs_path (str): The GCS blob path to dump the query results to.
  41         dataset_location (str): The location to create the dataset. Defaults to `US`.
  42         job_config (dict): The full config spec for the query job.
  43     Returns:
  44         The API representation of the completed query job.
  45     """
  46     client = bigquery.Client(project=project_id)
  47     if not job_config:
  48         job_config = bigquery.QueryJobConfig()
  49         job_config.create_disposition = bigquery.job.CreateDisposition.CREATE_IF_NEEDED
  50         job_config.write_disposition = bigquery.job.WriteDisposition.WRITE_TRUNCATE
  51     job_id = None
  52     def cancel():
  53         if job_id:
  54             client.cancel_job(job_id)
  55     with KfpExecutionContext(on_cancel=cancel) as ctx:
  56         job_id = 'query_' + ctx.context_id()
  57         query_job = _get_job(client, job_id)
  58         table_ref = None
  59         if not query_job:
  60             dataset_ref = _prepare_dataset_ref(client, dataset_id, output_gcs_path, 
  61                 dataset_location)
  62             if dataset_ref:
  63                 if not table_id:
  64                     table_id = job_id
  65                 table_ref = dataset_ref.table(table_id)
  66                 job_config.destination = table_ref
  67             query_job = client.query(query, job_config, job_id=job_id)
  68         _display_job_link(project_id, job_id)
  69         query_job.result() # Wait for query to finish
  70         if output_gcs_path:
  71             job_id = 'extract_' + ctx.context_id()
  72             extract_job = _get_job(client, job_id)
  73             logging.info('Extracting data from table {} to {}.'.format(str(table_ref), output_gcs_path))
  74             if not extract_job:
  75                 extract_job = client.extract_table(table_ref, output_gcs_path)
  76             extract_job.result()  # Wait for export to finish
  77         _dump_outputs(query_job, output_gcs_path, table_ref)
  78         return query_job.to_api_repr()
  79 
  80 def _get_job(client, job_id):
  81     try:
  82         return client.get_job(job_id)
  83     except exceptions.NotFound:
  84         return None
  85 
  86 def _prepare_dataset_ref(client, dataset_id, output_gcs_path, dataset_location):
  87     if not output_gcs_path and not dataset_id:
  88         return None
  89     
  90     if not dataset_id:
  91         dataset_id = 'kfp_tmp_dataset'
  92     dataset_ref = client.dataset(dataset_id)
  93     dataset = _get_dataset(client, dataset_ref)
  94     if not dataset:
  95         logging.info('Creating dataset {}'.format(dataset_id))
  96         dataset = _create_dataset(client, dataset_ref, dataset_location)
  97     return dataset_ref
  98 
  99 def _get_dataset(client, dataset_ref):
  100     try:
  101         return client.get_dataset(dataset_ref)
  102     except exceptions.NotFound:
  103         return None
  104 
  105 def _create_dataset(client, dataset_ref, location):
  106     dataset = bigquery.Dataset(dataset_ref)
  107     dataset.location = location
  108     return client.create_dataset(dataset)
  109 
  110 def _display_job_link(project_id, job_id):
  111     display.display(display.Link(
  112         href= 'https://console.cloud.google.com/bigquery?project={}'
  113             '&j={}&page=queryresults'.format(project_id, job_id),
  114         text='Query Details'
  115     ))
  116 
  117 def _dump_outputs(job, output_path, table_ref):
  118     gcp_common.dump_file(KFP_OUTPUT_PATH + 'bigquery/query-job.json', 
  119         json.dumps(job.to_api_repr()))
  120     if not output_path:
  121         output_path = ''
  122     gcp_common.dump_file(KFP_OUTPUT_PATH + 'bigquery/query-output-path.txt', 
  123         output_path)
  124     (dataset_id, table_id) = (table_ref.dataset_id, table_ref.table_id) if table_ref else ('', '')
  125     gcp_common.dump_file(KFP_OUTPUT_PATH + 'bigquery/query-dataset-id.txt', 
  126         dataset_id)
  127     gcp_common.dump_file(KFP_OUTPUT_PATH + 'bigquery/query-table-id.txt', 
  128         table_id)
  129
  ```
  --- END FILES ---
  
  Please first localize the bug based on the issue statement, and then generate a patch according to the `git diff` format fenced by three backticks.
  
  Here is an example:
  
  ```diff
  diff --git a/examples/server_async.py b/examples/server_async.py
  --- a/examples/server_async.py
  +++ b/examples/server_async.py
  @@ -313,4 +313,4 @@
   
   
   if __name__ == "__main__":
  -    asyncio.run(run_async_server("."), debug=True)
  +    asyncio.run(run_async_server(), debug=True)
  diff --git a/examples/server_sync.py b/examples/server_sync.py
  --- a/examples/server_sync.py
  +++ b/examples/server_sync.py
  @@ -313,5 +313,5 @@
   
   
   if __name__ == "__main__":
  -    server = run_sync_server(".")
  +    server = run_sync_server()
       server.shutdown()
  
  ```
  
golden_diff: diff --git a/component_sdk/python/kfp_component/google/bigquery/_query.py b/component_sdk/python/kfp_component/google/bigquery/_query.py
  --- a/component_sdk/python/kfp_component/google/bigquery/_query.py
  +++ b/component_sdk/python/kfp_component/google/bigquery/_query.py
  @@ -43,7 +43,7 @@
       Returns:
           The API representation of the completed query job.
       """
  -    client = bigquery.Client(project=project_id)
  +    client = bigquery.Client(project=project_id, location=dataset_location)
       if not job_config:
           job_config = bigquery.QueryJobConfig()
           job_config.create_disposition = bigquery.job.CreateDisposition.CREATE_IF_NEEDED
  
verification_info: {"golden_diff": "diff --git a/component_sdk/python/kfp_component/google/bigquery/_query.py b/component_sdk/python/kfp_component/google/bigquery/_query.py\n--- a/component_sdk/python/kfp_component/google/bigquery/_query.py\n+++ b/component_sdk/python/kfp_component/google/bigquery/_query.py\n@@ -43,7 +43,7 @@\n     Returns:\n         The API representation of the completed query job.\n     \"\"\"\n-    client = bigquery.Client(project=project_id)\n+    client = bigquery.Client(project=project_id, location=dataset_location)\n     if not job_config:\n         job_config = bigquery.QueryJobConfig()\n         job_config.create_disposition = bigquery.job.CreateDisposition.CREATE_IF_NEEDED\n", "issue": "Bug: Dataset not found in US\nI'm getting an error while running the BigQuery kfp component to export some data from BQ to GCS. It says: `google.api_core.exceptions.NotFound: 404 Not found: xxx was not found in location US` which is correct, because the dataset is in EU, but I couldn't find any way to change the default behaviour of looking for datasets in the US.\r\n\r\nOn our regular codebase we specify the location on the BigQuery client initialisation, but the code here just uses the default: https://github.com/danicat/pipelines/blob/master/component_sdk/python/kfp_component/google/bigquery/_query.py\r\n\r\nIt is a one line change in the code, but I don't know how to rebuild the ml-pipeline container. The developer guide doesn't mention how to do it (but it does explain how to build containers for other components of the solution).\r\n\r\nPlease help! :)\n", "before_files": [{"content": "# Copyright 2018 Google LLC\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#      http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nimport json\nimport logging\n\nfrom google.cloud import bigquery\nfrom google.api_core import exceptions\n\nfrom kfp_component.core import KfpExecutionContext, display\nfrom .. import common as gcp_common\n\n# TODO(hongyes): make this path configurable as a environment variable\nKFP_OUTPUT_PATH = '/tmp/kfp/output/'\n\ndef query(query, project_id, dataset_id=None, table_id=None, \n    output_gcs_path=None, dataset_location='US', job_config=None):\n    \"\"\"Submit a query to Bigquery service and dump outputs to Bigquery table or \n    a GCS blob.\n    \n    Args:\n        query (str): The query used by Bigquery service to fetch the results.\n        project_id (str): The project to execute the query job.\n        dataset_id (str): The ID of the persistent dataset to keep the results\n            of the query. If the dataset does not exist, the operation will \n            create a new one.\n        table_id (str): The ID of the table to keep the results of the query. If\n            absent, the operation will generate a random id for the table.\n        output_gcs_path (str): The GCS blob path to dump the query results to.\n        dataset_location (str): The location to create the dataset. Defaults to `US`.\n        job_config (dict): The full config spec for the query job.\n    Returns:\n        The API representation of the completed query job.\n    \"\"\"\n    client = bigquery.Client(project=project_id)\n    if not job_config:\n        job_config = bigquery.QueryJobConfig()\n        job_config.create_disposition = bigquery.job.CreateDisposition.CREATE_IF_NEEDED\n        job_config.write_disposition = bigquery.job.WriteDisposition.WRITE_TRUNCATE\n    job_id = None\n    def cancel():\n        if job_id:\n            client.cancel_job(job_id)\n    with KfpExecutionContext(on_cancel=cancel) as ctx:\n        job_id = 'query_' + ctx.context_id()\n        query_job = _get_job(client, job_id)\n        table_ref = None\n        if not query_job:\n            dataset_ref = _prepare_dataset_ref(client, dataset_id, output_gcs_path, \n                dataset_location)\n            if dataset_ref:\n                if not table_id:\n                    table_id = job_id\n                table_ref = dataset_ref.table(table_id)\n                job_config.destination = table_ref\n            query_job = client.query(query, job_config, job_id=job_id)\n        _display_job_link(project_id, job_id)\n        query_job.result() # Wait for query to finish\n        if output_gcs_path:\n            job_id = 'extract_' + ctx.context_id()\n            extract_job = _get_job(client, job_id)\n            logging.info('Extracting data from table {} to {}.'.format(str(table_ref), output_gcs_path))\n            if not extract_job:\n                extract_job = client.extract_table(table_ref, output_gcs_path)\n            extract_job.result()  # Wait for export to finish\n        _dump_outputs(query_job, output_gcs_path, table_ref)\n        return query_job.to_api_repr()\n\ndef _get_job(client, job_id):\n    try:\n        return client.get_job(job_id)\n    except exceptions.NotFound:\n        return None\n\ndef _prepare_dataset_ref(client, dataset_id, output_gcs_path, dataset_location):\n    if not output_gcs_path and not dataset_id:\n        return None\n    \n    if not dataset_id:\n        dataset_id = 'kfp_tmp_dataset'\n    dataset_ref = client.dataset(dataset_id)\n    dataset = _get_dataset(client, dataset_ref)\n    if not dataset:\n        logging.info('Creating dataset {}'.format(dataset_id))\n        dataset = _create_dataset(client, dataset_ref, dataset_location)\n    return dataset_ref\n\ndef _get_dataset(client, dataset_ref):\n    try:\n        return client.get_dataset(dataset_ref)\n    except exceptions.NotFound:\n        return None\n\ndef _create_dataset(client, dataset_ref, location):\n    dataset = bigquery.Dataset(dataset_ref)\n    dataset.location = location\n    return client.create_dataset(dataset)\n\ndef _display_job_link(project_id, job_id):\n    display.display(display.Link(\n        href= 'https://console.cloud.google.com/bigquery?project={}'\n            '&j={}&page=queryresults'.format(project_id, job_id),\n        text='Query Details'\n    ))\n\ndef _dump_outputs(job, output_path, table_ref):\n    gcp_common.dump_file(KFP_OUTPUT_PATH + 'bigquery/query-job.json', \n        json.dumps(job.to_api_repr()))\n    if not output_path:\n        output_path = ''\n    gcp_common.dump_file(KFP_OUTPUT_PATH + 'bigquery/query-output-path.txt', \n        output_path)\n    (dataset_id, table_id) = (table_ref.dataset_id, table_ref.table_id) if table_ref else ('', '')\n    gcp_common.dump_file(KFP_OUTPUT_PATH + 'bigquery/query-dataset-id.txt', \n        dataset_id)\n    gcp_common.dump_file(KFP_OUTPUT_PATH + 'bigquery/query-table-id.txt', \n        table_id)\n", "path": "component_sdk/python/kfp_component/google/bigquery/_query.py"}], "after_files": [{"content": "# Copyright 2018 Google LLC\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#      http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nimport json\nimport logging\n\nfrom google.cloud import bigquery\nfrom google.api_core import exceptions\n\nfrom kfp_component.core import KfpExecutionContext, display\nfrom .. import common as gcp_common\n\n# TODO(hongyes): make this path configurable as a environment variable\nKFP_OUTPUT_PATH = '/tmp/kfp/output/'\n\ndef query(query, project_id, dataset_id=None, table_id=None, \n    output_gcs_path=None, dataset_location='US', job_config=None):\n    \"\"\"Submit a query to Bigquery service and dump outputs to Bigquery table or \n    a GCS blob.\n    \n    Args:\n        query (str): The query used by Bigquery service to fetch the results.\n        project_id (str): The project to execute the query job.\n        dataset_id (str): The ID of the persistent dataset to keep the results\n            of the query. If the dataset does not exist, the operation will \n            create a new one.\n        table_id (str): The ID of the table to keep the results of the query. If\n            absent, the operation will generate a random id for the table.\n        output_gcs_path (str): The GCS blob path to dump the query results to.\n        dataset_location (str): The location to create the dataset. Defaults to `US`.\n        job_config (dict): The full config spec for the query job.\n    Returns:\n        The API representation of the completed query job.\n    \"\"\"\n    client = bigquery.Client(project=project_id, location=dataset_location)\n    if not job_config:\n        job_config = bigquery.QueryJobConfig()\n        job_config.create_disposition = bigquery.job.CreateDisposition.CREATE_IF_NEEDED\n        job_config.write_disposition = bigquery.job.WriteDisposition.WRITE_TRUNCATE\n    job_id = None\n    def cancel():\n        if job_id:\n            client.cancel_job(job_id)\n    with KfpExecutionContext(on_cancel=cancel) as ctx:\n        job_id = 'query_' + ctx.context_id()\n        query_job = _get_job(client, job_id)\n        table_ref = None\n        if not query_job:\n            dataset_ref = _prepare_dataset_ref(client, dataset_id, output_gcs_path, \n                dataset_location)\n            if dataset_ref:\n                if not table_id:\n                    table_id = job_id\n                table_ref = dataset_ref.table(table_id)\n                job_config.destination = table_ref\n            query_job = client.query(query, job_config, job_id=job_id)\n        _display_job_link(project_id, job_id)\n        query_job.result() # Wait for query to finish\n        if output_gcs_path:\n            job_id = 'extract_' + ctx.context_id()\n            extract_job = _get_job(client, job_id)\n            logging.info('Extracting data from table {} to {}.'.format(str(table_ref), output_gcs_path))\n            if not extract_job:\n                extract_job = client.extract_table(table_ref, output_gcs_path)\n            extract_job.result()  # Wait for export to finish\n        _dump_outputs(query_job, output_gcs_path, table_ref)\n        return query_job.to_api_repr()\n\ndef _get_job(client, job_id):\n    try:\n        return client.get_job(job_id)\n    except exceptions.NotFound:\n        return None\n\ndef _prepare_dataset_ref(client, dataset_id, output_gcs_path, dataset_location):\n    if not output_gcs_path and not dataset_id:\n        return None\n    \n    if not dataset_id:\n        dataset_id = 'kfp_tmp_dataset'\n    dataset_ref = client.dataset(dataset_id)\n    dataset = _get_dataset(client, dataset_ref)\n    if not dataset:\n        logging.info('Creating dataset {}'.format(dataset_id))\n        dataset = _create_dataset(client, dataset_ref, dataset_location)\n    return dataset_ref\n\ndef _get_dataset(client, dataset_ref):\n    try:\n        return client.get_dataset(dataset_ref)\n    except exceptions.NotFound:\n        return None\n\ndef _create_dataset(client, dataset_ref, location):\n    dataset = bigquery.Dataset(dataset_ref)\n    dataset.location = location\n    return client.create_dataset(dataset)\n\ndef _display_job_link(project_id, job_id):\n    display.display(display.Link(\n        href= 'https://console.cloud.google.com/bigquery?project={}'\n            '&j={}&page=queryresults'.format(project_id, job_id),\n        text='Query Details'\n    ))\n\ndef _dump_outputs(job, output_path, table_ref):\n    gcp_common.dump_file(KFP_OUTPUT_PATH + 'bigquery/query-job.json', \n        json.dumps(job.to_api_repr()))\n    if not output_path:\n        output_path = ''\n    gcp_common.dump_file(KFP_OUTPUT_PATH + 'bigquery/query-output-path.txt', \n        output_path)\n    (dataset_id, table_id) = (table_ref.dataset_id, table_ref.table_id) if table_ref else ('', '')\n    gcp_common.dump_file(KFP_OUTPUT_PATH + 'bigquery/query-dataset-id.txt', \n        dataset_id)\n    gcp_common.dump_file(KFP_OUTPUT_PATH + 'bigquery/query-table-id.txt', \n        table_id)\n", "path": "component_sdk/python/kfp_component/google/bigquery/_query.py"}]}