Fixing MultiFormObjectDecorator to work with SS 3.1

code/extensions/MultiFormObjectDecorator.php

@@ -23,16 +23,27 @@ class MultiFormObjectDecorator extends DataExtension {
 	private static $has_one = array(
 		'MultiFormSession' => 'MultiFormSession',
 	);
-	
+
+	/**
+	 * Augment any queries to MultiFormObjectDecorator and only
+	 * return anything that isn't considered temporary.
+	 */
 	public function augmentSQL(SQLQuery &$query) {
-		// If you're querying by ID, ignore the sub-site - this is a bit ugly...
+		$where = $query->getWhere();
+		if(!$where && !$this->wantsTemporary($query)) {
+			$from = array_values($query->getFrom());
+			$query->addWhere("{$from[0]}.MultiFormIsTemporary = 0");
+			return;
+		}
+
 		if(
-			strpos($query->where[0], ".`ID` = ") === false 
-			&& strpos($query->where[0], ".ID = ") === false 
-			&& strpos($query->where[0], "ID = ") !== 0
+			strpos($where[0], ".`ID` = ") === false 
+			&& strpos($where[0], ".ID = ") === false 
+			&& strpos($where[0], "ID = ") !== 0
 			&& !$this->wantsTemporary($query)
 		) {
-			$query->where[] = "\"{$query->from[0]}\".\"MultiFormIsTemporary\" = 0"; 
+			$from = array_values($query->getFrom());
+			$query->addWhere("{$from[0]}.MultiFormIsTemporary = 0");
 		}
 	}
 	
@@ -45,8 +56,10 @@ class MultiFormObjectDecorator extends DataExtension {
 	 * @return boolean
 	 */
 	protected function wantsTemporary($query) {
-		foreach($query->where as $whereClause) {
-			if($whereClause == "\"{$query->from[0]}\".\"MultiFormIsTemporary\" = 1") {
+		foreach($query->getWhere() as $whereClause) {
+			$from = array_values($query->getFrom());
+			// SQLQuery will automatically add double quotes and single quotes to values, so check against that.
+			if($whereClause == "{$from[0]}.\"MultiFormIsTemporary\" = '1'") {
 				return true;
 			}
 		}

tests/MultiFormObjectDecoratorTest.php

@@ -0,0 +1,44 @@
+<?php
+class MultiFormObjectDecoratorTest extends SapphireTest {
+
+	protected static $fixture_file = 'MultiFormObjectDecoratorTest.yml';
+
+	protected $requiredExtensions = array(
+		'MultiFormObjectDecorator_DataObject' => array('MultiFormObjectDecorator')
+	);
+
+	protected $extraDataObjects = array(
+		'MultiFormObjectDecorator_DataObject'
+	);
+
+	public function testTemporaryDataFilteredQuery() {
+		$records = MultiFormObjectDecorator_DataObject::get()
+			->map('Name')
+			->toArray();
+
+		$this->assertContains('Test 1', $records);
+		$this->assertContains('Test 2', $records);
+		$this->assertNotContains('Test 3', $records);
+
+	}
+
+	public function testTemporaryDataQuery() {
+		$records = MultiFormObjectDecorator_DataObject::get()
+			->filter(array('MultiFormIsTemporary' => 1))
+			->map('Name')
+			->toArray();
+
+		$this->assertNotContains('Test 1', $records);
+		$this->assertNotContains('Test 2', $records);
+		$this->assertContains('Test 3', $records);
+	}
+
+}
+
+class MultiFormObjectDecorator_DataObject extends DataObject {
+
+	private static $db = array(
+		'Name' => 'Varchar'
+	);
+
+}

tests/MultiFormObjectDecoratorTest.yml

@@ -0,0 +1,10 @@
+MultiFormObjectDecorator_DataObject:
+  test-data-1:
+    Name: Test 1
+    MultiFormIsTemporary: 0
+  test-data-2:
+    Name: Test 2
+    MultiFormIsTemporary: 0
+  test-data-3:
+    Name: Test 3
+    MultiFormIsTemporary: 1

tests/MultiFormTest.php

@@ -79,7 +79,7 @@ class MultiFormTest extends FunctionalTest {
 		// A new session is generated, even though we made up the identifier
 		$this->assertInstanceOf('MultiFormSession', $this->form->session);
 	}
-	
+
 }
 class MultiFormTest_Controller extends Controller implements TestOnly {