From b0e951dbf7624cded0da6039a4dd08304efccfd6 Mon Sep 17 00:00:00 2001
From: Jeff Deville <jeffdeville@gmail.com>
Date: Fri, 9 Sep 2016 14:33:14 -0400
Subject: [PATCH 1/3] Created a macro that will enforce existence of tenant ids
 on models where the model is tenanted

---
 config/test.exs                         |   4 +
 lib/apartmentex/repo.ex                 | 175 ++++++++++++++++++
 lib/apartmentex/tenant_missing_error.ex |   3 +
 test/lib/apartmentex/repo_test.exs      | 233 ++++++++++++++++++++++++
 test/support/test_repo.ex               | 126 +++++++++++++
 test/test_helper.exs                    |   1 +
 6 files changed, 542 insertions(+)
 create mode 100644 lib/apartmentex/repo.ex
 create mode 100644 lib/apartmentex/tenant_missing_error.ex
 create mode 100644 test/lib/apartmentex/repo_test.exs
 create mode 100644 test/support/test_repo.ex

diff --git a/config/test.exs b/config/test.exs
index d8a27da..1fe9be2 100644
--- a/config/test.exs
+++ b/config/test.exs
@@ -6,4 +6,8 @@ config :apartmentex, Apartmentex.TestPostgresRepo,
   adapter: Ecto.Adapters.Postgres,
   pool: Ecto.Adapters.SQL.Sandbox
 
+config :apartmentex, Ecto.TestRepo,
+  url: "ecto://user:pass@local/hello"
+
+
 config :logger, level: :warn
diff --git a/lib/apartmentex/repo.ex b/lib/apartmentex/repo.ex
new file mode 100644
index 0000000..808887d
--- /dev/null
+++ b/lib/apartmentex/repo.ex
@@ -0,0 +1,175 @@
+defmodule Apartmentex.Repo do
+  defmacro __using__(opts) do
+    quote bind_quoted: [opts: opts] do
+      @behaviour Ecto.Repo
+      alias Apartmentex.TenantMissingError
+
+      @repo Keyword.fetch!(opts, :repo)
+      @untenanted [Ecto.Migration.SchemaMigration] ++ Keyword.get(opts, :untenanted, [])
+
+      # From Ecto.Repo
+      defdelegate __adapter__, to: @repo
+      defdelegate __log__(entry), to: @repo
+      defdelegate config(), to: @repo
+      defdelegate start_link(opts \\ []), to: @repo
+      defdelegate stop(pid, timeout \\ 5000), to: @repo
+      defdelegate transaction(fun_or_multi, opts \\ []), to: @repo
+      defdelegate in_transaction?(), to: @repo
+      defdelegate rollback(value), to: @repo
+      defdelegate aggregate(queryable, aggregate, field, opts \\ []), to: @repo
+      defdelegate preload(struct_or_structs, preloads, opts \\ []), to: @repo
+
+      # From Ecto.Adapters.SQL
+      defdelegate __pool__, to: @repo
+      defdelegate __sql__, to: @repo
+
+      def all(queryable, opts \\ []) do
+        assert_tenant(queryable)
+        @repo.all(queryable, opts)
+      end
+
+      def get(queryable, id, opts \\ []) do
+        assert_tenant(queryable)
+        @repo.get(queryable, id, opts)
+      end
+
+      def get!(queryable, id, opts \\ []) do
+        assert_tenant(queryable)
+        @repo.get!(queryable, id, opts)
+      end
+
+      def get_by(queryable, clauses, opts \\ []) do
+        assert_tenant(queryable)
+        @repo.get_by(queryable, clauses, opts)
+      end
+
+      def get_by!(queryable, clauses, opts \\ []) do
+        assert_tenant(queryable)
+        @repo.get_by!(queryable, clauses, opts)
+      end
+
+      def one(queryable, opts \\ []) do
+        assert_tenant(queryable)
+        @repo.one(queryable, opts)
+      end
+
+      def one!(queryable, opts \\ []) do
+        assert_tenant(queryable)
+        @repo.one!(queryable, opts)
+      end
+
+      @insert_all_error """
+      For insert_all
+        - For tenanted tables
+            - Your first parameter must be a tuple with the prefix, and the table name
+        - For non-tenanted tables
+            - Your first parameter may not be the string name of the table, because we can't
+              check the associated model to see if it requires a tenant.
+      """
+      def insert_all(schema_or_source, entries, opts \\ [])
+      def insert_all({nil, source} = schema_or_source, entries, opts) do
+        if requires_tenant?(source) do
+          raise TenantMissingError, message: @insert_all_error
+        end
+        @repo.insert_all(schema_or_source, entries, opts)
+      end
+
+      def insert_all({_prefix, _source} = schema_or_source, entries, opts), do: @repo.insert_all(schema_or_source, entries, opts)
+      def insert_all(schema_or_source, entries, opts) when is_binary(schema_or_source), do: raise TenantMissingError, message: @insert_all_error
+      def insert_all(schema_or_source, entries, opts) when is_atom(schema_or_source) do
+        if requires_tenant?(schema_or_source) do
+          raise TenantMissingError, message: @insert_all_error
+        end
+        @repo.insert_all(schema_or_source, entries, opts)
+      end
+
+      def insert_all(schema_or_source, entries, opts) do
+        if requires_tenant?(schema_or_source) do
+          raise TenantMissingError, message: @insert_all_error
+        end
+        @repo.insert_all(schema_or_source, entries, opts)
+      end
+
+      def update_all(queryable, updates, opts \\ []) do
+        assert_tenant(queryable)
+        @repo.update_all(queryable, updates, opts)
+      end
+
+      def delete_all(queryable, opts \\ []) do
+        assert_tenant(queryable)
+        @repo.delete_all(queryable, opts)
+      end
+
+      def insert(struct, opts \\ []) do
+        assert_tenant(struct)
+        @repo.insert(struct, opts)
+      end
+
+      def update(struct, opts \\ []) do
+        assert_tenant(struct)
+        @repo.update(struct, opts)
+      end
+
+      def insert_or_update(changeset, opts \\ []) do
+        assert_tenant(changeset)
+        @repo.insert_or_update(changeset, opts)
+      end
+
+      def delete(struct, opts \\ []) do
+        assert_tenant(struct)
+        @repo.delete(struct, opts)
+      end
+
+      def insert!(struct, opts \\ []) do
+        assert_tenant(struct)
+        @repo.insert!(struct, opts)
+      end
+
+      def update!(struct, opts \\ []) do
+        assert_tenant(struct)
+        @repo.update!(struct, opts)
+      end
+
+      def insert_or_update!(changeset, opts \\ []) do
+        assert_tenant(changeset)
+        @repo.insert_or_update!(changeset, opts)
+      end
+
+      def delete!(struct, opts \\ []) do
+        assert_tenant(struct)
+        @repo.delete!(struct, opts)
+      end
+
+      defp assert_tenant(%Ecto.Changeset{} = changeset) do
+        assert_tenant(changeset.data)
+      end
+
+      defp assert_tenant(%{__meta__: _} = model) do
+        if requires_tenant?(model) && !has_prefix?(model) do
+          raise TenantMissingError, message: "No tenant specified in #{model.__struct__}"
+        end
+      end
+
+      defp assert_tenant(queryable) do
+        query = Ecto.Queryable.to_query(queryable)
+        if requires_tenant?(query) && !has_prefix?(query) do
+          raise TenantMissingError, message: "No tenant specified in #{get_model_from_query(query)}"
+        end
+      end
+
+      defp has_prefix?(%{__meta__: _} = model) do
+        if Ecto.get_meta(model, :prefix), do: true, else: false
+      end
+
+
+      defp get_model_from_query(%{from: {_, model}}), do: model
+
+      defp requires_tenant?(%{from: {_, model}}), do: not model in @untenanted
+      defp requires_tenant?(%{__struct__: model}), do: not model in @untenanted
+      defp requires_tenant?(model), do: not model in @untenanted
+
+      defp has_prefix?(%{prefix: nil}), do: false
+      defp has_prefix?(%{prefix: _}), do: true
+    end
+  end
+end
diff --git a/lib/apartmentex/tenant_missing_error.ex b/lib/apartmentex/tenant_missing_error.ex
new file mode 100644
index 0000000..5fe7245
--- /dev/null
+++ b/lib/apartmentex/tenant_missing_error.ex
@@ -0,0 +1,3 @@
+defmodule Apartmentex.TenantMissingError do
+  defexception message: "No tenant specified"
+end
diff --git a/test/lib/apartmentex/repo_test.exs b/test/lib/apartmentex/repo_test.exs
new file mode 100644
index 0000000..38f674b
--- /dev/null
+++ b/test/lib/apartmentex/repo_test.exs
@@ -0,0 +1,233 @@
+defmodule Apartmentex.Test.UntenantedRepo do
+  use Apartmentex.Repo, repo: Ecto.TestRepo, untenanted: [Apartmentex.Note]
+end
+
+defmodule Apartmentex.Test.TenantedRepo do
+  use Apartmentex.Repo, repo: Ecto.TestRepo
+end
+
+defmodule Apartmentex.ApartmentexTest do
+  use ExUnit.Case
+  import Apartmentex.PrefixBuilder
+  alias Apartmentex.{Note,TenantMissingError}
+  alias Apartmentex.Test.{TenantedRepo,UntenantedRepo}
+
+  @tenant_id 2
+  @error_message "No tenant specified in Elixir.Apartmentex.Note"
+
+  def set_tenant(%Ecto.Changeset{} = changeset, tenant) do
+    %{changeset | data: set_tenant(changeset.data, tenant)}
+  end
+
+  def set_tenant(%{__meta__: _} = model, tenant) do
+    Ecto.put_meta(model,  prefix: build_prefix(tenant))
+  end
+
+  def set_tenant(queryable, tenant) do
+    queryable
+    |> Ecto.Queryable.to_query
+    |> Map.put(:prefix, build_prefix(tenant))
+  end
+
+  def scoped_note_query do
+    Note
+    |> Ecto.Queryable.to_query
+    |> set_tenant(@tenant_id)
+  end
+
+  def scoped_note do
+    %Note{id: 1, body: "body"} |> set_tenant(@tenant_id)
+  end
+
+  def scoped_changeset do
+    Note.changeset(scoped_note, %{body: "body"})
+  end
+
+  test ".all/2 verifies tenant existence" do
+    assert_raise TenantMissingError, @error_message, fn ->
+      TenantedRepo.all(Note, [])
+    end
+    assert TenantedRepo.all(scoped_note_query, []) == [1]
+    assert UntenantedRepo.all(Note, []) == [1]
+  end
+
+  test ".get/2 verifies tenant existence" do
+    assert_raise TenantMissingError, @error_message, fn ->
+      TenantedRepo.get(Note, 1)
+    end
+
+    assert TenantedRepo.get(scoped_note_query, 1) == 1
+    assert UntenantedRepo.get(Note, 1) == 1
+  end
+
+  test ".get!/2 verifies tenant existence" do
+    assert_raise TenantMissingError, @error_message, fn ->
+      TenantedRepo.get(Note, 1)
+    end
+
+    assert TenantedRepo.get(scoped_note_query, 1) == 1
+    assert UntenantedRepo.get(Note, 1) == 1
+  end
+
+  test ".get_by(queryable, clauses, opts \\ []) verifies tenant existence" do
+    assert_raise TenantMissingError, @error_message, fn ->
+      TenantedRepo.get_by(Note, body: "immaterial")
+    end
+
+    assert TenantedRepo.get_by(scoped_note_query, body: "immaterial") == 1
+    assert UntenantedRepo.get_by(Note, body: "immaterial") == 1
+  end
+
+  test ".get_by!(queryable, clauses, opts \\ []) verifies tenant existence" do
+    assert_raise TenantMissingError, @error_message, fn ->
+      TenantedRepo.get_by!(Note, body: "immaterial")
+    end
+
+    assert TenantedRepo.get_by!(scoped_note_query, body: "immaterial") == 1
+    assert UntenantedRepo.get_by!(Note, body: "immaterial") == 1
+  end
+
+  test ".one(queryable, opts \\ []) verifies tenant existence" do
+    assert_raise TenantMissingError, @error_message, fn ->
+      TenantedRepo.one(Note)
+    end
+
+    assert TenantedRepo.one(scoped_note_query) == 1
+    assert UntenantedRepo.one(Note) == 1
+  end
+
+  test ".one!(queryable, opts \\ []) verifies tenant existence" do
+    assert_raise TenantMissingError, @error_message, fn ->
+      TenantedRepo.one!(Note)
+    end
+
+    assert TenantedRepo.one!(scoped_note_query) == 1
+    assert UntenantedRepo.one!(Note) == 1
+  end
+
+  test ".insert_all(schema_or_source, entries, opts \\ []) verifies tenant existence" do
+    assert_raise TenantMissingError, ~r/For insert_all/, fn ->
+      TenantedRepo.insert_all(Note, [%Note{body: "body0"}])
+    end
+    assert_raise TenantMissingError, ~r/For insert_all/, fn ->
+      TenantedRepo.insert_all("notes", [%Note{body: "body0"}])
+    end
+    assert_raise TenantMissingError, ~r/For insert_all/, fn ->
+      TenantedRepo.insert_all({nil, "notes"}, [%Note{body: "body0"}])
+    end
+    assert_raise TenantMissingError, ~r/For insert_all/, fn ->
+      TenantedRepo.insert_all({nil, Note}, [%Note{body: "body0"}])
+    end
+
+
+    assert TenantedRepo.insert_all({@tenant_id, "notes"}, [%{body: "body"}]) == {1, nil}
+    assert TenantedRepo.insert_all({@tenant_id, Note}, [%{body: "body"}]) == {1, nil}
+    assert UntenantedRepo.insert_all(Note, [%{body: "body0"}]) == {1, nil}
+  end
+
+  test ".update_all(queryable, updates, opts \\ []) verifies tenant existence" do
+    assert_raise TenantMissingError, @error_message, fn ->
+      TenantedRepo.update_all(Note, set: [body: "new"])
+    end
+
+    assert TenantedRepo.update_all(scoped_note_query, set: [body: "new"])
+    assert UntenantedRepo.update_all(Note, set: [body: "new"]) == {1, nil}
+  end
+
+  test ".delete_all(queryable, opts \\ []) verifies tenant existence" do
+    assert_raise TenantMissingError, @error_message, fn ->
+      TenantedRepo.delete_all(Note)
+    end
+
+    assert TenantedRepo.delete_all(scoped_note_query) == {1, nil}
+    assert UntenantedRepo.delete_all(Note) == {1, nil}
+  end
+
+  test ".insert(struct, opts \\ []) verifies tenant existence" do
+    assert_raise TenantMissingError, @error_message, fn ->
+      TenantedRepo.insert(Note.changeset(%Note{}, %{}))
+    end
+    assert_raise TenantMissingError, @error_message, fn ->
+      TenantedRepo.insert(%Note{})
+    end
+
+    assert {:ok, _} = TenantedRepo.insert(scoped_note)
+    assert {:ok, _} = TenantedRepo.insert(Note.changeset(scoped_note, %{}))
+    assert {:ok, _} = UntenantedRepo.insert(%Note{})
+  end
+
+  test ".update(struct, opts \\ []) verifies tenant existence" do
+    assert_raise TenantMissingError, @error_message, fn ->
+      TenantedRepo.update(Note.changeset(%Note{id: 1}, %{body: "body"}))
+    end
+
+    assert TenantedRepo.update(scoped_changeset)
+    assert UntenantedRepo.update(Note.changeset(%Note{id: 1}, %{body: "body"}))
+  end
+
+  test ".insert_or_update(changeset, opts \\ []) verifies tenant existence" do
+    assert_raise TenantMissingError, @error_message, fn ->
+      TenantedRepo.insert_or_update(Note.changeset(%Note{id: 1}, %{body: "body"}))
+    end
+
+    assert {:ok, _} = TenantedRepo.insert_or_update(scoped_changeset)
+    assert {:ok, _} = UntenantedRepo.insert_or_update(Note.changeset(%Note{id: 1}, %{body: "body"}))
+  end
+
+  test ".delete(struct, opts \\ []) verifies tenant existence" do
+    assert_raise TenantMissingError, @error_message, fn ->
+      TenantedRepo.delete(%Note{id: 1})
+    end
+    assert_raise TenantMissingError, @error_message, fn ->
+      TenantedRepo.delete(Note.changeset(%Note{id: 1}, %{}))
+    end
+
+    assert {:ok, _} = TenantedRepo.delete(scoped_note)
+    assert {:ok, _} = TenantedRepo.delete(scoped_changeset)
+    assert {:ok, _} = UntenantedRepo.delete(%Note{id: 1})
+  end
+
+  test ".insert!(struct, opts \\ []) verifies tenant existence" do
+    assert_raise TenantMissingError, @error_message, fn ->
+      TenantedRepo.insert!(Note.changeset(%Note{}, %{}))
+    end
+    assert_raise TenantMissingError, @error_message, fn ->
+      TenantedRepo.insert!(%Note{})
+    end
+
+    assert TenantedRepo.insert!(scoped_note)
+    assert TenantedRepo.insert!(Note.changeset(scoped_note, %{}))
+    assert UntenantedRepo.insert!(%Note{})
+  end
+
+  test ".update!(struct, opts \\ []) verifies tenant existence" do
+    assert_raise TenantMissingError, @error_message, fn ->
+      TenantedRepo.update!(Note.changeset(%Note{id: 1}, %{body: "body"}))
+    end
+
+    assert TenantedRepo.update!(scoped_changeset)
+    assert UntenantedRepo.update!(Note.changeset(%Note{id: 1}, %{body: "body"}))
+  end
+
+  test ".insert_or_update!(changeset, opts \\ []) verifies tenant existence" do
+    assert_raise TenantMissingError, @error_message, fn ->
+      TenantedRepo.insert_or_update!(Note.changeset(%Note{id: 1}, %{body: "body"}))
+    end
+
+    assert TenantedRepo.insert_or_update!(scoped_changeset)
+    assert UntenantedRepo.insert_or_update!(Note.changeset(%Note{id: 1}, %{body: "body"}))
+  end
+
+  test ".delete!(struct, opts \\ []) verifies tenant existence" do
+    assert_raise TenantMissingError, @error_message, fn ->
+      TenantedRepo.delete!(%Note{id: 1})
+    end
+    assert_raise TenantMissingError, @error_message, fn ->
+      TenantedRepo.delete!(Note.changeset(%Note{id: 1}, %{}))
+    end
+
+    assert TenantedRepo.delete!(scoped_note)
+    assert TenantedRepo.delete!(scoped_changeset)
+    assert UntenantedRepo.delete!(%Note{id: 1})
+  end
+end
diff --git a/test/support/test_repo.ex b/test/support/test_repo.ex
new file mode 100644
index 0000000..da4eeaa
--- /dev/null
+++ b/test/support/test_repo.ex
@@ -0,0 +1,126 @@
+# NOTE: This is a copy of https://github.com/elixir-ecto/ecto/blob/master/test/support/test_repo.exs
+
+defmodule Ecto.TestAdapter do
+  @behaviour Ecto.Adapter
+
+  alias Ecto.Migration.SchemaMigration
+
+  defmacro __before_compile__(_opts), do: :ok
+
+  def ensure_all_started(_, _) do
+    {:ok, []}
+  end
+
+  def child_spec(_repo, opts) do
+    :apartmentex   = opts[:otp_app]
+    "user"  = opts[:username]
+    "pass"  = opts[:password]
+    "hello" = opts[:database]
+    "local" = opts[:hostname]
+
+    Supervisor.Spec.worker(Task, [fn -> :timer.sleep(:infinity) end])
+  end
+
+  ## Types
+
+  def loaders(:binary_id, type), do: [Ecto.UUID, type]
+  def loaders(_primitive, type), do: [type]
+
+  def dumpers(:binary_id, type), do: [type, Ecto.UUID]
+  def dumpers(_primitive, type), do: [type]
+
+  def autogenerate(:id), do: nil
+  def autogenerate(:embed_id), do: Ecto.UUID.autogenerate
+  def autogenerate(:binary_id), do: Ecto.UUID.autogenerate
+
+  ## Queryable
+
+  def prepare(operation, query), do: {:nocache, {operation, query}}
+
+  def execute(_repo, _, {:nocache, {:all, %{from: {_, SchemaMigration}}}}, _, _, _) do
+    {length(migrated_versions()),
+     Enum.map(migrated_versions(), &List.wrap/1)}
+  end
+
+  def execute(_repo, _, {:nocache, {:all, _}}, _, _, _) do
+    {1, [[1]]}
+  end
+
+  def execute(_repo, _meta, {:nocache, {:delete_all, %{from: {_, SchemaMigration}}}}, [version], _, _) do
+    Process.put(:migrated_versions, List.delete(migrated_versions(), version))
+    {1, nil}
+  end
+
+  def execute(_repo, _meta, {:nocache, {op, %{from: {source, _}}}}, _params, _preprocess, _opts) do
+    send self(), {op, source}
+    {1, nil}
+  end
+
+  ## Schema
+
+  def insert_all(_repo, %{source: {_, source}}, _header, rows, _returning, _opts) do
+    send self(), {:insert_all, source, rows}
+    {1, nil}
+  end
+
+  def insert(_repo, %{source: {nil, "schema_migrations"}}, val, _, _) do
+    version = Keyword.fetch!(val, :version)
+    Process.put(:migrated_versions, [version|migrated_versions()])
+    {:ok, [version: 1]}
+  end
+
+  def insert(_repo, %{context: nil}, _fields, return, _opts),
+    do: send(self(), :insert) && {:ok, Enum.zip(return, 1..length(return))}
+  def insert(_repo, %{context: {:invalid, _}=res}, _fields, _return, _opts),
+    do: res
+
+  # Notice the list of changes is never empty.
+  def update(_repo, %{context: nil}, [_|_], _filters, return, _opts),
+    do: send(self(), :update) && {:ok, Enum.zip(return, 1..length(return))}
+  def update(_repo, %{context: {:invalid, _}=res}, [_|_], _filters, _return, _opts),
+    do: res
+
+  def delete(_repo, _schema_meta, _filter, _opts),
+    do: send(self(), :delete) && {:ok, []}
+
+  ## Transactions
+
+  def transaction(_repo, _opts, fun) do
+    # Makes transactions "trackable" in tests
+    send self(), {:transaction, fun}
+    try do
+      {:ok, fun.()}
+    catch
+      :throw, {:ecto_rollback, value} ->
+        {:error, value}
+    end
+  end
+
+  def rollback(_repo, value) do
+    send self(), {:rollback, value}
+    throw {:ecto_rollback, value}
+  end
+
+  ## Migrations
+
+  def supports_ddl_transaction? do
+    Process.get(:supports_ddl_transaction?) || false
+  end
+
+  def execute_ddl(_repo, command, _) do
+    Process.put(:last_command, command)
+    :ok
+  end
+
+  defp migrated_versions do
+    Process.get(:migrated_versions) || []
+  end
+end
+
+Application.put_env(:ecto, Ecto.TestRepo, [user: "invalid"])
+
+defmodule Ecto.TestRepo do
+  use Ecto.Repo, otp_app: :apartmentex, adapter: Ecto.TestAdapter
+end
+
+Ecto.TestRepo.start_link(url: "ecto://user:pass@local/hello")
diff --git a/test/test_helper.exs b/test/test_helper.exs
index d357cb2..b82963c 100644
--- a/test/test_helper.exs
+++ b/test/test_helper.exs
@@ -8,6 +8,7 @@ Mix.Task.run "ecto.drop", ["quiet", "-r", "Apartmentex.TestPostgresRepo"]
 Mix.Task.run "ecto.create", ["quiet", "-r", "Apartmentex.TestPostgresRepo"]
 
 Apartmentex.TestPostgresRepo.start_link
+Ecto.TestRepo.start_link(url: "ecto://user:pass@local/hello")
 
 ExUnit.start()
 

From 53d5a7e237623c203b45fa1a994ecf7448f5fc41 Mon Sep 17 00:00:00 2001
From: Jeff Deville <jeffdeville@gmail.com>
Date: Thu, 8 Sep 2016 16:57:24 -0400
Subject: [PATCH 2/3] provide an alternative to replacing usages of Repo with
 methods with a different signature, to increase compatibility with other
 libraries

---
 lib/apartmentex.ex                | 45 +++++++++++--------------------
 lib/apartmentex/repo_additions.ex | 18 +++++++++++++
 test/apartmentex_test.exs         | 25 +++++++++++++++++
 3 files changed, 59 insertions(+), 29 deletions(-)
 create mode 100644 lib/apartmentex/repo_additions.ex

diff --git a/lib/apartmentex.ex b/lib/apartmentex.ex
index 517a040..de05a8e 100644
--- a/lib/apartmentex.ex
+++ b/lib/apartmentex.ex
@@ -1,6 +1,5 @@
 defmodule Apartmentex do
-  alias Ecto.Changeset
-  import Apartmentex.PrefixBuilder
+  import Apartmentex.RepoAdditions
 
   defdelegate drop_tenant(repo, tenant), to: Apartmentex.TenantActions
   defdelegate migrate_tenant(repo, tenant, direction \\ :up, opts \\ []), to: Apartmentex.TenantActions
@@ -8,7 +7,7 @@ defmodule Apartmentex do
 
   def all(repo, queryable, tenant, opts \\ []) when is_list(opts) do
     queryable
-    |> add_prefix_to_query(tenant)
+    |> set_tenant(tenant)
     |> repo.all(opts)
   end
 
@@ -17,7 +16,7 @@ defmodule Apartmentex do
   """
   def get(repo, queryable, id, tenant, opts \\ []) do
     queryable
-    |> add_prefix_to_query(tenant)
+    |> set_tenant(tenant)
     |> repo.get(id, opts)
   end
 
@@ -26,19 +25,19 @@ defmodule Apartmentex do
   """
   def get!(repo, queryable, id, tenant, opts \\ []) do
     queryable
-    |> add_prefix_to_query(tenant)
+    |> set_tenant(tenant)
     |> repo.get!(id, opts)
   end
 
   def get_by(repo, queryable, clauses, tenant, opts \\ []) do
     queryable
-    |> add_prefix_to_query(tenant)
+    |> set_tenant(tenant)
     |> repo.get_by(clauses, opts)
   end
 
   def get_by!(repo, queryable, clauses, tenant, opts \\ []) do
     queryable
-    |> add_prefix_to_query(tenant)
+    |> set_tenant(tenant)
     |> repo.get_by!(clauses, opts)
   end
 
@@ -47,7 +46,7 @@ defmodule Apartmentex do
   """
   def one(repo, queryable, tenant, opts \\ []) do
     queryable
-    |> add_prefix_to_query(tenant)
+    |> set_tenant(tenant)
     |> repo.one(opts)
   end
 
@@ -56,7 +55,7 @@ defmodule Apartmentex do
   """
   def one!(repo, queryable, tenant, opts \\ []) do
     queryable
-    |> add_prefix_to_query(tenant)
+    |> set_tenant(tenant)
     |> repo.one!(opts)
   end
 
@@ -67,13 +66,13 @@ defmodule Apartmentex do
   """
   def insert(repo, model_or_changeset, tenant, opts \\ []) do
     model_or_changeset
-    |> add_prefix(tenant)
+    |> set_tenant(tenant)
     |> repo.insert(opts)
   end
 
   def insert!(repo, model_or_changeset, tenant, opts \\ []) do
     model_or_changeset
-    |> add_prefix(tenant)
+    |> set_tenant(tenant)
     |> repo.insert!(opts)
   end
 
@@ -82,13 +81,13 @@ defmodule Apartmentex do
   """
   def update(repo, model_or_changeset, tenant, opts \\ []) do
     model_or_changeset
-    |> add_prefix(tenant)
+    |> set_tenant(tenant)
     |> repo.update(opts)
   end
 
   def update!(repo, model_or_changeset, tenant, opts \\ []) do
     model_or_changeset
-    |> add_prefix(tenant)
+    |> set_tenant(tenant)
     |> repo.update!(opts)
   end
 
@@ -97,7 +96,7 @@ defmodule Apartmentex do
   """
   def update_all(repo, queryable, updates, tenant, opts \\ []) do
     queryable
-    |> add_prefix_to_query(tenant)
+    |> set_tenant(tenant)
     |> repo.update_all(updates, opts)
   end
 
@@ -106,7 +105,7 @@ defmodule Apartmentex do
   """
   def delete(repo, model_or_changeset, tenant, opts \\ []) do
     model_or_changeset
-    |> add_prefix(tenant)
+    |> set_tenant(tenant)
     |> repo.delete(opts)
   end
 
@@ -115,7 +114,7 @@ defmodule Apartmentex do
   """
   def delete!(repo, model_or_changeset, tenant, opts \\ []) do
     model_or_changeset
-    |> add_prefix(tenant)
+    |> set_tenant(tenant)
     |> repo.delete!(opts)
   end
 
@@ -124,23 +123,11 @@ defmodule Apartmentex do
   """
   def delete_all(repo, queryable, tenant, opts \\ []) do
     queryable
-    |> add_prefix_to_query(tenant)
+    |> set_tenant(tenant)
     |> repo.delete_all(opts)
   end
 
   #helpers
 
-  defp add_prefix(%Changeset{} = changeset, tenant) do
-    %{changeset | data: add_prefix(changeset.data, tenant)}
-  end
-
-  defp add_prefix(%{__struct__: _} = model, tenant) do
-    Ecto.put_meta(model,  prefix: build_prefix(tenant))
-  end
 
-  defp add_prefix_to_query(queryable, tenant) do
-    queryable
-    |> Ecto.Queryable.to_query
-    |> Map.put(:prefix, build_prefix(tenant))
-  end
 end
diff --git a/lib/apartmentex/repo_additions.ex b/lib/apartmentex/repo_additions.ex
new file mode 100644
index 0000000..6bab361
--- /dev/null
+++ b/lib/apartmentex/repo_additions.ex
@@ -0,0 +1,18 @@
+defmodule Apartmentex.RepoAdditions do
+  alias Ecto.Changeset
+  import Apartmentex.PrefixBuilder
+
+  def set_tenant(%Changeset{} = changeset, tenant) do
+    %{changeset | data: set_tenant(changeset.data, tenant)}
+  end
+
+  def set_tenant(%{__meta__: _} = model, tenant) do
+    Ecto.put_meta(model,  prefix: build_prefix(tenant))
+  end
+
+  def set_tenant(queryable, tenant) do
+    queryable
+    |> Ecto.Queryable.to_query
+    |> Map.put(:prefix, build_prefix(tenant))
+  end
+end
diff --git a/test/apartmentex_test.exs b/test/apartmentex_test.exs
index 566fa38..88b0a9c 100644
--- a/test/apartmentex_test.exs
+++ b/test/apartmentex_test.exs
@@ -3,6 +3,7 @@ defmodule Apartmentex.ApartmentexTest do
 
   alias Apartmentex.Note
   alias Apartmentex.TestPostgresRepo
+  import Apartmentex.RepoAdditions
 
   @tenant_id 2
 
@@ -159,4 +160,28 @@ defmodule Apartmentex.ApartmentexTest do
     updated_notes = Apartmentex.all(TestPostgresRepo, Note, @tenant_id)
     assert Enum.map(updated_notes, & &1.body) == ["updated", "updated"]
   end
+
+  test ".set_tenant/2 struct adds the tenant prefix" do
+    prefix = %Note{}
+    |> set_tenant(@tenant_id)
+    |> Ecto.get_meta(:prefix)
+    assert prefix == "tenant_2"
+  end
+
+  test ".set_tenant/2 changeset adds the tenant prefix" do
+    prefix = Note.changeset(%Note{}, %{})
+    |> set_tenant(@tenant_id)
+    |> Map.fetch!(:data)
+    |> Ecto.get_meta(:prefix)
+
+    assert prefix == "tenant_2"
+  end
+
+  test ".set_tenant/2 queryable adds the tenant prefix" do
+    prefix = Note
+    |> set_tenant(@tenant_id)
+    |> Map.fetch!(:prefix)
+
+    assert prefix == "tenant_#{@tenant_id}"
+  end
 end

From bad33fe0c6d61585038fc31f7fb626111251ad58 Mon Sep 17 00:00:00 2001
From: Jeff Deville <jeffdeville@gmail.com>
Date: Thu, 8 Sep 2016 17:34:58 -0400
Subject: [PATCH 3/3] Also provide the ability to extract the tenant id from
 the prefix

---
 lib/apartmentex/prefix_builder.ex  |  4 ++++
 test/apartmentex_test.exs          |  9 +++++++--
 test/lib/apartmentex/repo_test.exs | 16 +---------------
 3 files changed, 12 insertions(+), 17 deletions(-)

diff --git a/lib/apartmentex/prefix_builder.ex b/lib/apartmentex/prefix_builder.ex
index 018a034..d00a810 100644
--- a/lib/apartmentex/prefix_builder.ex
+++ b/lib/apartmentex/prefix_builder.ex
@@ -12,4 +12,8 @@ defmodule Apartmentex.PrefixBuilder do
   def build_prefix(tenant) do
     @schema_prefix <> Integer.to_string(tenant.id)
   end
+
+  def extract_tenant(table_prefix) do
+    String.replace_prefix(table_prefix, @schema_prefix, "")
+  end
 end
diff --git a/test/apartmentex_test.exs b/test/apartmentex_test.exs
index 88b0a9c..db15022 100644
--- a/test/apartmentex_test.exs
+++ b/test/apartmentex_test.exs
@@ -165,7 +165,7 @@ defmodule Apartmentex.ApartmentexTest do
     prefix = %Note{}
     |> set_tenant(@tenant_id)
     |> Ecto.get_meta(:prefix)
-    assert prefix == "tenant_2"
+    assert prefix == "tenant_#{@tenant_id}"
   end
 
   test ".set_tenant/2 changeset adds the tenant prefix" do
@@ -174,7 +174,7 @@ defmodule Apartmentex.ApartmentexTest do
     |> Map.fetch!(:data)
     |> Ecto.get_meta(:prefix)
 
-    assert prefix == "tenant_2"
+    assert prefix == "tenant_#{@tenant_id}"
   end
 
   test ".set_tenant/2 queryable adds the tenant prefix" do
@@ -184,4 +184,9 @@ defmodule Apartmentex.ApartmentexTest do
 
     assert prefix == "tenant_#{@tenant_id}"
   end
+
+  test ".extract_tenant/1 removes the prefix from the schema" do
+    assert Apartmentex.PrefixBuilder.extract_tenant("tenant_#{@tenant_id}") == "#{@tenant_id}"
+    assert Apartmentex.PrefixBuilder.extract_tenant("tenant_somestring") == "somestring"
+  end
 end
diff --git a/test/lib/apartmentex/repo_test.exs b/test/lib/apartmentex/repo_test.exs
index 38f674b..8e3d9a6 100644
--- a/test/lib/apartmentex/repo_test.exs
+++ b/test/lib/apartmentex/repo_test.exs
@@ -8,27 +8,13 @@ end
 
 defmodule Apartmentex.ApartmentexTest do
   use ExUnit.Case
-  import Apartmentex.PrefixBuilder
+  import Apartmentex.RepoAdditions
   alias Apartmentex.{Note,TenantMissingError}
   alias Apartmentex.Test.{TenantedRepo,UntenantedRepo}
 
   @tenant_id 2
   @error_message "No tenant specified in Elixir.Apartmentex.Note"
 
-  def set_tenant(%Ecto.Changeset{} = changeset, tenant) do
-    %{changeset | data: set_tenant(changeset.data, tenant)}
-  end
-
-  def set_tenant(%{__meta__: _} = model, tenant) do
-    Ecto.put_meta(model,  prefix: build_prefix(tenant))
-  end
-
-  def set_tenant(queryable, tenant) do
-    queryable
-    |> Ecto.Queryable.to_query
-    |> Map.put(:prefix, build_prefix(tenant))
-  end
-
   def scoped_note_query do
     Note
     |> Ecto.Queryable.to_query