And the inevitable missing svk add

[dbsrgits/DBIx-Class.git] / lib / DBIx / Class / Schema.pm

diff --git a/lib/DBIx/Class/Schema.pm b/lib/DBIx/Class/Schema.pm
index d900a27..8385e1d 100644 (file)
--- a/lib/DBIx/Class/Schema.pm
+++ b/lib/DBIx/Class/Schema.pm
@@ -67,10 +67,11 @@ compose_connection to create/modify all the existing database classes.
 =cut
 
 sub register_class {
-  my ($class, $name, $to_register) = @_;
-  my %reg = %{$class->class_registrations};
+  my ($self, $name, $to_register) = @_;
+  my %reg = %{$self->class_registrations};
   $reg{$name} = $to_register;
-  $class->class_registrations(\%reg);
+  $self->class_registrations(\%reg);
+  $to_register->result_source->schema($self);
 }
 
 =head2 registered_classes
@@ -108,7 +109,8 @@ Returns the result source object for the registered name
 
 sub source {
   my ($self, $class) = @_;
-  return $self->class_registrations->{$class}->result_source;
+  return $self->class_registrations->{$class}->result_source
+    if $self->class_registrations->{$class};
 }
 
 =head2 resultset