remove_columns() sanitization by Oleg Pronin

diff --git a/lib/DBIx/Class/ResultSource.pm b/lib/DBIx/Class/ResultSource.pm
index fd7da0f..4f62cba 100644 (file)
--- a/lib/DBIx/Class/ResultSource.pm
+++ b/lib/DBIx/Class/ResultSource.pm
@@ -318,22 +318,18 @@ broken result source.
 =cut
 
 sub remove_columns {
-  my ($self, @cols) = @_;
-
-  return unless $self->_ordered_columns;
-
-  my $columns = $self->_columns;
-  my @remaining;
+  my ($self, @to_remove) = @_;
 
-  foreach my $col (@{$self->_ordered_columns}) {
-    push @remaining, $col unless grep(/$col/, @cols);
-  }
+  my $columns = $self->_columns
+    or return;
 
-  foreach (@cols) {
+  my %to_remove;
+  for (@to_remove) {
     delete $columns->{$_};
-  };
+    ++$to_remove{$_};
+  }
 
-  $self->_ordered_columns(\@remaining);
+  $self->_ordered_columns([ grep { not $to_remove{$_} } @{$self->_ordered_columns} ]);
 }
 
 sub remove_column { shift->remove_columns(@_); } # DO NOT CHANGE THIS TO GLOB

diff --git a/t/60core.t b/t/60core.t
index a46be69..d41b551 100644 (file)
--- a/t/60core.t
+++ b/t/60core.t
@@ -7,7 +7,7 @@ use DBICTest;
 
 my $schema = DBICTest->init_schema();
 
-plan tests => 84;
+plan tests => 86;
 
 eval { require DateTime::Format::MySQL };
 my $NO_DTFM = $@ ? 1 : 0;
@@ -335,10 +335,29 @@ ok(!$@, "stringify to false value doesn't cause error");
 
 # test remove_columns
 {
-  is_deeply([$schema->source('CD')->columns], [qw/cdid artist title year genreid single_track/]);
-  $schema->source('CD')->remove_columns('year');
-  is_deeply([$schema->source('CD')->columns], [qw/cdid artist title genreid single_track/]);
-  ok(! exists $schema->source('CD')->_columns->{'year'}, 'year still exists in _columns');
+  is_deeply(
+    [$schema->source('CD')->columns],
+    [qw/cdid artist title year genreid single_track/],
+    'initial columns',
+  );
+
+  $schema->source('CD')->remove_columns('coolyear'); #should not delete year
+  is_deeply(
+    [$schema->source('CD')->columns],
+    [qw/cdid artist title year genreid single_track/],
+    'nothing removed when removing a non-existent column',
+  );
+
+  $schema->source('CD')->remove_columns('genreid', 'year');
+  is_deeply(
+    [$schema->source('CD')->columns],
+    [qw/cdid artist title single_track/],
+    'removed two columns',
+  );
+
+  my $priv_columns = $schema->source('CD')->_columns;
+  ok(! exists $priv_columns->{'year'}, 'year purged from _columns');
+  ok(! exists $priv_columns->{'genreid'}, 'genreid purged from _columns');
 }
 
 # test get_inflated_columns with objects