feat: initial implementation of bq reservations for endpoints

lib/logflare/backends/adaptor/bigquery_adaptor.ex

@@ -531,6 +531,22 @@ defmodule Logflare.Backends.Adaptor.BigQueryAdaptor do
     ]
   end
 
+  @spec parse_and_select_reservation(String.t() | nil) :: String.t() | nil
+  defp parse_and_select_reservation(nil), do: nil
+  defp parse_and_select_reservation(""), do: nil
+
+  defp parse_and_select_reservation(reservations_string) when is_binary(reservations_string) do
+    reservations_string
+    |> String.split("\n")
+    |> Enum.map(&String.trim/1)
+    |> Enum.reject(&(&1 == ""))
+    |> case do
+      [] -> nil
+      [single] -> single
+      reservations -> Enum.random(reservations)
+    end
+  end
+
   @spec execute_query_with_context(
           user_id :: integer(),
           query_string :: String.t(),
@@ -578,6 +594,16 @@ defmodule Logflare.Backends.Adaptor.BigQueryAdaptor do
             )
         ]
 
+    # At API query time, select endpoint-level reservation if configured
+    endpoint_reservation = parse_and_select_reservation(endpoint_query.bigquery_reservations)
+
+    query_opts =
+      if endpoint_reservation do
+        Keyword.put(query_opts, :reservation, endpoint_reservation)
+      else
+        query_opts
+      end
+
     execute_user_query(user, query_string, bq_params, query_opts)
   end
 

lib/logflare/endpoints/query.ex

@@ -31,7 +31,8 @@ defmodule Logflare.Endpoints.Query do
              :max_limit,
              :enable_auth,
              :labels,
-             :redact_pii
+             :redact_pii,
+             :bigquery_reservations
            ]}
   typed_schema "endpoint_queries" do
     field(:token, Ecto.UUID, autogenerate: true)
@@ -47,6 +48,7 @@ defmodule Logflare.Endpoints.Query do
     field(:enable_auth, :boolean, default: true)
     field(:redact_pii, :boolean, default: false)
     field(:labels, :string)
+    field(:bigquery_reservations, :string)
     field(:parsed_labels, :map, virtual: true)
     field(:metrics, :map, virtual: true)
 
@@ -81,7 +83,8 @@ defmodule Logflare.Endpoints.Query do
       :language,
       :description,
       :backend_id,
-      :labels
+      :labels,
+      :bigquery_reservations
     ])
     |> infer_language_from_backend()
     |> validate_required([:name, :query, :language])
@@ -102,7 +105,8 @@ defmodule Logflare.Endpoints.Query do
       :language,
       :description,
       :backend_id,
-      :labels
+      :labels,
+      :bigquery_reservations
     ])
     |> infer_language_from_backend()
     |> validate_query(:query)

lib/logflare_web/live/endpoints/actions/show.html.heex

@@ -63,6 +63,10 @@
       <li :if={@show_endpoint.labels} class="list-group-item">
         <span><strong>labels:</strong> {@show_endpoint.labels}</span>
       </li>
+      <li :if={@show_endpoint.language == :bq_sql && @show_endpoint.bigquery_reservations && @show_endpoint.bigquery_reservations != ""} class="list-group-item">
+        <span><strong>BigQuery reservations:</strong></span>
+        <pre class="tw-whitespace-pre-wrap tw-text-sm tw-mt-1 tw-mb-0 tw-bg-transparent tw-border-0 tw-p-0">{@show_endpoint.bigquery_reservations}</pre>
+      </li>
     </ul>
   </div>
 

lib/logflare_web/live/endpoints/components/endpoint_form.html.heex

@@ -46,6 +46,20 @@
       </small>
     </div>
 
+    <div :if={@determined_language == :bq_sql} class="form-group">
+      {label(f, :bigquery_reservations, "BigQuery Reservations (optional)")}
+      {textarea(f, :bigquery_reservations,
+        class: "form-control",
+        rows: 4,
+        placeholder: "projects/123/locations/us/reservations/my-reservation-1\nprojects/123/locations/us/reservations/my-reservation-2",
+        "phx-update": "ignore"
+      )}
+      {error_tag(f, :bigquery_reservations)}
+      <small class="form-text text-muted">
+        Specify BigQuery reservation paths (one per line). When this endpoint is queried, one reservation will be randomly selected. Empty lines are ignored.
+      </small>
+    </div>
+
     <.header_with_anchor text="Query" />
     <section class="tw-flex tw-flex-col">
       <div class="form-group">

priv/repo/migrations/20251210155948_add_bigquery_reservations_to_endpoint_queries.exs

@@ -0,0 +1,9 @@
+defmodule Logflare.Repo.Migrations.AddBigqueryReservationsToEndpointQueries do
+  use Ecto.Migration
+
+  def change do
+    alter table(:endpoint_queries) do
+      add :bigquery_reservations, :text
+    end
+  end
+end

test/logflare/endpoints_test.exs

@@ -565,4 +565,156 @@ defmodule Logflare.EndpointsTest do
       assert Endpoints.derive_language_from_backend_id(backend.id) == :pg_sql
     end
   end
+
+  describe "endpoint-level bigquery reservations" do
+    test "run_query/1 uses single endpoint reservation" do
+      pid = self()
+
+      expect(GoogleApi.BigQuery.V2.Api.Jobs, :bigquery_jobs_query, 1, fn _conn, _proj_id, opts ->
+        send(pid, {:reservation, opts[:body].reservation})
+        {:ok, TestUtils.gen_bq_response([%{"testing" => "123"}])}
+      end)
+
+      user = insert(:user)
+      reservation = "projects/123/locations/us/reservations/my-endpoint-res"
+
+      endpoint =
+        insert(:endpoint,
+          user: user,
+          query: "select current_datetime() as testing",
+          bigquery_reservations: reservation
+        )
+
+      assert {:ok, %{rows: [%{"testing" => _}]}} = Endpoints.run_query(endpoint)
+
+      assert_receive {:reservation, ^reservation}
+    end
+
+    test "run_query/1 selects from multiple endpoint reservations" do
+      pid = self()
+
+      expect(GoogleApi.BigQuery.V2.Api.Jobs, :bigquery_jobs_query, 1, fn _conn, _proj_id, opts ->
+        send(pid, {:reservation, opts[:body].reservation})
+        {:ok, TestUtils.gen_bq_response([%{"testing" => "123"}])}
+      end)
+
+      user = insert(:user)
+
+      reservations =
+        "projects/123/locations/us/reservations/res1\nprojects/123/locations/us/reservations/res2"
+
+      endpoint =
+        insert(:endpoint,
+          user: user,
+          query: "select current_datetime() as testing",
+          bigquery_reservations: reservations
+        )
+
+      assert {:ok, %{rows: [%{"testing" => _}]}} = Endpoints.run_query(endpoint)
+
+      # Verify one of the configured reservations is used
+      assert_receive {:reservation, res}
+
+      assert res in [
+               "projects/123/locations/us/reservations/res1",
+               "projects/123/locations/us/reservations/res2"
+             ]
+    end
+
+    test "run_query/1 ignores empty lines in reservations" do
+      pid = self()
+
+      expect(GoogleApi.BigQuery.V2.Api.Jobs, :bigquery_jobs_query, 1, fn _conn, _proj_id, opts ->
+        send(pid, {:reservation, opts[:body].reservation})
+        {:ok, TestUtils.gen_bq_response([%{"testing" => "123"}])}
+      end)
+
+      user = insert(:user)
+
+      # Reservations with empty lines and whitespace
+      reservations = """
+      projects/123/locations/us/reservations/valid-res
+
+
+      """
+
+      endpoint =
+        insert(:endpoint,
+          user: user,
+          query: "select current_datetime() as testing",
+          bigquery_reservations: reservations
+        )
+
+      assert {:ok, %{rows: [%{"testing" => _}]}} = Endpoints.run_query(endpoint)
+
+      # Should use the only valid reservation
+      assert_receive {:reservation, "projects/123/locations/us/reservations/valid-res"}
+    end
+
+    test "run_query/1 does not set reservation when bigquery_reservations is nil" do
+      pid = self()
+
+      expect(GoogleApi.BigQuery.V2.Api.Jobs, :bigquery_jobs_query, 1, fn _conn, _proj_id, opts ->
+        send(pid, {:reservation, opts[:body].reservation})
+        {:ok, TestUtils.gen_bq_response([%{"testing" => "123"}])}
+      end)
+
+      user = insert(:user)
+
+      endpoint =
+        insert(:endpoint,
+          user: user,
+          query: "select current_datetime() as testing",
+          bigquery_reservations: nil
+        )
+
+      assert {:ok, %{rows: [%{"testing" => _}]}} = Endpoints.run_query(endpoint)
+
+      assert_receive {:reservation, nil}
+    end
+
+    test "run_query/1 does not set reservation when bigquery_reservations is empty string" do
+      pid = self()
+
+      expect(GoogleApi.BigQuery.V2.Api.Jobs, :bigquery_jobs_query, 1, fn _conn, _proj_id, opts ->
+        send(pid, {:reservation, opts[:body].reservation})
+        {:ok, TestUtils.gen_bq_response([%{"testing" => "123"}])}
+      end)
+
+      user = insert(:user)
+
+      endpoint =
+        insert(:endpoint,
+          user: user,
+          query: "select current_datetime() as testing",
+          bigquery_reservations: ""
+        )
+
+      assert {:ok, %{rows: [%{"testing" => _}]}} = Endpoints.run_query(endpoint)
+
+      assert_receive {:reservation, nil}
+    end
+
+    test "run_query/1 does not set reservation when bigquery_reservations contains only whitespace" do
+      pid = self()
+
+      expect(GoogleApi.BigQuery.V2.Api.Jobs, :bigquery_jobs_query, 1, fn _conn, _proj_id, opts ->
+        send(pid, {:reservation, opts[:body].reservation})
+        {:ok, TestUtils.gen_bq_response([%{"testing" => "123"}])}
+      end)
+
+      user = insert(:user)
+
+      endpoint =
+        insert(:endpoint,
+          user: user,
+          query: "select current_datetime() as testing",
+          bigquery_reservations: "   \n  \n   "
+        )
+
+      assert {:ok, %{rows: [%{"testing" => _}]}} = Endpoints.run_query(endpoint)
+
+      assert_receive {:reservation, nil}
+    end
+  end
 end