BatchQueryResult.php 7.0 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256
  1. <?php
  2. /**
  3. * @link http://www.yiiframework.com/
  4. * @copyright Copyright (c) 2008 Yii Software LLC
  5. * @license http://www.yiiframework.com/license/
  6. */
  7. namespace yii\db;
  8. use yii\base\Component;
  9. /**
  10. * BatchQueryResult represents a batch query from which you can retrieve data in batches.
  11. *
  12. * You usually do not instantiate BatchQueryResult directly. Instead, you obtain it by
  13. * calling [[Query::batch()]] or [[Query::each()]]. Because BatchQueryResult implements the [[\Iterator]] interface,
  14. * you can iterate it to obtain a batch of data in each iteration. For example,
  15. *
  16. * ```php
  17. * $query = (new Query)->from('user');
  18. * foreach ($query->batch() as $i => $users) {
  19. * // $users represents the rows in the $i-th batch
  20. * }
  21. * foreach ($query->each() as $user) {
  22. * }
  23. * ```
  24. *
  25. * @author Qiang Xue <qiang.xue@gmail.com>
  26. * @since 2.0
  27. */
  28. class BatchQueryResult extends Component implements \Iterator
  29. {
  30. /**
  31. * @event Event an event that is triggered when the batch query is reset.
  32. * @see reset()
  33. * @since 2.0.41
  34. */
  35. const EVENT_RESET = 'reset';
  36. /**
  37. * @event Event an event that is triggered when the last batch has been fetched.
  38. * @since 2.0.41
  39. */
  40. const EVENT_FINISH = 'finish';
  41. /**
  42. * @var Connection the DB connection to be used when performing batch query.
  43. * If null, the "db" application component will be used.
  44. */
  45. public $db;
  46. /**
  47. * @var Query the query object associated with this batch query.
  48. * Do not modify this property directly unless after [[reset()]] is called explicitly.
  49. */
  50. public $query;
  51. /**
  52. * @var int the number of rows to be returned in each batch.
  53. */
  54. public $batchSize = 100;
  55. /**
  56. * @var bool whether to return a single row during each iteration.
  57. * If false, a whole batch of rows will be returned in each iteration.
  58. */
  59. public $each = false;
  60. /**
  61. * @var DataReader the data reader associated with this batch query.
  62. */
  63. private $_dataReader;
  64. /**
  65. * @var array the data retrieved in the current batch
  66. */
  67. private $_batch;
  68. /**
  69. * @var mixed the value for the current iteration
  70. */
  71. private $_value;
  72. /**
  73. * @var string|int the key for the current iteration
  74. */
  75. private $_key;
  76. /**
  77. * @var int MSSQL error code for exception that is thrown when last batch is size less than specified batch size
  78. * @see https://github.com/yiisoft/yii2/issues/10023
  79. */
  80. private $mssqlNoMoreRowsErrorCode = -13;
  81. /**
  82. * Destructor.
  83. */
  84. public function __destruct()
  85. {
  86. // make sure cursor is closed
  87. $this->reset();
  88. }
  89. /**
  90. * Resets the batch query.
  91. * This method will clean up the existing batch query so that a new batch query can be performed.
  92. */
  93. public function reset()
  94. {
  95. if ($this->_dataReader !== null) {
  96. $this->_dataReader->close();
  97. }
  98. $this->_dataReader = null;
  99. $this->_batch = null;
  100. $this->_value = null;
  101. $this->_key = null;
  102. $this->trigger(self::EVENT_RESET);
  103. }
  104. /**
  105. * Resets the iterator to the initial state.
  106. * This method is required by the interface [[\Iterator]].
  107. */
  108. public function rewind()
  109. {
  110. $this->reset();
  111. $this->next();
  112. }
  113. /**
  114. * Moves the internal pointer to the next dataset.
  115. * This method is required by the interface [[\Iterator]].
  116. */
  117. public function next()
  118. {
  119. if ($this->_batch === null || !$this->each || $this->each && next($this->_batch) === false) {
  120. $this->_batch = $this->fetchData();
  121. reset($this->_batch);
  122. }
  123. if ($this->each) {
  124. $this->_value = current($this->_batch);
  125. if ($this->query->indexBy !== null) {
  126. $this->_key = key($this->_batch);
  127. } elseif (key($this->_batch) !== null) {
  128. $this->_key = $this->_key === null ? 0 : $this->_key + 1;
  129. } else {
  130. $this->_key = null;
  131. }
  132. } else {
  133. $this->_value = $this->_batch;
  134. $this->_key = $this->_key === null ? 0 : $this->_key + 1;
  135. }
  136. }
  137. /**
  138. * Fetches the next batch of data.
  139. * @return array the data fetched
  140. * @throws Exception
  141. */
  142. protected function fetchData()
  143. {
  144. if ($this->_dataReader === null) {
  145. $this->_dataReader = $this->query->createCommand($this->db)->query();
  146. }
  147. $rows = $this->getRows();
  148. return $this->query->populate($rows);
  149. }
  150. /**
  151. * Reads and collects rows for batch
  152. * @return array
  153. * @since 2.0.23
  154. */
  155. protected function getRows()
  156. {
  157. $rows = [];
  158. $count = 0;
  159. try {
  160. while ($count++ < $this->batchSize) {
  161. if ($row = $this->_dataReader->read()) {
  162. $rows[] = $row;
  163. } else {
  164. // we've reached the end
  165. $this->trigger(self::EVENT_FINISH);
  166. break;
  167. }
  168. }
  169. } catch (\PDOException $e) {
  170. $errorCode = isset($e->errorInfo[1]) ? $e->errorInfo[1] : null;
  171. if ($this->getDbDriverName() !== 'sqlsrv' || $errorCode !== $this->mssqlNoMoreRowsErrorCode) {
  172. throw $e;
  173. }
  174. }
  175. return $rows;
  176. }
  177. /**
  178. * Returns the index of the current dataset.
  179. * This method is required by the interface [[\Iterator]].
  180. * @return int the index of the current row.
  181. */
  182. public function key()
  183. {
  184. return $this->_key;
  185. }
  186. /**
  187. * Returns the current dataset.
  188. * This method is required by the interface [[\Iterator]].
  189. * @return mixed the current dataset.
  190. */
  191. public function current()
  192. {
  193. return $this->_value;
  194. }
  195. /**
  196. * Returns whether there is a valid dataset at the current position.
  197. * This method is required by the interface [[\Iterator]].
  198. * @return bool whether there is a valid dataset at the current position.
  199. */
  200. public function valid()
  201. {
  202. return !empty($this->_batch);
  203. }
  204. /**
  205. * Gets db driver name from the db connection that is passed to the `batch()`, if it is not passed it uses
  206. * connection from the active record model
  207. * @return string|null
  208. */
  209. private function getDbDriverName()
  210. {
  211. if (isset($this->db->driverName)) {
  212. return $this->db->driverName;
  213. }
  214. if (!empty($this->_batch)) {
  215. $key = array_keys($this->_batch)[0];
  216. if (isset($this->_batch[$key]->db->driverName)) {
  217. return $this->_batch[$key]->db->driverName;
  218. }
  219. }
  220. return null;
  221. }
  222. /**
  223. * Unserialization is disabled to prevent remote code execution in case application
  224. * calls unserialize() on user input containing specially crafted string.
  225. * @see CVE-2020-15148
  226. * @since 2.0.38
  227. */
  228. public function __wakeup()
  229. {
  230. throw new \BadMethodCallException('Cannot unserialize ' . __CLASS__);
  231. }
  232. }