From a4fdb1c328bc0728d73d5ecf625971d31ca76ef3 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Luis=20Ram=C3=B3n=20L=C3=B3pez?= <lrlopez@gmail.com>
Date: Thu, 2 Jan 2014 23:27:58 +0100
Subject: INSERT queries are now compound primary keys aware

---
 idiorm.php                | 46 ++++++++++++++++++++++++++++++++++++++++++----
 test/QueryBuilderTest.php | 14 ++++++++++++++
 2 files changed, 56 insertions(+), 4 deletions(-)

diff --git a/idiorm.php b/idiorm.php
index 3ea98a7..e4d7cae 100644
--- a/idiorm.php
+++ b/idiorm.php
@@ -776,6 +776,18 @@
             return $this;
         }
 
+        /**
+         * Counts the number of columns that belong to the primary
+         * key and their value is null.
+         */
+        public function count_null_id_columns() {
+            if (is_array($this->_get_id_column_name())) {
+                return count(array_filter($this->id(), 'is_null'));
+            } else {
+                return is_null($this->id()) ? 1 : 0;
+            }
+        }
+
         /**
          * Add a column to the list of columns returned by the SELECT
          * query. This defaults to '*'. The second optional argument is
@@ -1607,12 +1619,27 @@
          * (table names, column names etc). This method can
          * also deal with dot-separated identifiers eg table.column
          */
-        protected function _quote_identifier($identifier) {
+        protected function _quote_one_identifier($identifier) {
             $parts = explode('.', $identifier);
             $parts = array_map(array($this, '_quote_identifier_part'), $parts);
             return join('.', $parts);
         }
 
+        /**
+         * Quote a string that is used as an identifier
+         * (table names, column names etc) or an array containing
+         * multiple identifiers. This method can also deal with
+         * dot-separated identifiers eg table.column
+         */
+        protected function _quote_identifier($identifier) {
+            if (is_array($identifier)) {
+                $result = array_map(array($this, '_quote_one_identifier'), $identifier);
+                return join(', ', $result);
+            } else {
+                return $this->_quote_one_identifier($identifier);
+            }
+        }
+
         /**
          * This method performs the actual quoting of a single
          * part of an identifier, using the identifier quote
@@ -1853,12 +1880,23 @@
             // If we've just inserted a new record, set the ID of this object
             if ($this->_is_new) {
                 $this->_is_new = false;
-                if (is_null($this->id())) {
+                if ($this->count_null_id_columns() != 0) {
                     $db = self::get_db($this->_connection_name);
                     if($db->getAttribute(PDO::ATTR_DRIVER_NAME) == 'pgsql') {
-                        $this->_data[$this->_get_id_column_name()] = self::get_last_statement()->fetchColumn();
+                        // it may return several columns if a compound primary
+                        // key is used
+                        $row = self::get_last_statement()->fetch(PDO::FETCH_ASSOC);
+                        foreach($row as $key => $value) {
+                            $this->_data[$key] = $value;
+                        }
                     } else {
-                        $this->_data[$this->_get_id_column_name()] = $db->lastInsertId();
+                        $column = $this->_get_id_column_name();
+                        // if the primary key is compound, assign the last inserted id
+                        // to the first column
+                        if (is_array($column)) {
+                            $column = array_slice($column, 0, 1);
+                        }
+                        $this->_data[$column] = $db->lastInsertId();
                     }
                 }
             }
diff --git a/test/QueryBuilderTest.php b/test/QueryBuilderTest.php
index 2f60c59..cd3a168 100644
--- a/test/QueryBuilderTest.php
+++ b/test/QueryBuilderTest.php
@@ -514,6 +514,20 @@ class QueryBuilderTest extends PHPUnit_Framework_TestCase {
         $this->assertEquals($expected, ORM::get_last_query());
     }
 
+    public function test_quote_identifier_part() {
+        $widget = ORM::for_table('widget')->find_one(1);
+        $widget->set('added', '2013-01-04');
+        $widget->save();
+        $expected = "UPDATE `widget` SET `added` = '2013-01-04' WHERE `id` = '1'";
+        $this->assertEquals($expected, ORM::get_last_query());
+    }
+    
+    public function test_quote_multiple_identifiers_part() {
+        $record = ORM::for_table('widget')->use_id_column(array('id1', 'id2'))->create();
+        $expected = "`id1`, `id2`";
+        $this->assertEquals($expected, $record->_quote_identifier($record->_get_id_column_name()));
+    }
+
     /**
      * Regression tests
      */
-- 
cgit v1.2.3