Alter the query instead of the repo

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

lib/apartmentex.ex

@@ -1,14 +1,13 @@
 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
   defdelegate new_tenant(repo, tenant), to: Apartmentex.TenantActions
 
   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

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

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

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

lib/apartmentex/tenant_missing_error.ex

@@ -0,0 +1,3 @@
+defmodule Apartmentex.TenantMissingError do
+  defexception message: "No tenant specified"
+end