diff --git a/lib/igniter.ex b/lib/igniter.ex
index 662418e..63fe886 100644
--- a/lib/igniter.ex
+++ b/lib/igniter.ex
@@ -25,8 +25,8 @@ defmodule AshPostgres.Igniter do
       {igniter, {:ok, value}} when is_binary(value) or is_nil(value) ->
         {:ok, igniter, value}
 
-      _ ->
-        :error
+      {igniter, _} ->
+        {:error, igniter}
     end
   end
 
@@ -37,8 +37,8 @@ defmodule AshPostgres.Igniter do
       {igniter, {:ok, value}} when is_atom(value) ->
         {:ok, igniter, value}
 
-      _ ->
-        :error
+      {igniter, _} ->
+        {:error, igniter}
     end
   end
 
diff --git a/lib/mix/tasks/ash_postgres.gen.resources.ex b/lib/mix/tasks/ash_postgres.gen.resources.ex
index 7360542..346bbda 100644
--- a/lib/mix/tasks/ash_postgres.gen.resources.ex
+++ b/lib/mix/tasks/ash_postgres.gen.resources.ex
@@ -5,7 +5,7 @@ defmodule Mix.Tasks.AshPostgres.Gen.Resources do
 
   @shortdoc "Generates or updates resources based on a database schema"
 
-  @doc """
+  @moduledoc """
   #{@shortdoc}
 
   ## Example
diff --git a/lib/resource_generator/resource_generator.ex b/lib/resource_generator/resource_generator.ex
index 2b4b395..81ef6af 100644
--- a/lib/resource_generator/resource_generator.ex
+++ b/lib/resource_generator/resource_generator.ex
@@ -1,4 +1,5 @@
 defmodule AshPostgres.ResourceGenerator do
+  @moduledoc false
   alias AshPostgres.ResourceGenerator.Spec
 
   require Logger
@@ -6,26 +7,9 @@ defmodule AshPostgres.ResourceGenerator do
   def generate(igniter, repos, domain, opts \\ []) do
     {igniter, resources} = Ash.Resource.Igniter.list_resources(igniter)
 
-    resources =
-      Task.async_stream(resources, fn resource ->
-        {resource, AshPostgres.Igniter.repo(igniter, resource),
-         AshPostgres.Igniter.table(igniter, resource)}
-      end)
-      |> Enum.map(fn {:ok, {resource, repo, table}} ->
-        repo =
-          case repo do
-            {:ok, _igniter, repo} -> repo
-            _ -> nil
-          end
-
-        table =
-          case table do
-            {:ok, _igniter, table} -> table
-            _ -> nil
-          end
-
-        {resource, repo, table}
-      end)
+    # This is a hack. We should be looking at compiled resources
+    # unlikely to ever matter given how this task will be used though.
+    resources = Enum.filter(resources, &Code.ensure_loaded?/1)
 
     igniter = Igniter.include_all_elixir_files(igniter)
 
@@ -140,8 +124,6 @@ defmodule AshPostgres.ResourceGenerator do
   end
 
   defp check_constraints(%{check_constraints: check_constraints}, _) do
-    IO.inspect(check_constraints)
-
     check_constraints =
       Enum.map_join(check_constraints, "\n", fn check_constraint ->
         """
@@ -158,9 +140,8 @@ defmodule AshPostgres.ResourceGenerator do
 
   defp skip_unique_indexes(%{indexes: indexes}) do
     indexes
-    |> Enum.filter(& &1.unique?)
-    |> Enum.filter(fn %{columns: columns} ->
-      Enum.all?(columns, &Regex.match?(~r/^[0-9a-zA-Z_]+$/, &1))
+    |> Enum.filter(fn %{unique?: unique?, columns: columns} ->
+      unique? && Enum.all?(columns, &Regex.match?(~r/^[0-9a-zA-Z_]+$/, &1))
     end)
     |> Enum.reject(&index_as_identity?/1)
     |> case do
@@ -176,9 +157,8 @@ defmodule AshPostgres.ResourceGenerator do
 
   defp identity_index_names(%{indexes: indexes}) do
     indexes
-    |> Enum.filter(& &1.unique?)
-    |> Enum.filter(fn %{columns: columns} ->
-      Enum.all?(columns, &Regex.match?(~r/^[0-9a-zA-Z_]+$/, &1))
+    |> Enum.filter(fn %{unique?: unique?, columns: columns} ->
+      unique? && Enum.all?(columns, &Regex.match?(~r/^[0-9a-zA-Z_]+$/, &1))
     end)
     |> case do
       [] ->
@@ -195,9 +175,8 @@ defmodule AshPostgres.ResourceGenerator do
 
   defp add_identities(str, %{indexes: indexes}) do
     indexes
-    |> Enum.filter(& &1.unique?)
-    |> Enum.filter(fn %{columns: columns} ->
-      Enum.all?(columns, &Regex.match?(~r/^[0-9a-zA-Z_]+$/, &1))
+    |> Enum.filter(fn %{unique?: unique?, columns: columns} ->
+      unique? && Enum.all?(columns, &Regex.match?(~r/^[0-9a-zA-Z_]+$/, &1))
     end)
     |> Enum.map(fn index ->
       name = index.name
@@ -431,10 +410,8 @@ defmodule AshPostgres.ResourceGenerator do
   defp custom_indexes(table_spec, true) do
     table_spec.indexes
     |> Enum.reject(fn index ->
-      !index.unique? || (&index_as_identity?/1)
-    end)
-    |> Enum.reject(fn index ->
-      Enum.any?(index.columns, &String.contains?(&1, "("))
+      !index.unique? || (&index_as_identity?/1) ||
+        Enum.any?(index.columns, &String.contains?(&1, "("))
     end)
     |> case do
       [] ->
diff --git a/lib/resource_generator/sensitive_data.ex b/lib/resource_generator/sensitive_data.ex
index 10fa369..77b0e45 100644
--- a/lib/resource_generator/sensitive_data.ex
+++ b/lib/resource_generator/sensitive_data.ex
@@ -1,4 +1,5 @@
 defmodule AshPostgres.ResourceGenerator.SensitiveData do
+  @moduledoc false
   # I got this from ChatGPT, but this is a best effort transformation
   # anyway.
   @sensitive_patterns [
@@ -65,7 +66,7 @@ defmodule AshPostgres.ResourceGenerator.SensitiveData do
     ~r/.*_token/i
   ]
 
-  def is_sensitive?(column_name) do
+  def sensitive?(column_name) do
     Enum.any?(@sensitive_patterns, fn pattern ->
       Regex.match?(pattern, column_name)
     end)
diff --git a/lib/resource_generator/spec.ex b/lib/resource_generator/spec.ex
index 4a20c1b..fe9fa24 100644
--- a/lib/resource_generator/spec.ex
+++ b/lib/resource_generator/spec.ex
@@ -1,4 +1,5 @@
 defmodule AshPostgres.ResourceGenerator.Spec do
+  @moduledoc false
   require Logger
 
   defstruct [
@@ -15,6 +16,7 @@ defmodule AshPostgres.ResourceGenerator.Spec do
   ]
 
   defmodule Attribute do
+    @moduledoc false
     defstruct [
       :name,
       :type,
@@ -31,6 +33,7 @@ defmodule AshPostgres.ResourceGenerator.Spec do
   end
 
   defmodule ForeignKey do
+    @moduledoc false
     defstruct [
       :constraint_name,
       :match_type,
@@ -44,14 +47,17 @@ defmodule AshPostgres.ResourceGenerator.Spec do
   end
 
   defmodule Index do
+    @moduledoc false
     defstruct [:name, :columns, :unique?, :nulls_distinct, :where_clause, :using, :include]
   end
 
   defmodule CheckConstraint do
+    @moduledoc false
     defstruct [:name, :column, :expression]
   end
 
   defmodule Relationship do
+    @moduledoc false
     defstruct [
       :name,
       :type,
@@ -71,7 +77,6 @@ defmodule AshPostgres.ResourceGenerator.Spec do
       Ecto.Migrator.with_repo(repo, fn repo ->
         repo
         |> table_specs(opts)
-        |> verify_found_tables(opts)
         |> Enum.group_by(&Enum.take(&1, 2), fn [_, _, field, type, default, size, allow_nil?] ->
           name = Macro.underscore(field)
 
@@ -421,39 +426,6 @@ defmodule AshPostgres.ResourceGenerator.Spec do
     raise "Unexpected character: #{inspect(other)} at #{inspect(stack)} with #{inspect(field)} - #{inspect(acc)}"
   end
 
-  defp verify_found_tables(specs, opts) do
-    if opts[:tables] do
-      not_found =
-        Enum.reject(opts[:tables], fn table ->
-          if String.ends_with?(table, ".") do
-            true
-          else
-            case String.split(table, ".") do
-              [schema, table] ->
-                Enum.any?(specs, fn spec ->
-                  Enum.at(spec, 0) == table and Enum.at(spec, 1) == schema
-                end)
-
-              [table] ->
-                Enum.any?(specs, fn spec ->
-                  Enum.at(spec, 0) == table
-                end)
-            end
-          end
-        end)
-
-      case not_found do
-        [] ->
-          specs
-
-        tables ->
-          raise "The following tables did not exist: #{inspect(tables)}"
-      end
-    else
-      specs
-    end
-  end
-
   defp build_attributes(attributes, table_name, schema, repo, opts) do
     attributes
     |> set_primary_key(table_name, schema, repo)
@@ -546,9 +518,9 @@ defmodule AshPostgres.ResourceGenerator.Spec do
     |> Enum.flat_map(fn {repo, specs} ->
       do_add_relationships(
         specs,
-        Enum.flat_map(resources, fn {resource, resource_repo, table} ->
-          if resource_repo == repo do
-            [{resource, table}]
+        Enum.flat_map(resources, fn resource ->
+          if AshPostgres.DataLayer.Info.repo(resource) == repo do
+            [{resource, AshPostgres.DataLayer.Info.table(resource)}]
           else
             []
           end
@@ -918,7 +890,7 @@ defmodule AshPostgres.ResourceGenerator.Spec do
     Enum.map(attributes, fn attribute ->
       %{
         attribute
-        | sensitive?: AshPostgres.ResourceGenerator.SensitiveData.is_sensitive?(attribute.name)
+        | sensitive?: AshPostgres.ResourceGenerator.SensitiveData.sensitive?(attribute.name)
       }
     end)
   end
diff --git a/test/resource_generator_test.exs b/test/resource_generator_test.exs
index f501d26..50bfa7d 100644
--- a/test/resource_generator_test.exs
+++ b/test/resource_generator_test.exs
@@ -3,6 +3,7 @@ defmodule AshPostgres.RelWithParentFilterTest do
 
   setup do
     AshPostgres.TestRepo.query!("DROP TABLE IF EXISTS example_table")
+
     AshPostgres.TestRepo.query!("CREATE TABLE example_table (
       id UUID DEFAULT uuid_generate_v4() PRIMARY KEY,
       name VARCHAR(255),
@@ -10,18 +11,45 @@ defmodule AshPostgres.RelWithParentFilterTest do
       email VARCHAR(255)
     )")
 
-    on_exit(fn ->
-      AshPostgres.TestRepo.query!("DROP TABLE IF EXISTS example_table")
-    end)
+    :ok
   end
 
   test "a resource is generated from a table" do
-    Igniter.new()
-    |> Igniter.compose_task("ash_postgres.gen.resources", [
-      "MyApp.Accounts",
-      "--tables",
-      "example_table",
-      "--yes"
-    ])
+    resource =
+      Igniter.new()
+      |> Igniter.compose_task("ash_postgres.gen.resources", [
+        "MyApp.Accounts",
+        "--tables",
+        "example_table",
+        "--yes"
+      ])
+      |> Igniter.prepare_for_write()
+      |> Map.get(:rewrite)
+      |> Rewrite.source!("lib/my_app/accounts/example_table.ex")
+      |> Rewrite.Source.get(:content)
+
+    assert String.trim(resource) ==
+             String.trim("""
+             defmodule MyApp.Accounts.ExampleTable do
+               use Ash.Resource,
+                 domain: MyApp.Accounts,
+                 data_layer: AshPostgres.DataLayer
+
+               postgres do
+                 table "example_table"
+                 repo AshPostgres.TestRepo
+               end
+
+               attributes do
+                 uuid_primary_key(:id)
+                 attribute(:name, :string)
+                 attribute(:age, :integer)
+
+                 attribute :email, :string do
+                   sensitive?(true)
+                 end
+               end
+             end
+             """)
   end
 end
diff --git a/test/support/resources/post.ex b/test/support/resources/post.ex
index e4f64c2..fabbdc6 100644
--- a/test/support/resources/post.ex
+++ b/test/support/resources/post.ex
@@ -50,6 +50,7 @@ defmodule HasNoComments do
 end
 
 defmodule CiCategory do
+  @moduledoc false
   use Ash.Type.NewType,
     subtype_of: :ci_string
 end