perf: increase expiration time for auth context (#2268)

* perf: increase expiration time for auth context

* perf: increase auth cache ttl, add oauth_access_tokens to publication_tables

* chore: fix invalid expiration error

* chore: use alter publication

* chore: correct alter publication syntax

lib/logflare/auth/cache.ex

@@ -25,7 +25,7 @@ defmodule Logflare.Auth.Cache do
                  Utils.cache_limit(100_000)
                ]
                |> Enum.filter(& &1),
-             expiration: Utils.cache_expiration_sec(30, 15)
+             expiration: Utils.cache_expiration_min(5, 2)
            ]
          ]}
     }

lib/logflare/context_cache.ex

@@ -118,6 +118,7 @@ defmodule Logflare.ContextCache do
   end
 
   defp select_key(%_{id: id}), do: id
+  defp select_key({:ok, %_{id: id}}), do: id
   defp select_key(true), do: "true"
   defp select_key(nil), do: :not_found
   defp select_key(_), do: :unknown

lib/logflare/context_cache/cache_buster.ex

@@ -116,6 +116,14 @@ defmodule Logflare.ContextCache.CacheBuster do
     {Logflare.TeamUsers, String.to_integer(id)}
   end
 
+  defp handle_record(%UpdatedRecord{
+         relation: {_schema, "oauth_access_tokens"},
+         record: %{"id" => id}
+       })
+       when is_binary(id) do
+    {Logflare.Auth, String.to_integer(id)}
+  end
+
   defp handle_record(%NewRecord{
          relation: {_schema, "billing_accounts"},
          record: %{"id" => _id}
@@ -175,6 +183,14 @@ defmodule Logflare.ContextCache.CacheBuster do
     {Logflare.TeamUsers, :not_found}
   end
 
+  defp handle_record(%NewRecord{
+         relation: {_schema, "oauth_access_tokens"},
+         record: %{"id" => _id}
+       }) do
+    # When new records are created they were previously cached as `nil` so we need to bust the :not_found keys
+    {Logflare.Auth, :not_found}
+  end
+
   defp handle_record(%DeletedRecord{
          relation: {_schema, "billing_accounts"},
          old_record: %{"id" => id}
@@ -233,6 +249,15 @@ defmodule Logflare.ContextCache.CacheBuster do
     {Logflare.TeamUsers, String.to_integer(id)}
   end
 
+  defp handle_record(%DeletedRecord{
+         relation: {_schema, "oauth_access_tokens"},
+         old_record: %{"id" => id}
+       })
+       when is_binary(id) do
+    # Must do `alter table rules replica identity full` to get full records on deletes otherwise all fields are null
+    {Logflare.Auth, String.to_integer(id)}
+  end
+
   defp handle_record(_record) do
     :noop
   end

priv/repo/migrations/20241204120824_recreate_tables_publication.exs

@@ -0,0 +1,19 @@
+defmodule Logflare.Repo.Migrations.RecreateTablesPublication do
+  use Ecto.Migration
+  use Ecto.Migration
+
+  @publications Application.get_env(:logflare, Logflare.ContextCache.CacheBuster)[:publications]
+  @table "oauth_access_tokens"
+
+  def up do
+    for p <- @publications do
+      execute("ALTER PUBLICATION #{p} ADD TABLE #{@table};")
+    end
+  end
+
+  def down do
+    for p <- @publications do
+      execute("ALTER PUBLICATION #{p} DROP TABLE #{@table};")
+    end
+  end
+end

test/logflare/auth/cache_test.exs

@@ -3,6 +3,7 @@ defmodule Logflare.Auth.CacheTest do
   alias Logflare.Auth
   alias Logflare.Factory
   alias Logflare.Partners.Partner
+  alias Logflare.ContextCache
 
   setup do
     user = Factory.insert(:user)
@@ -47,6 +48,17 @@ defmodule Logflare.Auth.CacheTest do
     assert {:error, :unauthorized} = Auth.Cache.verify_access_token(key)
   end
 
+  test "cache busting", %{user: user} do
+    {:ok, key} = Auth.create_access_token(user)
+
+    Auth
+    |> expect(:verify_access_token, 2, fn _ -> {:ok, key} end)
+
+    assert {:ok, _} = Auth.Cache.verify_access_token(key.token)
+    ContextCache.bust_keys([{Auth, key.id}])
+    assert {:ok, _} = Auth.Cache.verify_access_token(key.token)
+  end
+
   defp access_token_fixture(user_or_team_or_partner) do
     {:ok, key} = Auth.create_access_token(user_or_team_or_partner)
     key

test/test_helper.exs

@@ -17,6 +17,7 @@ Mimic.copy(Logflare.Source.RateCounterServer)
 Mimic.copy(Logflare.Source.BigQuery.Schema)
 Mimic.copy(Logflare.SystemMetrics.AllLogsLogged)
 Mimic.copy(Logflare.Sources.Cache)
+Mimic.copy(Logflare.Auth)
 Mimic.copy(Logflare.SingleTenant)
 Mimic.copy(Logflare.Backends.Adaptor.WebhookAdaptor)
 Mimic.copy(Logflare.Backends.Adaptor.WebhookAdaptor.Client)