<?php<liu21st@gmail.com>declare (strict_types = 1);
namespace think\db\connector;
use Closure;
use MongoDB\BSON\ObjectID;
use MongoDB\Driver\BulkWrite;
use MongoDB\Driver\Command;
use MongoDB\Driver\Cursor;
use MongoDB\Driver\Exception\AuthenticationException;
use MongoDB\Driver\Exception\BulkWriteException;
use MongoDB\Driver\Exception\ConnectionException;
use MongoDB\Driver\Exception\InvalidArgumentException;
use MongoDB\Driver\Exception\RuntimeException;
use MongoDB\Driver\Manager;
use MongoDB\Driver\Query as MongoQuery;
use MongoDB\Driver\ReadPreference;
use MongoDB\Driver\WriteConcern;
use think\db\BaseQuery;
use think\db\builder\Mongo as Builder;
use think\db\Connection;
use think\db\exception\DbException as Exception;
use think\db\Mongo as Query;
use function implode;
use function is_array;
class Mongo extends Connection
{
protected $dbName = '';
protected $typeMap = 'array';
protected $mongo; protected $cursor; protected $session_uuid; protected $sessions = [];
protected $builder;
protected $config = [
'type' => '',
'hostname' => '',
'database' => '',
'is_replica_set' => false,
'username' => '',
'password' => '',
'hostport' => '',
'dsn' => '',
'params' => [],
'charset' => 'utf8',
'pk' => '_id',
'pk_type' => 'ObjectID',
'prefix' => '',
'deploy' => 0,
'rw_separate' => false,
'master_num' => 1,
'slave_no' => '',
'fields_strict' => true,
'fields_cache' => false,
'trigger_sql' => true,
'auto_timestamp' => false,
'datetime_format' => 'Y-m-d H:i:s',
'pk_convert_id' => false,
'type_map' => ['root' => 'array', 'document' => 'array'],
];
public function getQueryClass(): string
{
return Query::class;
}
public function getBuilder()
{
return $this->builder;
}
public function getBuilderClass(): string
{
return Builder::class;
}
public function connect(array $config = [], $linkNum = 0)
{
if (!isset($this->links[$linkNum])) {
if (empty($config)) {
$config = $this->config;
} else {
$config = array_merge($this->config, $config);
}
$this->dbName = $config['database'];
$this->typeMap = $config['type_map'];
if ($config['pk_convert_id'] && '_id' == $config['pk']) {
$this->config['pk'] = 'id';
}
if (empty($config['dsn'])) {
$config['dsn'] = 'mongodb://' . ($config['username'] ? "{$config['username']}" : '') . ($config['password'] ? ":{$config['password']}@" : '') . $config['hostname'] . ($config['hostport'] ? ":{$config['hostport']}" : '');
}
$startTime = microtime(true);
$this->links[$linkNum] = new Manager($config['dsn'], $config['params']);
if (!empty($config['trigger_sql'])) {
$this->trigger('CONNECT:[ UseTime:' . number_format(microtime(true) - $startTime, 6) . 's ] ' . $config['dsn']);
}
}
return $this->links[$linkNum];
}
public function getMongo()
{
return $this->mongo ?: null;
}
public function db(string $db = null)
{
if (is_null($db)) {
return $this->dbName;
} else {
$this->dbName = $db;
}
}
public function cursor($query)
{
$options = $query->parseOptions();
$mongoQuery = $this->builder->select($query);
$master = $query->getOptions('master') ? true : false;
return $this->getCursor($query, $mongoQuery, $master);
}
public function getCursor(BaseQuery $query, $mongoQuery, bool $master = false): Cursor
{
$this->initConnect($master);
$this->db->updateQueryTimes();
$options = $query->getOptions();
$namespace = $options['table'];
if (false === strpos($namespace, '.')) {
$namespace = $this->dbName . '.' . $namespace;
}
if (!empty($this->queryStr)) {
$this->queryStr = 'db' . strstr($namespace, '.') . '.' . $this->queryStr;
}
if ($mongoQuery instanceof Closure) {
$mongoQuery = $mongoQuery($query);
}
$readPreference = $options['readPreference'] ?? null;
$this->queryStartTime = microtime(true);
if ($session = $this->getSession()) {
$this->cursor = $this->mongo->executeQuery($namespace, $query, [
'readPreference' => is_null($readPreference) ? new ReadPreference(ReadPreference::RP_PRIMARY) : $readPreference,
'session' => $session,
]);
} else {
$this->cursor = $this->mongo->executeQuery($namespace, $mongoQuery, $readPreference);
}
if (!empty($this->config['trigger_sql'])) {
$this->trigger('', $master);
}
return $this->cursor;
}
public function query(MongoQuery $query)
{
return $this->mongoQuery($this->newQuery(), $query);
}
public function execute(BulkWrite $bulk)
{
return $this->mongoExecute($this->newQuery(), $bulk);
}
protected function mongoQuery(BaseQuery $query, $mongoQuery): array
{
$options = $query->parseOptions();
if ($query->getOptions('cache')) {
$cacheItem = $this->parseCache($query, $query->getOptions('cache'));
$key = $cacheItem->getKey();
if ($this->cache->has($key)) {
return $this->cache->get($key);
}
}
if ($mongoQuery instanceof Closure) {
$mongoQuery = $mongoQuery($query);
}
$master = $query->getOptions('master') ? true : false;
$this->getCursor($query, $mongoQuery, $master);
$resultSet = $this->getResult($options['typeMap']);
if (isset($cacheItem) && $resultSet) {
$cacheItem->set($resultSet);
$this->cacheData($cacheItem);
}
return $resultSet;
}
protected function mongoExecute(BaseQuery $query, BulkWrite $bulk)
{
$this->initConnect(true);
$this->db->updateQueryTimes();
$options = $query->getOptions();
$namespace = $options['table'];
if (false === strpos($namespace, '.')) {
$namespace = $this->dbName . '.' . $namespace;
}
if (!empty($this->queryStr)) {
$this->queryStr = 'db' . strstr($namespace, '.') . '.' . $this->queryStr;
}
$writeConcern = $options['writeConcern'] ?? null;
$this->queryStartTime = microtime(true);
if ($session = $this->getSession()) {
$writeResult = $this->mongo->executeBulkWrite($namespace, $bulk, [
'session' => $session,
'writeConcern' => is_null($writeConcern) ? new WriteConcern(1) : $writeConcern,
]);
} else {
$writeResult = $this->mongo->executeBulkWrite($namespace, $bulk, $writeConcern);
}
if (!empty($this->config['trigger_sql'])) {
$this->trigger();
}
$this->numRows = $writeResult->getMatchedCount();
if ($query->getOptions('cache')) {
$cacheItem = $this->parseCache($query, $query->getOptions('cache'));
$key = $cacheItem->getKey();
$tag = $cacheItem->getTag();
if (isset($key) && $this->cache->has($key)) {
$this->cache->delete($key);
} elseif (!empty($tag) && method_exists($this->cache, 'tag')) {
$this->cache->tag($tag)->clear();
}
}
return $writeResult;
}
public function command(Command $command, string $dbName = '', ReadPreference $readPreference = null, $typeMap = null, bool $master = false): array
{
$this->initConnect($master);
$this->db->updateQueryTimes();
$this->queryStartTime = microtime(true);
$dbName = $dbName ?: $this->dbName;
if (!empty($this->queryStr)) {
$this->queryStr = 'db.' . $this->queryStr;
}
if ($session = $this->getSession()) {
$this->cursor = $this->mongo->executeCommand($dbName, $command, [
'readPreference' => is_null($readPreference) ? new ReadPreference(ReadPreference::RP_PRIMARY) : $readPreference,
'session' => $session,
]);
} else {
$this->cursor = $this->mongo->executeCommand($dbName, $command, $readPreference);
}
if (!empty($this->config['trigger_sql'])) {
$this->trigger('', $master);
}
return $this->getResult($typeMap);
}
protected function getResult($typeMap = null): array
{
if (is_null($typeMap)) {
$typeMap = $this->typeMap;
}
$typeMap = is_string($typeMap) ? ['root' => $typeMap] : $typeMap;
$this->cursor->setTypeMap($typeMap);
$result = $this->cursor->toArray();
if ($this->getConfig('pk_convert_id')) {
foreach ($result as &$data) {
$this->convertObjectID($data);
}
}
$this->numRows = count($result);
return $result;
}
protected function convertObjectID(array &$data): void
{
if (isset($data['_id']) && is_object($data['_id'])) {
$data['id'] = $data['_id']->__toString();
unset($data['_id']);
}
}
public function mongoLog(string $type, $data, array $options = [])
{
if (!$this->config['trigger_sql']) {
return;
}
if (is_array($data)) {
array_walk_recursive($data, function (&$value) {
if ($value instanceof ObjectID) {
$value = $value->__toString();
}
});
}
switch (strtolower($type)) {
case 'aggregate':
$this->queryStr = 'runCommand(' . ($data ? json_encode($data) : '') . ');';
break;
case 'find':
$this->queryStr = $type . '(' . ($data ? json_encode($data) : '') . ')';
if (isset($options['sort'])) {
$this->queryStr .= '.sort(' . json_encode($options['sort']) . ')';
}
if (isset($options['skip'])) {
$this->queryStr .= '.skip(' . $options['skip'] . ')';
}
if (isset($options['limit'])) {
$this->queryStr .= '.limit(' . $options['limit'] . ')';
}
$this->queryStr .= ';';
break;
case 'insert':
case 'remove':
$this->queryStr = $type . '(' . ($data ? json_encode($data) : '') . ');';
break;
case 'update':
$this->queryStr = $type . '(' . json_encode($options) . ',' . json_encode($data) . ');';
break;
case 'cmd':
$this->queryStr = $data . '(' . json_encode($options) . ');';
break;
}
$this->options = $options;
}
public function getLastSql(): string
{
return $this->queryStr;
}
public function close()
{
$this->mongo = null;
$this->cursor = null;
$this->linkRead = null;
$this->linkWrite = null;
$this->links = [];
}
protected function initConnect(bool $master = true): void
{
if (!empty($this->config['deploy'])) {
if ($master) {
if (!$this->linkWrite) {
$this->linkWrite = $this->multiConnect(true);
}
$this->mongo = $this->linkWrite;
} else {
if (!$this->linkRead) {
$this->linkRead = $this->multiConnect(false);
}
$this->mongo = $this->linkRead;
}
} elseif (!$this->mongo) {
$this->mongo = $this->connect();
}
}
protected function multiConnect(bool $master = false): Manager
{
$config = [];
foreach (['username', 'password', 'hostname', 'hostport', 'database', 'dsn'] as $name) {
$config[$name] = is_string($this->config[$name]) ? explode(',', $this->config[$name]) : $this->config[$name];
}
$m = floor(mt_rand(0, $this->config['master_num'] - 1));
if ($this->config['rw_separate']) {
if ($master) {
if ($this->config['is_replica_set']) {
return $this->replicaSetConnect();
} else {
$r = $m;
}
} elseif (is_numeric($this->config['slave_no'])) {
$r = $this->config['slave_no'];
} else {
$r = floor(mt_rand($this->config['master_num'], count($config['hostname']) - 1));
}
} else {
$r = floor(mt_rand(0, count($config['hostname']) - 1));
}
$dbConfig = [];
foreach (['username', 'password', 'hostname', 'hostport', 'database', 'dsn'] as $name) {
$dbConfig[$name] = $config[$name][$r] ?? $config[$name][0];
}
return $this->connect($dbConfig, $r);
}
public function replicaSetConnect(): Manager
{
$this->dbName = $this->config['database'];
$this->typeMap = $this->config['type_map'];
$startTime = microtime(true);
$this->config['params']['replicaSet'] = $this->config['database'];
$manager = new Manager($this->buildUrl(), $this->config['params']);
if (!empty($config['trigger_sql'])) {
$this->trigger('CONNECT:ReplicaSet[ UseTime:' . number_format(microtime(true) - $startTime, 6) . 's ] ' . $this->config['dsn']);
}
return $manager;
}
private function buildUrl(): string
{
$url = 'mongodb://' . ($this->config['username'] ? "{$this->config['username']}" : '') . ($this->config['password'] ? ":{$this->config['password']}@" : '');
$hostList = is_string($this->config['hostname']) ? explode(',', $this->config['hostname']) : $this->config['hostname'];
$portList = is_string($this->config['hostport']) ? explode(',', $this->config['hostport']) : $this->config['hostport'];
for ($i = 0; $i < count($hostList); $i++) {
$url = $url . $hostList[$i] . ':' . $portList[0] . ',';
}
return rtrim($url, ",") . '/';
}
public function insert(BaseQuery $query, bool $getLastInsID = false)
{
$options = $query->parseOptions();
if (empty($options['data'])) {
throw new Exception('miss data to insert');
}
$bulk = $this->builder->insert($query);
$writeResult = $this->mongoExecute($query, $bulk);
$result = $writeResult->getInsertedCount();
if ($result) {
$data = $options['data'];
$lastInsId = $this->getLastInsID($query);
if ($lastInsId) {
$pk = $query->getPk();
$data[$pk] = $lastInsId;
}
$query->setOption('data', $data);
$this->db->trigger('after_insert', $query);
if ($getLastInsID) {
return $lastInsId;
}
}
return $result;
}
public function getLastInsID(BaseQuery $query)
{
$id = $this->builder->getLastInsID();
if (is_array($id)) {
array_walk($id, function (&$item, $key) {
if ($item instanceof ObjectID) {
$item = $item->__toString();
}
});
} elseif ($id instanceof ObjectID) {
$id = $id->__toString();
}
return $id;
}
public function insertAll(BaseQuery $query, array $dataSet = []): int
{
$query->parseOptions();
if (!is_array(reset($dataSet))) {
return 0;
}
$bulk = $this->builder->insertAll($query, $dataSet);
$writeResult = $this->mongoExecute($query, $bulk);
return $writeResult->getInsertedCount();
}
public function update(BaseQuery $query): int
{
$query->parseOptions();
$bulk = $this->builder->update($query);
$writeResult = $this->mongoExecute($query, $bulk);
$result = $writeResult->getModifiedCount();
if ($result) {
$this->db->trigger('after_update', $query);
}
return $result;
}
public function delete(BaseQuery $query): int
{
$query->parseOptions();
$bulk = $this->builder->delete($query);
$writeResult = $this->mongoExecute($query, $bulk);
$result = $writeResult->getDeletedCount();
if ($result) {
$this->db->trigger('after_delete', $query);
}
return $result;
}
public function select(BaseQuery $query): array
{
$resultSet = $this->db->trigger('before_select', $query);
if (!$resultSet) {
$resultSet = $this->mongoQuery($query, function ($query) {
return $this->builder->select($query);
});
}
return $resultSet;
}
public function find(BaseQuery $query): array
{
$result = $this->db->trigger('before_find', $query);
if (!$result) {
$resultSet = $this->mongoQuery($query, function ($query) {
return $this->builder->select($query, true);
});
$result = $resultSet[0] ?? [];
}
return $result;
}
public function value(BaseQuery $query, string $field, $default = null)
{
$options = $query->parseOptions();
if (isset($options['projection'])) {
$query->removeOption('projection');
}
$query->setOption('projection', (array) $field);
if (!empty($options['cache'])) {
$cacheItem = $this->parseCache($query, $options['cache']);
$key = $cacheItem->getKey();
if ($this->cache->has($key)) {
return $this->cache->get($key);
}
}
$mongoQuery = $this->builder->select($query, true);
if (isset($options['projection'])) {
$query->setOption('projection', $options['projection']);
} else {
$query->removeOption('projection');
}
$resultSet = $this->mongoQuery($query, $mongoQuery);
if (!empty($resultSet)) {
$data = array_shift($resultSet);
$result = $data[$field];
} else {
$result = false;
}
if (isset($cacheItem) && false !== $result) {
$cacheItem->set($result);
$this->cacheData($cacheItem);
}
return false !== $result ? $result : $default;
}
public function column(BaseQuery $query, $field, string $key = ''): array
{
$options = $query->parseOptions();
if (isset($options['projection'])) {
$query->removeOption('projection');
}
if (is_array($field)) {
$field = implode(',', $field);
}
if ($key && '*' != $field) {
$projection = $key . ',' . $field;
} else {
$projection = $field;
}
$query->field($projection);
if (!empty($options['cache'])) {
$cacheItem = $this->parseCache($query, $options['cache']);
$key = $cacheItem->getKey();
if ($this->cache->has($key)) {
return $this->cache->get($key);
}
}
$mongoQuery = $this->builder->select($query);
if (isset($options['projection'])) {
$query->setOption('projection', $options['projection']);
} else {
$query->removeOption('projection');
}
$resultSet = $this->mongoQuery($query, $mongoQuery);
if (('*' == $field || strpos($field, ',')) && $key) {
$result = array_column($resultSet, null, $key);
} elseif (!empty($resultSet)) {
$result = array_column($resultSet, $field, $key);
} else {
$result = [];
}
if (isset($cacheItem)) {
$cacheItem->set($result);
$this->cacheData($cacheItem);
}
return $result;
}
public function cmd(BaseQuery $query, $command, $extra = null, string $db = ''): array
{
if (is_array($command) || is_object($command)) {
$this->mongoLog('cmd', 'cmd', $command);
$command = new Command($command);
} else {
$command = $this->builder->$command($query, $extra);
}
return $this->command($command, $db);
}
public function getTableFields($tableName): array
{
return [];
}
public function transaction(callable $callback)
{
$this->startTrans();
try {
$result = null;
if (is_callable($callback)) {
$result = call_user_func_array($callback, [$this]);
}
$this->commit();
return $result;
} catch (\Exception $e) {
$this->rollback();
throw $e;
} catch (\Throwable $e) {
$this->rollback();
throw $e;
}
}
public function startTrans()
{
$this->initConnect(true);
$this->session_uuid = uniqid();
$this->sessions[$this->session_uuid] = $this->getMongo()->startSession();
$this->sessions[$this->session_uuid]->startTransaction([]);
}
public function commit()
{
if ($session = $this->getSession()) {
$session->commitTransaction();
$this->setLastSession();
}
}
public function rollback()
{
if ($session = $this->getSession()) {
$session->abortTransaction();
$this->setLastSession();
}
}
/**
* 结束当前会话,设置上一个会话为当前会话
* @author klinson <klinson@163.com>
*/
protected function setLastSession()
{
if ($session = $this->getSession()) {
$session->endSession();
unset($this->sessions[$this->session_uuid]);
if (empty($this->sessions)) {
$this->session_uuid = null;
} else {
end($this->sessions);
$this->session_uuid = key($this->sessions);
}
}
}
/**
* 获取当前会话
* @return \MongoDB\Driver\Session|null
* @author klinson <klinson@163.com>
*/
public function getSession()
{
return ($this->session_uuid && isset($this->sessions[$this->session_uuid]))
? $this->sessions[$this->session_uuid]
: null;
}
}