feat: managed service accounts for BYOB users (#2508)

* feat: managed service accounts for BYOB users

* chore: update staging and prod secrets with managed SA count

* fix: add handling of when system is enabled but byob is disabled

* chore: fix failing tests

* chore: simplify logic, formatting

* fix: update function head to match :query

* chore: version bump

* chore: fix failing test

* feat: add iam policy updating on config change

Author: Ziinc

VERSION

@@ -1 +1 @@
-1.16.0
+1.16.1

cloudbuild/.prod.env.enc

@@ -98,6 +98,15 @@ Assign `BigQuery Data Owner` and `BigQuery Job User` permissions to the Logflare
 
 ![BigQuery Job User Permissions](./bq-job-user-permissions.png)
 
+:::note
+If using [managed service accounts](#managed-service-accounts), include the following additional roles:
+
+- `roles/resourcemanager.projectIamAdmin`
+- `roles/iam.serviceAccountCreator`
+- `roles/iam.serviceAccountTokenCreator`
+
+:::
+
 #### Step 3: Update Account Settings in Logflare
 
 Find the GCP project ID in the [dashboard](https://console.cloud.google.com/home/dashboard)
@@ -118,6 +127,24 @@ You can also optionally update your sources' TTL to tweak how long you want to r
 The steps for setting up self-hosted Logflare requires different BigQuery configurations, please refer to the [self-hosting](/self-hosting) documentation for more details.
 :::
 
+### Managed Service Accounts
+
+When query volume for your instance is high, you may experience BigQuery rate limiting for their REST API. This occurs when BigQuery receives over 100 requests per second per user. This adversely affects the Logflare Endpoints functionality as well as Logflare Search functionality.
+
+By default, all BYOB users have managed service accounts disabled. In order to enable routing requests through mangaged service accounts, go to **Accounts > BigQuery Backend > Managed Service Accounts** and enable the setting by ticking the checkbox.
+
+Ensure that the Logflare service account has the following roles:
+
+- `roles/bigquery.admin` **(no change)**
+- `roles/resourcemanager.projectIamAdmin` **(new)** - used to set the IAM policy with the managed service accounts
+- `roles/iam.serviceAccountCreator` **(new)** - used to create new managed service accounts
+- `roles/iam.serviceAccountTokenCreator` **(new)** - used to generate short-lived tokens for authenticating with BigQuery REST API
+
+Enabling this option will have the following effects:
+
+- BigQuery REST API requests made will be routed through multiple managed service accounts. Each managed service account has the `logflare-managed-` prefix followed by the index (for example `logflare-managed-0`).
+- IAM policy for this project will be completely managed by Logflare, and permissions for managed service accounts will be handled in a non-destructive manner.
+
 ## Querying in BigQuery
 
 You can directly execute SQL queries in BigQuery instead of through the Logflare UI. This would be helpful for generating reports that require aggregations, or to perform queries across multiple BigQuery tables.

cloudbuild/.staging.env.enc

@@ -14,6 +14,7 @@ defmodule Logflare.Backends.Adaptor.BigQueryAdaptor do
   alias Logflare.Backends
   alias Logflare.Google.BigQuery.GenUtils
   alias Logflare.Google.CloudResourceManager
+  alias Logflare.Google
   use Supervisor
   require Logger
 
@@ -161,6 +162,14 @@ defmodule Logflare.Backends.Adaptor.BigQueryAdaptor do
     "#{@service_account_prefix}-#{service_account_index}"
   end
 
+  @doc """
+  Returns a list of all managed service account ids
+  """
+  def managed_service_account_ids() do
+    for i <- 0..(managed_service_account_pool_size() - 1),
+        do: managed_service_account_id(i)
+  end
+
   @doc """
   Lists all managed service accounts.
 
@@ -395,7 +404,19 @@ defmodule Logflare.Backends.Adaptor.BigQueryAdaptor do
     iex> update_iam_policy()
     :ok
   """
-  def update_iam_policy() do
+  def update_iam_policy(user \\ nil) do
     CloudResourceManager.set_iam_policy(async: false)
+
+    if Map.get(user || %{}, :bigquery_project_id) do
+      # byob project, maybe append managed SA to policy
+      append_managed_sa_to_iam_policy(user) |> dbg()
+    end
   end
+
+  defdelegate get_iam_policy(user), to: CloudResourceManager
+  defdelegate append_managed_sa_to_iam_policy(user), to: CloudResourceManager
+  defdelegate append_managed_service_accounts(project_id, policy), to: CloudResourceManager
+  defdelegate patch_dataset_access(user), to: Google.BigQuery
+
+  defdelegate get_conn(conn_type), to: GenUtils
 end

docs/docs.logflare.com/docs/backends/bigquery/index.md

@@ -52,7 +52,7 @@ defmodule Logflare.BqRepo do
       |> then(fn map -> struct(QueryRequest, map) end)
 
     result =
-      GenUtils.get_conn(:query)
+      GenUtils.get_conn({:query, user})
       |> Api.Jobs.bigquery_jobs_query(project_id, body: query_request)
       |> GenUtils.maybe_parse_google_api_result()
 

lib/logflare/backends/adaptor/bigquery_adaptor.ex

@@ -78,21 +78,30 @@ defmodule Logflare.Google.BigQuery.GenUtils do
 
   Uses `Logflare.FinchDefault` by default
   """
-  @typep conn_type :: :ingest | :query | :default
+  @typep conn_type :: :ingest | {:query, User.t()} | :default
   @spec get_conn(conn_type()) :: Tesla.Env.client()
   def get_conn(conn_type \\ :default) do
+    system_managed_sa_enabled = BigQueryAdaptor.managed_service_accounts_enabled?()
     # use pid as the partition hash
-    partition_count =
-      if conn_type == :query do
-        BigQueryAdaptor.managed_service_account_partition_count()
-      else
-        BigQueryAdaptor.ingest_service_account_partition_count()
+    {use_managed_sa?, partition_count} =
+      case conn_type do
+        {:query,
+         %_{bigquery_project_id: project_id, bigquery_enable_managed_service_accounts: true}}
+        when system_managed_sa_enabled == true and project_id != nil ->
+          {true, BigQueryAdaptor.managed_service_account_partition_count()}
+
+        {:query, %_{bigquery_project_id: nil}}
+        when system_managed_sa_enabled == true ->
+          {true, BigQueryAdaptor.managed_service_account_partition_count()}
+
+        _ ->
+          {false, BigQueryAdaptor.ingest_service_account_partition_count()}
       end
 
     partition = :erlang.phash2(self(), partition_count)
 
     {name, metadata} =
-      if conn_type == :query && BigQueryAdaptor.managed_service_accounts_enabled?() do
+      if use_managed_sa? == true do
         pool_size = BigQueryAdaptor.managed_service_account_pool_size()
 
         sa_index = :erlang.phash2(self(), pool_size)
@@ -144,7 +153,7 @@ defmodule Logflare.Google.BigQuery.GenUtils do
     ).adapter
   end
 
-  defp build_tesla_adapter_call(:query) do
+  defp build_tesla_adapter_call({:query, _}) do
     Tesla.client(
       [],
       {Tesla.Adapter.Finch, name: Logflare.FinchQuery, receive_timeout: 60_000}

lib/logflare/ecto/bigquery/bq_repo.ex

@@ -6,6 +6,7 @@ defmodule Logflare.Google.CloudResourceManager do
   alias GoogleApi.CloudResourceManager.V1.Model
   alias Logflare.Google.BigQuery.GenUtils
   alias Logflare.Users
+  alias Logflare.User
   alias Logflare.TeamUsers
   alias Logflare.Utils.Tasks
   alias Logflare.Backends.Adaptor.BigQueryAdaptor
@@ -27,6 +28,80 @@ defmodule Logflare.Google.CloudResourceManager do
     )
   end
 
+  def get_iam_policy(%User{bigquery_project_id: nil}), do: {:error, :no_project_id}
+
+  def get_iam_policy(user) do
+    conn = GenUtils.get_conn()
+
+    Api.Projects.cloudresourcemanager_projects_get_iam_policy(
+      conn,
+      user.bigquery_project_id,
+      body: %Model.GetIamPolicyRequest{}
+    )
+  end
+
+  def append_managed_sa_to_iam_policy(%User{bigquery_project_id: nil}),
+    do: {:error, :no_project_id}
+
+  def append_managed_sa_to_iam_policy(%User{bigquery_enable_managed_service_accounts: false}),
+    do: {:error, :managed_service_accounts_disabled}
+
+  def append_managed_sa_to_iam_policy(user) do
+    with {:enabled?, true} <- {:enabled?, BigQueryAdaptor.managed_service_accounts_enabled?()},
+         {:ok, policy} <- get_iam_policy(user),
+         {:contains?, _policy, false} <-
+           {:contains?, policy, contains_managed_service_accounts?(policy)} do
+      append_managed_service_accounts(user.bigquery_project_id, policy)
+    else
+      {:contains?, policy, _} -> {:ok, policy}
+      {:enabled?, false} -> {:error, :managed_service_accounts_disabled}
+      {:error, _err} = err -> err
+    end
+  end
+
+  # returns false if missing any of the managed service accounts
+  defp contains_managed_service_accounts?(%Model.Policy{bindings: bindings}) do
+    ids = BigQueryAdaptor.managed_service_account_ids()
+    members = Enum.flat_map(bindings, fn %{members: members} -> members end)
+
+    Enum.all?(ids, fn id ->
+      Enum.any?(members, fn member ->
+        member =~ id
+      end)
+    end)
+  end
+
+  def append_managed_service_accounts(project_id, %Model.Policy{bindings: bindings})
+      when project_id != nil do
+    conn = GenUtils.get_conn()
+
+    members =
+      for %{email: name} <- BigQueryAdaptor.list_managed_service_accounts() do
+        "serviceAccount:" <> name
+      end
+
+    new_binding = %Model.Binding{members: members, role: "roles/bigquery.admin"}
+
+    policy = %Model.Policy{bindings: [new_binding | bindings]}
+    body = %Model.SetIamPolicyRequest{policy: policy}
+
+    case Api.Projects.cloudresourcemanager_projects_set_iam_policy(conn, project_id, body: body) do
+      {:ok, response} ->
+        Logger.info(
+          "Appended managed service accounts to IAM policy for #{project_id} successful"
+        )
+
+        {:ok, response}
+
+      {:error, err} ->
+        Logger.error(
+          "Append managed service accounts to IAM policy for #{project_id} unknown error: #{inspect(err)}"
+        )
+
+        {:error, err}
+    end
+  end
+
   def set_iam_policy(opts \\ [async: true])
 
   def set_iam_policy(async: true) do

lib/logflare/google/bigquery/gen_utils/gen_utils.ex

@@ -78,6 +78,7 @@ defmodule Logflare.User do
     field :bigquery_dataset_id, :string
     field :bigquery_udfs_hash, :string
     field :bigquery_processed_bytes_limit, :integer
+    field :bigquery_enable_managed_service_accounts, :boolean, default: false
     field :api_quota, :integer, default: @default_user_api_quota
     field :valid_google_account, :boolean
     field :provider_uid, :string
@@ -113,6 +114,7 @@ defmodule Logflare.User do
     :bigquery_dataset_location,
     :bigquery_dataset_id,
     :bigquery_processed_bytes_limit,
+    :bigquery_enable_managed_service_accounts,
     :valid_google_account,
     :provider_uid,
     :company,

lib/logflare/google/resource_manager.ex

@@ -3,7 +3,7 @@ defmodule Logflare.Users do
 
   import Ecto.Query
   alias Logflare.Google.BigQuery
-  alias Logflare.Google.CloudResourceManager
+  alias Logflare.Backends.Adaptor.BigQueryAdaptor
   alias Logflare.Repo
   alias Logflare.Source.Supervisor
   alias Logflare.Sources
@@ -219,7 +219,7 @@ defmodule Logflare.Users do
     case Repo.delete(user) do
       {:ok, _user} = response ->
         BigQuery.delete_dataset(user)
-        CloudResourceManager.set_iam_policy()
+        BigQueryAdaptor.update_iam_policy()
         response
 
       {:error, _reason} = response ->