#include #include #include #include #include #include #include #include #include #include #include #include #include #include #include #include #include namespace DB { namespace ErrorCodes { extern const int NUMBER_OF_ARGUMENTS_DOESNT_MATCH; extern const int MONGODB_CANNOT_AUTHENTICATE; extern const int BAD_ARGUMENTS; } StorageMongoDB::StorageMongoDB( const StorageID & table_id_, const std::string & host_, uint16_t port_, const std::string & database_name_, const std::string & collection_name_, const std::string & username_, const std::string & password_, const std::string & options_, const ColumnsDescription & columns_, const ConstraintsDescription & constraints_, const String & comment) : IStorage(table_id_) , database_name(database_name_) , collection_name(collection_name_) , username(username_) , password(password_) , uri("mongodb://" + host_ + ":" + std::to_string(port_) + "/" + database_name_ + "?" + options_) { StorageInMemoryMetadata storage_metadata; storage_metadata.setColumns(columns_); storage_metadata.setConstraints(constraints_); storage_metadata.setComment(comment); setInMemoryMetadata(storage_metadata); } void StorageMongoDB::connectIfNotConnected() { std::lock_guard lock{connection_mutex}; if (!connection) { StorageMongoDBSocketFactory factory; connection = std::make_shared(uri, factory); } if (!authenticated) { Poco::URI poco_uri(uri); auto query_params = poco_uri.getQueryParameters(); auto auth_source = std::find_if(query_params.begin(), query_params.end(), [&](const std::pair & param) { return param.first == "authSource"; }); auto auth_db = database_name; if (auth_source != query_params.end()) auth_db = auth_source->second; #if POCO_VERSION >= 0x01070800 if (!username.empty() && !password.empty()) { Poco::MongoDB::Database poco_db(auth_db); if (!poco_db.authenticate(*connection, username, password, Poco::MongoDB::Database::AUTH_SCRAM_SHA1)) throw Exception("Cannot authenticate in MongoDB, incorrect user or password", ErrorCodes::MONGODB_CANNOT_AUTHENTICATE); } #else authenticate(*connection, database_name, username, password); #endif authenticated = true; } } class StorageMongoDBSink : public SinkToStorage { public: explicit StorageMongoDBSink( const std::string & collection_name_, const std::string & db_name_, const StorageMetadataPtr & metadata_snapshot_, std::shared_ptr connection_ ) : SinkToStorage(metadata_snapshot_->getSampleBlock()) , collection_name(collection_name_) , db_name(db_name_) , metadata_snapshot{metadata_snapshot_} , connection(connection_) { } String getName() const override { return "StorageMySQLSink"; } void consume(Chunk chunk) override { auto block = getHeader().cloneWithColumns(chunk.detachColumns()); Poco::MongoDB::Database db(db_name); Poco::MongoDB::Document::Ptr index = new Poco::MongoDB::Document(); const auto columns = block.getColumns(); const size_t num_rows = block.rows(), num_cols = block.columns(); const auto data_types = block.getDataTypes(); const auto data_names = block.getNames(); std::vector row(num_cols); for (const auto i : collections::range(0, num_rows)) { for (const auto j : collections::range(0, num_cols)) { WriteBufferFromOwnString ostr; data_types[j]->getDefaultSerialization()->serializeText(*columns[j], i, ostr, FormatSettings{}); row[j] = ostr.str(); index->add(data_names[j], row[j]); } } Poco::SharedPtr insert_request = db.createInsertRequest(collection_name); insert_request->documents().push_back(index); connection->sendRequest(*insert_request); } private: String collection_name; String db_name; StorageMetadataPtr metadata_snapshot; std::shared_ptr connection; }; Pipe StorageMongoDB::read( const Names & column_names, const StorageSnapshotPtr & storage_snapshot, SelectQueryInfo & /*query_info*/, ContextPtr /*context*/, QueryProcessingStage::Enum /*processed_stage*/, size_t max_block_size, unsigned) { connectIfNotConnected(); storage_snapshot->check(column_names); Block sample_block; for (const String & column_name : column_names) { auto column_data = storage_snapshot->metadata->getColumns().getPhysical(column_name); sample_block.insert({ column_data.type, column_data.name }); } return Pipe(std::make_shared(connection, createCursor(database_name, collection_name, sample_block), sample_block, max_block_size)); } SinkToStoragePtr StorageMongoDB::write(const ASTPtr & /*query*/, const StorageMetadataPtr & metadata_snapshot, ContextPtr /*context*/) { connectIfNotConnected(); return std::make_shared(collection_name, database_name, metadata_snapshot, connection); } StorageMongoDBConfiguration StorageMongoDB::getConfiguration(ASTs engine_args, ContextPtr context) { StorageMongoDBConfiguration configuration; if (auto named_collection = getExternalDataSourceConfiguration(engine_args, context)) { auto [common_configuration, storage_specific_args, _] = named_collection.value(); configuration.set(common_configuration); for (const auto & [arg_name, arg_value] : storage_specific_args) { if (arg_name == "options") configuration.options = checkAndGetLiteralArgument(arg_value, "options"); else throw Exception(ErrorCodes::BAD_ARGUMENTS, "Unexpected key-value argument." "Got: {}, but expected one of:" "host, port, username, password, database, table, options.", arg_name); } } else { if (engine_args.size() < 5 || engine_args.size() > 6) throw Exception( "Storage MongoDB requires from 5 to 6 parameters: MongoDB('host:port', database, collection, 'user', 'password' [, 'options']).", ErrorCodes::NUMBER_OF_ARGUMENTS_DOESNT_MATCH); for (auto & engine_arg : engine_args) engine_arg = evaluateConstantExpressionOrIdentifierAsLiteral(engine_arg, context); /// 27017 is the default MongoDB port. auto parsed_host_port = parseAddress(checkAndGetLiteralArgument(engine_args[0], "host:port"), 27017); configuration.host = parsed_host_port.first; configuration.port = parsed_host_port.second; configuration.database = checkAndGetLiteralArgument(engine_args[1], "database"); configuration.table = checkAndGetLiteralArgument(engine_args[2], "table"); configuration.username = checkAndGetLiteralArgument(engine_args[3], "username"); configuration.password = checkAndGetLiteralArgument(engine_args[4], "password"); if (engine_args.size() >= 6) configuration.options = checkAndGetLiteralArgument(engine_args[5], "database"); } context->getRemoteHostFilter().checkHostAndPort(configuration.host, toString(configuration.port)); return configuration; } void registerStorageMongoDB(StorageFactory & factory) { factory.registerStorage("MongoDB", [](const StorageFactory::Arguments & args) { auto configuration = StorageMongoDB::getConfiguration(args.engine_args, args.getLocalContext()); return std::make_shared( args.table_id, configuration.host, configuration.port, configuration.database, configuration.table, configuration.username, configuration.password, configuration.options, args.columns, args.constraints, args.comment); }, { .source_access_type = AccessType::MONGO, }); } }