Add ability to pull/push under different goth service accounts.

.gitignore

@@ -9,3 +9,5 @@ config/*credentials.json
 .#*
 /doc/
 .elixir_ls/
+.idea/
+*.iml

README.md

@@ -41,3 +41,40 @@ Diplomat.Query.new(
   %{name: "20,000 Leagues Under The Sea"}
 ) |> Diplomat.Query.execute
 ```
+
+
+#### Use multiple accounts with Diplomat
+Configure Goth with additional accounts.
+```elixir
+{:ok, alternative_account} = Jason.decode(File.read!("priv/goth/alternative-account.json"))
+Goth.Config.add_config(alternative_account)
+```
+
+Require Diplomat and use the with_account option to set current (and only current) process to use alternative process within block. 
+
+The account name will be the client_email value from the additional Goth configuration you added by default.
+
+```elixir
+require Diplomat
+# copy data from prod to stage environment
+
+# 1. Fetch data from production account
+prod_entities = Diplomat.with_account(alternative_account["client_email"]) do
+  Diplomat.Query.new(
+    "select * from `Book` where name = @name",
+    %{name: "20,000 Leagues Under The Sea"}
+  ) |> Diplomat.Query.execute
+end
+
+# 2. Write to stage/dev account (default environment)
+target_project = Diplomat.Client.project()
+stage_entities = Enum.map(prod_entities, fn(entity) ->
+      put_in(entity, [Access.key(:key), Access.key(:project_id)], target_project)
+    end)
+Diplomat.Entity.upsert(stage_entities)
+```
+
+
+
+
+

lib/diplomat.ex

@@ -9,4 +9,21 @@ defmodule Diplomat do
   defmodule Proto do
     use Protobuf, from: Path.expand("datastore_v1beta3.proto", __DIR__), doc: false
   end
+
+  defmacro with_account(account, [do: block]) do
+    account = Macro.expand(account, __ENV__)
+    quote do
+      Process.put(:diplomat_account_queue, [unquote(account)] ++ Process.get(:diplomat_account_queue, []))
+      Process.put(:diplomat_account, unquote(account))
+      diplomat__response = (unquote(block))
+      diplomat__account_queue = case Process.get(:diplomat_account_queue, []) do
+                                  [_|t] -> t
+                                  _ -> []
+                                end
+      Process.put(:diplomat_account_queue, diplomat__account_queue)
+      Process.put(:diplomat_account, List.first(diplomat__account_queue))
+      diplomat__response
+    end
+  end
+
 end

lib/diplomat/client.ex

@@ -164,16 +164,26 @@ defmodule Diplomat.Client do
 
   defp token_module, do: Application.get_env(:diplomat, :token_module, Goth.Token)
 
-  defp project do
-    {:ok, project_id} = Goth.Config.get(:project_id)
+  def diplomat_account() do
+   Process.get(:diplomat_account)
+  end
+
+  def project do
+    {:ok, project_id} = case diplomat_account() do
+                          nil -> Goth.Config.get(:project_id)
+                          account -> Goth.Config.get(account, :project_id)
+                        end
     project_id
   end
 
   defp api_scope, do: api_scope(@api_version)
   defp api_scope("v1"), do: "https://www.googleapis.com/auth/datastore"
 
   defp auth_header do
-    {:ok, token} = token_module().for_scope(api_scope())
+    {:ok, token} = case diplomat_account() do
+                     nil -> token_module().for_scope(api_scope())
+                     account -> token_module().for_scope({account, api_scope()})
+                   end
     {"Authorization", "#{token.type} #{token.token}"}
   end
 

lib/diplomat/key.ex

@@ -92,7 +92,7 @@ defmodule Diplomat.Key do
           nil
 
         _ ->
-          {:ok, global_project_id} = Goth.Config.get(:project_id)
+          global_project_id = Diplomat.Client.project()
 
           PbPartition.new(
             project_id: key.project_id || global_project_id,

lib/diplomat/query.ex

@@ -43,7 +43,7 @@ defmodule Diplomat.Query do
 
   @spec execute(t, String.t() | nil) :: [Entity.t()] | Client.error()
   def execute(%__MODULE__{} = q, namespace \\ nil) do
-    {:ok, project} = Goth.Config.get(:project_id)
+    project = Diplomat.Client.project()
 
     RunQueryRequest.new(
       query_type: {:gql_query, q |> Query.proto()},
@@ -54,7 +54,7 @@ defmodule Diplomat.Query do
 
   @spec execute_with_pagination(t, String.t() | nil) :: [QueryResultBatch.t()] | Client.error()
   def execute_with_pagination(%__MODULE__{} = q, namespace \\ nil) do
-    {:ok, project} = Goth.Config.get(:project_id)
+    project = Diplomat.Client.project()
 
     RunQueryRequest.new(
       query_type: {:gql_query, q |> Query.proto()},

test/diplomat/account_override_test.exs

@@ -0,0 +1,27 @@
+defmodule Diplomat.AccountOverrideTest do
+  use ExUnit.Case
+  require Diplomat
+
+  test "Verify account overrides are applied correctly" do
+    assert Diplomat.Client.diplomat_account() == nil
+    r = Diplomat.with_account(:alternative_account_email) do
+      assert Diplomat.Client.diplomat_account() == :alternative_account_email
+      a = Diplomat.with_account(:alternative_account2_email) do
+        assert Diplomat.Client.diplomat_account() == :alternative_account2_email
+        :alpha
+      end
+      assert a == :alpha
+      assert Diplomat.Client.diplomat_account() == :alternative_account_email
+      b = Diplomat.with_account(:alternative_account3_email) do
+        assert Diplomat.Client.diplomat_account() == :alternative_account3_email
+        :beta
+      end
+      assert b == :beta
+      assert Diplomat.Client.diplomat_account() == :alternative_account_email
+      :omega
+    end
+    assert Diplomat.Client.diplomat_account() == nil
+    assert r == :omega
+  end
+
+end