support timeoutMs param

livebook-dev · Nov 16, 2023 · 165db07 · 165db07
1 parent fd11b4b
commit 165db07
Show file tree

Hide file tree

Showing 2 changed files with 15 additions and 4 deletions.
diff --git a/lib/req_bigquery.ex b/lib/req_bigquery.ex
@@ -10,10 +10,11 @@ defmodule ReqBigQuery do
   alias Req.Request
   alias ReqBigQuery.Result
 
-  @allowed_options ~w(goth default_dataset_id project_id bigquery max_results use_legacy_sql)a
+  @allowed_options ~w(goth default_dataset_id project_id bigquery max_results use_legacy_sql timeout_ms)a
   @base_url "https://bigquery.googleapis.com/bigquery/v2"
   @max_results 10_000
   @use_legacy_sql false
+  @timeout_ms 10_000
 
   @doc """
   Attaches to Req request.
@@ -40,6 +41,11 @@ defmodule ReqBigQuery do
       If set to false, the query will use BigQuery's GoogleSQL: https://cloud.google.com/bigquery/sql-reference/
       The default value is false.
 
+    * `:timeout_ms` - Optional. How long to wait for the query to complete, in milliseconds, before the request times out and returns.
+      Note: The call is not guaranteed to wait for the specified timeout. 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).
+
   If you want to set any of these options when attaching the plugin, pass them as the second argument.
 
   ## Examples
@@ -114,6 +120,7 @@ defmodule ReqBigQuery do
       |> Keyword.put_new(:base_url, @base_url)
       |> Keyword.put_new(:max_results, @max_results)
       |> Keyword.put_new(:use_legacy_sql, @use_legacy_sql)
+      |> Keyword.put_new(:timeout_ms, @timeout_ms)
 
     request
     |> Request.prepend_request_steps(bigquery_run: &run/1)
@@ -134,6 +141,7 @@ defmodule ReqBigQuery do
         |> build_request_body(options[:default_dataset_id])
         |> Map.put(:maxResults, options[:max_results])
         |> Map.put(:useLegacySql, options[:use_legacy_sql])
+        |> Map.put(:timeoutMs, options[:timeout_ms])
 
       %{request | url: uri}
       |> Request.merge_options(auth: {:bearer, token}, json: json)

diff --git a/test/req_bigquery_test.exs b/test/req_bigquery_test.exs
@@ -19,7 +19,8 @@ defmodule ReqBigQueryTest do
                "defaultDataset" => %{"datasetId" => "my_awesome_dataset"},
                "query" => "select * from iris",
                "maxResults" => 10000,
-               "useLegacySql" => true
+               "useLegacySql" => true,
+               "timeoutMs" => 20000
              }
 
       assert URI.to_string(request.url) ==
@@ -51,7 +52,8 @@ defmodule ReqBigQueryTest do
       goth: ctx.test,
       project_id: "my_awesome_project_id",
       default_dataset_id: "my_awesome_dataset",
-      use_legacy_sql: true
+      use_legacy_sql: true,
+      timeout_ms: 20_000
     ]
 
     assert response =
@@ -96,7 +98,8 @@ defmodule ReqBigQueryTest do
                  }
                ],
                "useLegacySql" => false,
-               "maxResults" => 10000
+               "maxResults" => 10000,
+               "timeoutMs" => 10000
              }
 
       assert URI.to_string(request.url) ==