Selection::insert(): refactored record fetch after insert

Author: Unlink

src/Database/Table/Selection.php

@@ -36,9 +36,12 @@ class Selection extends Nette\Object implements \Iterator, IRowContainer, \Array
 	/** @var string|array|NULL primary key field name */
 	protected $primary;
 
-	/** @var string|bool primary column sequence name, FALSE for autodetection */
+	/** @var string|bool|NULL primary column sequence name, FALSE for autodetection */
 	protected $primarySequence = FALSE;
 
+	/** @var string|bool|NULL primary key field name with autoincrement */
+	protected $primarySequenceField = FALSE;
+
 	/** @var IRow[] data read from database in [primary key => IRow] format */
 	protected $rows;
 
@@ -134,7 +137,7 @@ public function getPrimary($need = TRUE)
 
 
 	/**
-	 * @return string
+	 * @return string|NULL
 	 */
 	public function getPrimarySequence()
 	{
@@ -156,6 +159,23 @@ public function setPrimarySequence($sequence)
 		return $this;
 	}
 
+	/**
+	 * @return string|NULL
+	 */
+	public function getPrimarySequenceField()
+	{
+		if ($this->primarySequenceField === FALSE) {
+			$this->primarySequenceField = NULL;
+			foreach ($this->context->getStructure()->getColumns($this->name) as $columnMeta) {
+				if ($columnMeta['autoincrement']) {
+					$this->primarySequenceField = $columnMeta['name'];
+					break;
+				}
+			}
+		}
+
+		return $this->primarySequenceField;
+	}
 
 	/**
 	 * @return string
@@ -745,32 +765,35 @@ public function insert($data)
 			return $return->getRowCount();
 		}
 
-		$primarySequenceName = $this->getPrimarySequence();
-		$primaryKey = $this->context->getInsertId(
-			!empty($primarySequenceName)
-				? $this->context->getConnection()->getSupplementalDriver()->delimite($primarySequenceName)
-				: $primarySequenceName
-		);
-		if ($primaryKey === FALSE) {
-			unset($this->refCache['referencing'][$this->getGeneralCacheKey()][$this->getSpecificCacheKey()]);
-			return $return->getRowCount();
+		if (Nette\Utils\Arrays::isList($data)) {
+			return $data;
 		}
 
-		if (is_array($this->getPrimary())) {
-			$primaryKey = [];
-
-			foreach ((array) $this->getPrimary() as $key) {
-				if (!isset($data[$key])) {
-					return $data;
-				}
-
+		$primaryKey = [];
+		$primarySequenceField = $this->getPrimarySequenceField();
+		foreach ((array) $this->getPrimary() as $key) {
+			if (isset($data[$key])) {
 				$primaryKey[$key] = $data[$key];
 			}
-			if (count($primaryKey) === 1) {
-				$primaryKey = reset($primaryKey);
+		}
+
+		if ($primarySequenceField !== NULL) {
+			$primarySequenceName = $this->getPrimarySequence();
+			$primaryKey[$primarySequenceField] = $this->context->getInsertId(
+				!empty($primarySequenceName)
+					? $this->context->getConnection()->getSupplementalDriver()->delimite($primarySequenceName)
+					: $primarySequenceName
+			);
+			if ($primaryKey[$primarySequenceField] === FALSE) {
+				unset($this->refCache['referencing'][$this->getGeneralCacheKey()][$this->getSpecificCacheKey()]);
+				return $return->getRowCount();
 			}
 		}
 
+		if (count($primaryKey) === 1) {
+			$primaryKey = reset($primaryKey);
+		}
+
 		$row = $this->createSelectionInstance()
 			->select('*')
 			->wherePrimary($primaryKey)