static PyObject* _cbson_get_more_message(PyObject* self, PyObject* args) {
/* NOTE just using a random number as the request_id */
int request_id = rand();
char* collection_name = NULL;
int collection_name_length;
int num_to_return;
long long cursor_id;
buffer_t buffer;
int length_location, message_length;
PyObject* result;
if (!PyArg_ParseTuple(args, "et#iL",
"utf-8",
&collection_name,
&collection_name_length,
&num_to_return,
&cursor_id)) {
return NULL;
}
buffer = buffer_new();
if (!buffer) {
PyErr_NoMemory();
PyMem_Free(collection_name);
return NULL;
}
// save space for message length
length_location = buffer_save_space(buffer, 4);
if (length_location == -1) {
PyMem_Free(collection_name);
PyErr_NoMemory();
return NULL;
}
if (!buffer_write_bytes(buffer, (const char*)&request_id, 4) ||
!buffer_write_bytes(buffer,
"\x00\x00\x00\x00"
"\xd5\x07\x00\x00"
"\x00\x00\x00\x00", 12) ||
!buffer_write_bytes(buffer,
collection_name,
collection_name_length + 1) ||
!buffer_write_bytes(buffer, (const char*)&num_to_return, 4) ||
!buffer_write_bytes(buffer, (const char*)&cursor_id, 8)) {
buffer_free(buffer);
PyMem_Free(collection_name);
return NULL;
}
PyMem_Free(collection_name);
message_length = buffer_get_position(buffer) - length_location;
memcpy(buffer_get_buffer(buffer) + length_location, &message_length, 4);
/* objectify buffer */
result = Py_BuildValue("i" BYTES_FORMAT_STRING, request_id,
buffer_get_buffer(buffer),
buffer_get_position(buffer));
buffer_free(buffer);
return result;
}
/* add a lastError message on the end of the buffer.
* returns 0 on failure */
static int add_last_error(buffer_t buffer, int request_id, PyObject* args) {
int message_start;
int document_start;
int message_length;
int document_length;
PyObject* key;
PyObject* value;
Py_ssize_t pos = 0;
PyObject* one;
message_start = buffer_save_space(buffer, 4);
if (message_start == -1) {
PyErr_NoMemory();
return 0;
}
if (!buffer_write_bytes(buffer, (const char*)&request_id, 4) ||
!buffer_write_bytes(buffer,
"\x00\x00\x00\x00" /* responseTo */
"\xd4\x07\x00\x00" /* opcode */
"\x00\x00\x00\x00" /* options */
"admin.$cmd\x00" /* collection name */
"\x00\x00\x00\x00" /* skip */
"\xFF\xFF\xFF\xFF", /* limit (-1) */
31)) {
return 0;
}
/* save space for length */
document_start = buffer_save_space(buffer, 4);
if (document_start == -1) {
PyErr_NoMemory();
return 0;
}
/* getlasterror: 1 */
one = PyLong_FromLong(1);
if (!write_pair(buffer, "getlasterror", 12, one, 0, 4, 1)) {
Py_DECREF(one);
return 0;
}
Py_DECREF(one);
/* getlasterror options */
while (PyDict_Next(args, &pos, &key, &value)) {
if (!decode_and_write_pair(buffer, key, value, 0, 4, 0)) {
return 0;
}
}
/* EOD */
if (!buffer_write_bytes(buffer, "\x00", 1)) {
return 0;
}
message_length = buffer_get_position(buffer) - message_start;
document_length = buffer_get_position(buffer) - document_start;
memcpy(buffer_get_buffer(buffer) + message_start, &message_length, 4);
memcpy(buffer_get_buffer(buffer) + document_start, &document_length, 4);
return 1;
}
static PyObject*
_send_write_command(PyObject* ctx, buffer_t buffer, int lst_len_loc,
int cmd_len_loc, unsigned char* errors,
PyObject* to_publish) {
PyObject* result;
int request_id = rand();
int position = buffer_get_position(buffer);
int length = position - lst_len_loc - 1;
memcpy(buffer_get_buffer(buffer) + lst_len_loc, &length, 4);
length = position - cmd_len_loc;
memcpy(buffer_get_buffer(buffer) + cmd_len_loc, &length, 4);
memcpy(buffer_get_buffer(buffer), &position, 4);
memcpy(buffer_get_buffer(buffer) + 4, &request_id, 4);
/* Send the current batch */
result = PyObject_CallMethod(ctx, "write_command",
"i" BYTES_FORMAT_STRING "O",
request_id,
buffer_get_buffer(buffer),
buffer_get_position(buffer),
to_publish);
if (result && PyDict_GetItemString(result, "writeErrors"))
*errors = 1;
return result;
}
static void write_doc(buffer_t buffer, VALUE hash, VALUE check_keys) {
buffer_position start_position = buffer_get_position(buffer);
buffer_position length_location = buffer_save_space(buffer, 4);
buffer_position length;
VALUE id_str = rb_str_new2("_id");
VALUE id_sym = ID2SYM(rb_intern("_id"));
if (length_location == -1) {
rb_raise(rb_eNoMemError, "failed to allocate memory in buffer.c");
}
if (rb_funcall(hash, rb_intern("has_key?"), 1, id_str) == Qtrue) {
VALUE id = rb_hash_aref(hash, id_str);
write_element_allow_id(id_str, id, pack_extra(buffer, check_keys), 1);
} else if (rb_funcall(hash, rb_intern("has_key?"), 1, id_sym) == Qtrue) {
VALUE id = rb_hash_aref(hash, id_sym);
write_element_allow_id(id_sym, id, pack_extra(buffer, check_keys), 1);
}
// we have to check for an OrderedHash and handle that specially
if (strcmp(rb_class2name(RBASIC(hash)->klass), "OrderedHash") == 0) {
VALUE keys = rb_funcall(hash, rb_intern("keys"), 0);
int i;
for(i = 0; i < RARRAY_LEN(keys); i++) {
VALUE key = RARRAY_PTR(keys)[i];
VALUE value = rb_hash_aref(hash, key);
write_element(key, value, pack_extra(buffer, check_keys));
}
} else {
rb_hash_foreach(hash, write_element, pack_extra(buffer, check_keys));
}
// write null byte and fill in length
SAFE_WRITE(buffer, &zero, 1);
length = buffer_get_position(buffer) - start_position;
// make sure that length doesn't exceed 4MB
if (length > 4 * 1024 * 1024) {
buffer_free(buffer);
rb_raise(InvalidDocument, "Document too large: BSON documents are limited to 4MB.");
return;
}
SAFE_WRITE_AT_POS(buffer, length_location, (const char*)&length, 4);
}
static VALUE method_serialize(VALUE self, VALUE doc, VALUE check_keys, VALUE move_id) {
VALUE result;
buffer_t buffer = buffer_new();
if (buffer == NULL) {
rb_raise(rb_eNoMemError, "failed to allocate memory in buffer.c");
}
write_doc(buffer, doc, check_keys, move_id);
result = rb_str_new(buffer_get_buffer(buffer), buffer_get_position(buffer));
if (bson_buffer_free(buffer) != 0) {
rb_raise(rb_eRuntimeError, "failed to free buffer");
}
return result;
}
static PyObject*
_send_insert(PyObject* self, PyObject* ctx,
PyObject* gle_args, buffer_t buffer,
char* coll_name, int coll_len, int request_id, int safe,
codec_options_t* options, PyObject* to_publish) {
if (safe) {
if (!add_last_error(self, buffer, request_id,
coll_name, coll_len, options, gle_args)) {
return NULL;
}
}
/* The max_doc_size parameter for legacy_write is the max size of any
* document in buffer. We enforced max size already, pass 0 here. */
return PyObject_CallMethod(ctx, "legacy_write",
"i" BYTES_FORMAT_STRING "iNO",
request_id,
buffer_get_buffer(buffer),
buffer_get_position(buffer),
0,
PyBool_FromLong((long)safe),
to_publish);
}
static void write_doc(buffer_t buffer, VALUE hash, VALUE check_keys, VALUE move_id) {
buffer_position start_position = buffer_get_position(buffer);
buffer_position length_location = buffer_save_space(buffer, 4);
buffer_position length;
int allow_id;
int (*write_function)(VALUE, VALUE, VALUE) = NULL;
VALUE id_str = rb_str_new2("_id");
VALUE id_sym = ID2SYM(rb_intern("_id"));
if (length_location == -1) {
rb_raise(rb_eNoMemError, "failed to allocate memory in buffer.c");
}
// write '_id' first if move_id is true. then don't allow an id to be written.
if(move_id == Qtrue) {
allow_id = 0;
if (rb_funcall(hash, rb_intern("has_key?"), 1, id_str) == Qtrue) {
VALUE id = rb_hash_aref(hash, id_str);
write_element_with_id(id_str, id, pack_extra(buffer, check_keys));
} else if (rb_funcall(hash, rb_intern("has_key?"), 1, id_sym) == Qtrue) {
VALUE id = rb_hash_aref(hash, id_sym);
write_element_with_id(id_sym, id, pack_extra(buffer, check_keys));
}
}
else {
allow_id = 1;
// Ensure that hash doesn't contain both '_id' and :_id
if ((rb_obj_classname(hash), "Hash") == 0) {
if ((rb_funcall(hash, rb_intern("has_key?"), 1, id_str) == Qtrue) &&
(rb_funcall(hash, rb_intern("has_key?"), 1, id_sym) == Qtrue)) {
VALUE oid_sym = rb_hash_delete(hash, id_sym);
rb_funcall(hash, rb_intern("[]="), 2, id_str, oid_sym);
}
}
}
if(allow_id == 1) {
write_function = write_element_with_id;
}
else {
write_function = write_element_without_id;
}
// we have to check for an OrderedHash and handle that specially
if (strcmp(rb_obj_classname(hash), "BSON::OrderedHash") == 0) {
VALUE keys = rb_funcall(hash, rb_intern("keys"), 0);
int i;
for(i = 0; i < RARRAY_LEN(keys); i++) {
VALUE key = rb_ary_entry(keys, i);
VALUE value = rb_hash_aref(hash, key);
write_function(key, value, pack_extra(buffer, check_keys));
}
} else if (rb_obj_is_kind_of(hash, RB_HASH) == Qtrue) {
rb_hash_foreach(hash, write_function, pack_extra(buffer, check_keys));
} else {
bson_buffer_free(buffer);
rb_raise(InvalidDocument, "BSON.serialize takes a Hash but got a %s", rb_obj_classname(hash));
}
// write null byte and fill in length
SAFE_WRITE(buffer, &zero, 1);
length = buffer_get_position(buffer) - start_position;
// make sure that length doesn't exceed 4MB
if (length > max_bson_size) {
bson_buffer_free(buffer);
rb_raise(InvalidDocument, "Document too large: BSON documents are limited to %d bytes.", max_bson_size);
return;
}
SAFE_WRITE_AT_POS(buffer, length_location, (const char*)&length, 4);
}
static int write_element(VALUE key, VALUE value, VALUE extra, int allow_id) {
buffer_t buffer = (buffer_t)NUM2LL(rb_ary_entry(extra, 0));
VALUE check_keys = rb_ary_entry(extra, 1);
if (TYPE(key) == T_SYMBOL) {
// TODO better way to do this... ?
key = rb_str_new2(rb_id2name(SYM2ID(key)));
}
if (TYPE(key) != T_STRING) {
bson_buffer_free(buffer);
rb_raise(rb_eTypeError, "keys must be strings or symbols");
}
if (allow_id == 0 && strcmp("_id", RSTRING_PTR(key)) == 0) {
return ST_CONTINUE;
}
if (check_keys == Qtrue) {
int i;
if (RSTRING_LEN(key) > 0 && RSTRING_PTR(key)[0] == '$') {
bson_buffer_free(buffer);
rb_raise(InvalidKeyName, "%s - key must not start with '$'", RSTRING_PTR(key));
}
for (i = 0; i < RSTRING_LEN(key); i++) {
if (RSTRING_PTR(key)[i] == '.') {
bson_buffer_free(buffer);
rb_raise(InvalidKeyName, "%s - key must not contain '.'", RSTRING_PTR(key));
}
}
}
switch(TYPE(value)) {
case T_BIGNUM:
{
if (rb_funcall(value, gt_operator, 1, LL2NUM(9223372036854775807LL)) == Qtrue ||
rb_funcall(value, lt_operator, 1, LL2NUM(-9223372036854775808ULL)) == Qtrue) {
bson_buffer_free(buffer);
rb_raise(rb_eRangeError, "MongoDB can only handle 8-byte ints");
}
}
// NOTE: falls through to T_FIXNUM code
case T_FIXNUM:
{
long long ll_value;
ll_value = NUM2LL(value);
if (ll_value > 2147483647LL ||
ll_value < -2147483648LL) {
write_name_and_type(buffer, key, 0x12);
SAFE_WRITE(buffer, (char*)&ll_value, 8);
} else {
int int_value;
write_name_and_type(buffer, key, 0x10);
int_value = (int)ll_value;
SAFE_WRITE(buffer, (char*)&int_value, 4);
}
break;
}
case T_TRUE:
{
write_name_and_type(buffer, key, 0x08);
SAFE_WRITE(buffer, &one, 1);
break;
}
case T_FALSE:
{
write_name_and_type(buffer, key, 0x08);
SAFE_WRITE(buffer, &zero, 1);
break;
}
case T_FLOAT:
{
double d = NUM2DBL(value);
write_name_and_type(buffer, key, 0x01);
SAFE_WRITE(buffer, (char*)&d, 8);
break;
}
case T_NIL:
{
write_name_and_type(buffer, key, 0x0A);
break;
}
case T_HASH:
{
write_name_and_type(buffer, key, 0x03);
write_doc(buffer, value, check_keys, Qfalse);
break;
}
case T_ARRAY:
{
buffer_position length_location, start_position, obj_length;
int items, i;
VALUE* values;
write_name_and_type(buffer, key, 0x04);
start_position = buffer_get_position(buffer);
// save space for length
length_location = buffer_save_space(buffer, 4);
if (length_location == -1) {
rb_raise(rb_eNoMemError, "failed to allocate memory in buffer.c");
}
items = RARRAY_LENINT(value);
for(i = 0; i < items; i++) {
char* name;
VALUE key;
INT2STRING(&name, i);
key = rb_str_new2(name);
write_element_with_id(key, rb_ary_entry(value, i), pack_extra(buffer, check_keys));
FREE_INTSTRING(name);
}
// write null byte and fill in length
SAFE_WRITE(buffer, &zero, 1);
obj_length = buffer_get_position(buffer) - start_position;
SAFE_WRITE_AT_POS(buffer, length_location, (const char*)&obj_length, 4);
break;
}
case T_STRING:
{
int length;
write_name_and_type(buffer, key, 0x02);
length = RSTRING_LENINT(value) + 1;
SAFE_WRITE(buffer, (char*)&length, 4);
write_utf8(buffer, value, 0);
SAFE_WRITE(buffer, &zero, 1);
break;
}
case T_SYMBOL:
{
const char* str_value = rb_id2name(SYM2ID(value));
int length = (int)strlen(str_value) + 1;
write_name_and_type(buffer, key, 0x0E);
SAFE_WRITE(buffer, (char*)&length, 4);
SAFE_WRITE(buffer, str_value, length);
break;
}
case T_OBJECT:
{
// TODO there has to be a better way to do these checks...
const char* cls = rb_obj_classname(value);
if (strcmp(cls, "BSON::Binary") == 0 ||
strcmp(cls, "ByteBuffer") == 0) {
const char subtype = strcmp(cls, "ByteBuffer") ?
(const char)FIX2INT(rb_funcall(value, rb_intern("subtype"), 0)) : 2;
VALUE string_data = rb_funcall(value, rb_intern("to_s"), 0);
int length = RSTRING_LENINT(string_data);
write_name_and_type(buffer, key, 0x05);
if (subtype == 2) {
const int other_length = length + 4;
SAFE_WRITE(buffer, (const char*)&other_length, 4);
SAFE_WRITE(buffer, &subtype, 1);
}
SAFE_WRITE(buffer, (const char*)&length, 4);
if (subtype != 2) {
SAFE_WRITE(buffer, &subtype, 1);
}
SAFE_WRITE(buffer, RSTRING_PTR(string_data), length);
break;
}
if (strcmp(cls, "BSON::ObjectId") == 0) {
VALUE as_array = rb_funcall(value, rb_intern("to_a"), 0);
int i;
write_name_and_type(buffer, key, 0x07);
for (i = 0; i < 12; i++) {
char byte = (char)FIX2INT(rb_ary_entry(as_array, i));
SAFE_WRITE(buffer, &byte, 1);
}
break;
}
if (strcmp(cls, "BSON::DBRef") == 0) {
buffer_position length_location, start_position, obj_length;
VALUE ns, oid;
write_name_and_type(buffer, key, 0x03);
start_position = buffer_get_position(buffer);
// save space for length
length_location = buffer_save_space(buffer, 4);
if (length_location == -1) {
rb_raise(rb_eNoMemError, "failed to allocate memory in buffer.c");
}
ns = rb_funcall(value, rb_intern("namespace"), 0);
write_element_with_id(rb_str_new2("$ref"), ns, pack_extra(buffer, Qfalse));
oid = rb_funcall(value, rb_intern("object_id"), 0);
write_element_with_id(rb_str_new2("$id"), oid, pack_extra(buffer, Qfalse));
// write null byte and fill in length
SAFE_WRITE(buffer, &zero, 1);
obj_length = buffer_get_position(buffer) - start_position;
SAFE_WRITE_AT_POS(buffer, length_location, (const char*)&obj_length, 4);
break;
}
if (strcmp(cls, "BSON::Code") == 0) {
buffer_position length_location, start_position, total_length;
int length;
VALUE code_str;
write_name_and_type(buffer, key, 0x0F);
start_position = buffer_get_position(buffer);
length_location = buffer_save_space(buffer, 4);
if (length_location == -1) {
rb_raise(rb_eNoMemError, "failed to allocate memory in buffer.c");
}
code_str = rb_funcall(value, rb_intern("code"), 0);
length = RSTRING_LENINT(code_str) + 1;
SAFE_WRITE(buffer, (char*)&length, 4);
SAFE_WRITE(buffer, RSTRING_PTR(code_str), length - 1);
SAFE_WRITE(buffer, &zero, 1);
write_doc(buffer, rb_funcall(value, rb_intern("scope"), 0), Qfalse, Qfalse);
total_length = buffer_get_position(buffer) - start_position;
SAFE_WRITE_AT_POS(buffer, length_location, (const char*)&total_length, 4);
break;
}
if (strcmp(cls, "BSON::MaxKey") == 0) {
write_name_and_type(buffer, key, 0x7f);
break;
}
if (strcmp(cls, "BSON::MinKey") == 0) {
write_name_and_type(buffer, key, 0xff);
break;
}
if (strcmp(cls, "BSON::Timestamp") == 0) {
write_name_and_type(buffer, key, 0x11);
int seconds = FIX2INT(
rb_funcall(value, rb_intern("seconds"), 0));
int increment = FIX2INT(
rb_funcall(value, rb_intern("increment"), 0));
SAFE_WRITE(buffer, (const char*)&increment, 4);
SAFE_WRITE(buffer, (const char*)&seconds, 4);
break;
}
if (strcmp(cls, "DateTime") == 0 || strcmp(cls, "Date") == 0 || strcmp(cls, "ActiveSupport::TimeWithZone") == 0) {
bson_buffer_free(buffer);
rb_raise(InvalidDocument, "%s is not currently supported; use a UTC Time instance instead.", cls);
break;
}
if(strcmp(cls, "Complex") == 0 || strcmp(cls, "Rational") == 0 || strcmp(cls, "BigDecimal") == 0) {
bson_buffer_free(buffer);
rb_raise(InvalidDocument, "Cannot serialize the Numeric type %s as BSON; only Bignum, Fixnum, and Float are supported.", cls);
break;
}
bson_buffer_free(buffer);
rb_raise(InvalidDocument, "Cannot serialize an object of class %s into BSON.", cls);
break;
}
case T_DATA:
{
const char* cls = rb_obj_classname(value);
if (strcmp(cls, "Time") == 0) {
double t = NUM2DBL(rb_funcall(value, rb_intern("to_f"), 0));
long long time_since_epoch = (long long)round(t * 1000);
write_name_and_type(buffer, key, 0x09);
SAFE_WRITE(buffer, (const char*)&time_since_epoch, 8);
break;
}
if(strcmp(cls, "BigDecimal") == 0) {
bson_buffer_free(buffer);
rb_raise(InvalidDocument, "Cannot serialize the Numeric type %s as BSON; only Bignum, Fixnum, and Float are supported.", cls);
break;
}
bson_buffer_free(buffer);
rb_raise(InvalidDocument, "Cannot serialize an object of class %s into BSON.", cls);
break;
}
case T_REGEXP:
{
VALUE pattern = RREGEXP_SRC(value);
long flags = RREGEXP_OPTIONS(value);
VALUE has_extra;
write_name_and_type(buffer, key, 0x0B);
write_utf8(buffer, pattern, 1);
SAFE_WRITE(buffer, &zero, 1);
if (flags & IGNORECASE) {
char ignorecase = 'i';
SAFE_WRITE(buffer, &ignorecase, 1);
}
if (flags & MULTILINE) {
char multiline = 'm';
char dotall = 's';
SAFE_WRITE(buffer, &multiline, 1);
SAFE_WRITE(buffer, &dotall, 1);
}
if (flags & EXTENDED) {
char extended = 'x';
SAFE_WRITE(buffer, &extended, 1);
}
has_extra = rb_funcall(value, rb_intern("respond_to?"), 1, rb_str_new2("extra_options_str"));
if (TYPE(has_extra) == T_TRUE) {
VALUE extra = rb_funcall(value, rb_intern("extra_options_str"), 0);
buffer_position old_position = buffer_get_position(buffer);
SAFE_WRITE(buffer, RSTRING_PTR(extra), RSTRING_LENINT(extra));
qsort(buffer_get_buffer(buffer) + old_position, RSTRING_LEN(extra), sizeof(char), cmp_char);
}
SAFE_WRITE(buffer, &zero, 1);
break;
}
default:
{
const char* cls = rb_obj_classname(value);
bson_buffer_free(buffer);
rb_raise(InvalidDocument, "Cannot serialize an object of class %s (type %d) into BSON.", cls, TYPE(value));
break;
}
}
return ST_CONTINUE;
}
/* add a lastError message on the end of the buffer.
* returns 0 on failure */
static int add_last_error(PyObject* self, buffer_t buffer,
int request_id, char* ns, int nslen, PyObject* args) {
struct module_state *state = GETSTATE(self);
int message_start;
int document_start;
int message_length;
int document_length;
PyObject* key;
PyObject* value;
Py_ssize_t pos = 0;
PyObject* one;
char *p = strchr(ns, '.');
/* Length of the database portion of ns. */
nslen = p ? (int)(p - ns) : nslen;
message_start = buffer_save_space(buffer, 4);
if (message_start == -1) {
PyErr_NoMemory();
return 0;
}
if (!buffer_write_bytes(buffer, (const char*)&request_id, 4) ||
!buffer_write_bytes(buffer,
"\x00\x00\x00\x00" /* responseTo */
"\xd4\x07\x00\x00" /* opcode */
"\x00\x00\x00\x00", /* options */
12) ||
!buffer_write_bytes(buffer,
ns, nslen) || /* database */
!buffer_write_bytes(buffer,
".$cmd\x00" /* collection name */
"\x00\x00\x00\x00" /* skip */
"\xFF\xFF\xFF\xFF", /* limit (-1) */
14)) {
return 0;
}
/* save space for length */
document_start = buffer_save_space(buffer, 4);
if (document_start == -1) {
PyErr_NoMemory();
return 0;
}
/* getlasterror: 1 */
if (!(one = PyLong_FromLong(1)))
return 0;
if (!write_pair(state->_cbson, buffer, "getlasterror", 12, one, 0, 4, 1)) {
Py_DECREF(one);
return 0;
}
Py_DECREF(one);
/* getlasterror options */
while (PyDict_Next(args, &pos, &key, &value)) {
if (!decode_and_write_pair(state->_cbson, buffer, key, value, 0, 4, 0)) {
return 0;
}
}
/* EOD */
if (!buffer_write_bytes(buffer, "\x00", 1)) {
return 0;
}
message_length = buffer_get_position(buffer) - message_start;
document_length = buffer_get_position(buffer) - document_start;
memcpy(buffer_get_buffer(buffer) + message_start, &message_length, 4);
memcpy(buffer_get_buffer(buffer) + document_start, &document_length, 4);
return 1;
}
static PyObject* _cbson_query_message(PyObject* self, PyObject* args) {
/* NOTE just using a random number as the request_id */
struct module_state *state = GETSTATE(self);
int request_id = rand();
unsigned int options;
char* collection_name = NULL;
int collection_name_length;
int begin, cur_size, max_size = 0;
int num_to_skip;
int num_to_return;
PyObject* query;
PyObject* field_selector = Py_None;
unsigned char uuid_subtype = 3;
buffer_t buffer;
int length_location, message_length;
PyObject* result;
if (!PyArg_ParseTuple(args, "Iet#iiO|Ob",
&options,
"utf-8",
&collection_name,
&collection_name_length,
&num_to_skip, &num_to_return,
&query, &field_selector, &uuid_subtype)) {
return NULL;
}
buffer = buffer_new();
if (!buffer) {
PyErr_NoMemory();
PyMem_Free(collection_name);
return NULL;
}
// save space for message length
length_location = buffer_save_space(buffer, 4);
if (length_location == -1) {
PyMem_Free(collection_name);
PyErr_NoMemory();
return NULL;
}
if (!buffer_write_bytes(buffer, (const char*)&request_id, 4) ||
!buffer_write_bytes(buffer, "\x00\x00\x00\x00\xd4\x07\x00\x00", 8) ||
!buffer_write_bytes(buffer, (const char*)&options, 4) ||
!buffer_write_bytes(buffer, collection_name,
collection_name_length + 1) ||
!buffer_write_bytes(buffer, (const char*)&num_to_skip, 4) ||
!buffer_write_bytes(buffer, (const char*)&num_to_return, 4)) {
buffer_free(buffer);
PyMem_Free(collection_name);
return NULL;
}
begin = buffer_get_position(buffer);
if (!write_dict(state->_cbson, buffer, query, 0, uuid_subtype, 1)) {
buffer_free(buffer);
PyMem_Free(collection_name);
return NULL;
}
max_size = buffer_get_position(buffer) - begin;
if (field_selector != Py_None) {
begin = buffer_get_position(buffer);
if (!write_dict(state->_cbson, buffer, field_selector, 0, uuid_subtype, 1)) {
buffer_free(buffer);
PyMem_Free(collection_name);
return NULL;
}
cur_size = buffer_get_position(buffer) - begin;
max_size = (cur_size > max_size) ? cur_size : max_size;
}
PyMem_Free(collection_name);
message_length = buffer_get_position(buffer) - length_location;
memcpy(buffer_get_buffer(buffer) + length_location, &message_length, 4);
/* objectify buffer */
result = Py_BuildValue("i" BYTES_FORMAT_STRING "i", request_id,
buffer_get_buffer(buffer),
buffer_get_position(buffer),
max_size);
buffer_free(buffer);
return result;
}
static PyObject* _cbson_update_message(PyObject* self, PyObject* args) {
/* NOTE just using a random number as the request_id */
struct module_state *state = GETSTATE(self);
int request_id = rand();
char* collection_name = NULL;
int collection_name_length;
int before, cur_size, max_size = 0;
PyObject* doc;
PyObject* spec;
unsigned char multi;
unsigned char upsert;
unsigned char safe;
unsigned char check_keys;
unsigned char uuid_subtype;
PyObject* last_error_args;
int options;
buffer_t buffer;
int length_location, message_length;
PyObject* result;
if (!PyArg_ParseTuple(args, "et#bbOObObb",
"utf-8",
&collection_name,
&collection_name_length,
&upsert, &multi, &spec, &doc, &safe,
&last_error_args, &check_keys, &uuid_subtype)) {
return NULL;
}
options = 0;
if (upsert) {
options += 1;
}
if (multi) {
options += 2;
}
buffer = buffer_new();
if (!buffer) {
PyErr_NoMemory();
PyMem_Free(collection_name);
return NULL;
}
// save space for message length
length_location = buffer_save_space(buffer, 4);
if (length_location == -1) {
PyMem_Free(collection_name);
PyErr_NoMemory();
return NULL;
}
if (!buffer_write_bytes(buffer, (const char*)&request_id, 4) ||
!buffer_write_bytes(buffer,
"\x00\x00\x00\x00"
"\xd1\x07\x00\x00"
"\x00\x00\x00\x00",
12) ||
!buffer_write_bytes(buffer,
collection_name,
collection_name_length + 1) ||
!buffer_write_bytes(buffer, (const char*)&options, 4)) {
buffer_free(buffer);
PyMem_Free(collection_name);
return NULL;
}
before = buffer_get_position(buffer);
if (!write_dict(state->_cbson, buffer, spec, 0, uuid_subtype, 1)) {
buffer_free(buffer);
PyMem_Free(collection_name);
return NULL;
}
max_size = buffer_get_position(buffer) - before;
before = buffer_get_position(buffer);
if (!write_dict(state->_cbson, buffer, doc, check_keys, uuid_subtype, 1)) {
buffer_free(buffer);
PyMem_Free(collection_name);
return NULL;
}
cur_size = buffer_get_position(buffer) - before;
max_size = (cur_size > max_size) ? cur_size : max_size;
message_length = buffer_get_position(buffer) - length_location;
memcpy(buffer_get_buffer(buffer) + length_location, &message_length, 4);
if (safe) {
if (!add_last_error(self, buffer, request_id, collection_name,
collection_name_length, last_error_args)) {
buffer_free(buffer);
PyMem_Free(collection_name);
return NULL;
}
}
PyMem_Free(collection_name);
/* objectify buffer */
result = Py_BuildValue("i" BYTES_FORMAT_STRING "i", request_id,
buffer_get_buffer(buffer),
buffer_get_position(buffer),
max_size);
buffer_free(buffer);
return result;
}
static PyObject* _cbson_insert_message(PyObject* self, PyObject* args) {
/* NOTE just using a random number as the request_id */
struct module_state *state = GETSTATE(self);
int request_id = rand();
char* collection_name = NULL;
int collection_name_length;
PyObject* docs;
PyObject* doc;
PyObject* iterator;
int before, cur_size, max_size = 0;
int options = 0;
unsigned char check_keys;
unsigned char safe;
unsigned char continue_on_error;
unsigned char uuid_subtype;
PyObject* last_error_args;
buffer_t buffer;
int length_location, message_length;
PyObject* result;
if (!PyArg_ParseTuple(args, "et#ObbObb",
"utf-8",
&collection_name,
&collection_name_length,
&docs, &check_keys, &safe,
&last_error_args,
&continue_on_error, &uuid_subtype)) {
return NULL;
}
if (continue_on_error) {
options += 1;
}
buffer = buffer_new();
if (!buffer) {
PyErr_NoMemory();
PyMem_Free(collection_name);
return NULL;
}
// save space for message length
length_location = buffer_save_space(buffer, 4);
if (length_location == -1) {
PyMem_Free(collection_name);
PyErr_NoMemory();
return NULL;
}
if (!buffer_write_bytes(buffer, (const char*)&request_id, 4) ||
!buffer_write_bytes(buffer,
"\x00\x00\x00\x00"
"\xd2\x07\x00\x00",
8) ||
!buffer_write_bytes(buffer, (const char*)&options, 4) ||
!buffer_write_bytes(buffer,
collection_name,
collection_name_length + 1)) {
PyMem_Free(collection_name);
buffer_free(buffer);
return NULL;
}
iterator = PyObject_GetIter(docs);
if (iterator == NULL) {
PyObject* InvalidOperation = _error("InvalidOperation");
if (InvalidOperation) {
PyErr_SetString(InvalidOperation, "input is not iterable");
Py_DECREF(InvalidOperation);
}
buffer_free(buffer);
PyMem_Free(collection_name);
return NULL;
}
while ((doc = PyIter_Next(iterator)) != NULL) {
before = buffer_get_position(buffer);
if (!write_dict(state->_cbson, buffer, doc, check_keys, uuid_subtype, 1)) {
Py_DECREF(doc);
Py_DECREF(iterator);
buffer_free(buffer);
PyMem_Free(collection_name);
return NULL;
}
Py_DECREF(doc);
cur_size = buffer_get_position(buffer) - before;
max_size = (cur_size > max_size) ? cur_size : max_size;
}
Py_DECREF(iterator);
if (PyErr_Occurred()) {
buffer_free(buffer);
PyMem_Free(collection_name);
return NULL;
}
if (!max_size) {
PyObject* InvalidOperation = _error("InvalidOperation");
if (InvalidOperation) {
PyErr_SetString(InvalidOperation, "cannot do an empty bulk insert");
Py_DECREF(InvalidOperation);
}
buffer_free(buffer);
PyMem_Free(collection_name);
return NULL;
}
message_length = buffer_get_position(buffer) - length_location;
memcpy(buffer_get_buffer(buffer) + length_location, &message_length, 4);
if (safe) {
if (!add_last_error(self, buffer, request_id, collection_name,
collection_name_length, last_error_args)) {
buffer_free(buffer);
PyMem_Free(collection_name);
return NULL;
}
}
PyMem_Free(collection_name);
/* objectify buffer */
result = Py_BuildValue("i" BYTES_FORMAT_STRING "i", request_id,
buffer_get_buffer(buffer),
buffer_get_position(buffer),
max_size);
buffer_free(buffer);
return result;
}
static PyObject*
_cbson_do_batched_write_command(PyObject* self, PyObject* args) {
struct module_state *state = GETSTATE(self);
long max_bson_size;
long max_cmd_size;
long max_write_batch_size;
long idx_offset = 0;
int idx = 0;
int cmd_len_loc;
int lst_len_loc;
int ns_len;
int ordered;
char *ns = NULL;
PyObject* max_bson_size_obj;
PyObject* max_write_batch_size_obj;
PyObject* command;
PyObject* doc;
PyObject* docs;
PyObject* ctx;
PyObject* iterator;
PyObject* result;
PyObject* results;
PyObject* to_publish = NULL;
unsigned char op;
unsigned char check_keys;
codec_options_t options;
unsigned char empty = 1;
unsigned char errors = 0;
buffer_t buffer;
if (!PyArg_ParseTuple(args, "et#bOObO&O", "utf-8",
&ns, &ns_len, &op, &command, &docs, &check_keys,
convert_codec_options, &options,
&ctx)) {
return NULL;
}
max_bson_size_obj = PyObject_GetAttrString(ctx, "max_bson_size");
#if PY_MAJOR_VERSION >= 3
max_bson_size = PyLong_AsLong(max_bson_size_obj);
#else
max_bson_size = PyInt_AsLong(max_bson_size_obj);
#endif
Py_XDECREF(max_bson_size_obj);
if (max_bson_size == -1) {
destroy_codec_options(&options);
PyMem_Free(ns);
return NULL;
}
/*
* Max BSON object size + 16k - 2 bytes for ending NUL bytes
* XXX: This should come from the server - SERVER-10643
*/
max_cmd_size = max_bson_size + 16382;
max_write_batch_size_obj = PyObject_GetAttrString(ctx, "max_write_batch_size");
#if PY_MAJOR_VERSION >= 3
max_write_batch_size = PyLong_AsLong(max_write_batch_size_obj);
#else
max_write_batch_size = PyInt_AsLong(max_write_batch_size_obj);
#endif
Py_XDECREF(max_write_batch_size_obj);
if (max_write_batch_size == -1) {
destroy_codec_options(&options);
PyMem_Free(ns);
return NULL;
}
/* Default to True */
ordered = !((PyDict_GetItemString(command, "ordered")) == Py_False);
if (!(results = PyList_New(0))) {
destroy_codec_options(&options);
PyMem_Free(ns);
return NULL;
}
if (!(to_publish = PyList_New(0))) {
destroy_codec_options(&options);
PyMem_Free(ns);
Py_DECREF(results);
return NULL;
}
if (!(buffer = _command_buffer_new(ns, ns_len))) {
destroy_codec_options(&options);
PyMem_Free(ns);
Py_DECREF(results);
Py_DECREF(to_publish);
return NULL;
}
PyMem_Free(ns);
/* Position of command document length */
cmd_len_loc = buffer_get_position(buffer);
if (!write_dict(state->_cbson, buffer, command, 0,
&options, 0)) {
goto cmdfail;
}
/* Write type byte for array */
*(buffer_get_buffer(buffer) + (buffer_get_position(buffer) - 1)) = 0x4;
switch (op) {
case _INSERT:
{
if (!buffer_write_bytes(buffer, "documents\x00", 10))
goto cmdfail;
break;
}
case _UPDATE:
{
/* MongoDB does key validation for update. */
check_keys = 0;
if (!buffer_write_bytes(buffer, "updates\x00", 8))
goto cmdfail;
break;
}
case _DELETE:
{
/* Never check keys in a delete command. */
check_keys = 0;
if (!buffer_write_bytes(buffer, "deletes\x00", 8))
goto cmdfail;
break;
}
default:
{
PyObject* InvalidOperation = _error("InvalidOperation");
if (InvalidOperation) {
PyErr_SetString(InvalidOperation, "Unknown command");
Py_DECREF(InvalidOperation);
}
goto cmdfail;
}
}
/* Save space for list document */
lst_len_loc = buffer_save_space(buffer, 4);
if (lst_len_loc == -1) {
PyErr_NoMemory();
goto cmdfail;
}
iterator = PyObject_GetIter(docs);
if (iterator == NULL) {
PyObject* InvalidOperation = _error("InvalidOperation");
if (InvalidOperation) {
PyErr_SetString(InvalidOperation, "input is not iterable");
Py_DECREF(InvalidOperation);
}
goto cmdfail;
}
while ((doc = PyIter_Next(iterator)) != NULL) {
int sub_doc_begin = buffer_get_position(buffer);
int cur_doc_begin;
int cur_size;
int enough_data = 0;
int enough_documents = 0;
char key[16];
empty = 0;
INT2STRING(key, idx);
if (!buffer_write_bytes(buffer, "\x03", 1) ||
!buffer_write_bytes(buffer, key, (int)strlen(key) + 1)) {
goto cmditerfail;
}
cur_doc_begin = buffer_get_position(buffer);
if (!write_dict(state->_cbson, buffer, doc,
check_keys, &options, 1)) {
goto cmditerfail;
}
/* We have enough data, maybe send this batch. */
enough_data = (buffer_get_position(buffer) > max_cmd_size);
enough_documents = (idx >= max_write_batch_size);
if (enough_data || enough_documents) {
buffer_t new_buffer;
cur_size = buffer_get_position(buffer) - cur_doc_begin;
/* This single document is too large for the command. */
if (!idx) {
if (op == _INSERT) {
_set_document_too_large(cur_size, max_bson_size);
} else {
PyObject* DocumentTooLarge = _error("DocumentTooLarge");
if (DocumentTooLarge) {
/*
* There's nothing intelligent we can say
* about size for update and remove.
*/
PyErr_SetString(DocumentTooLarge,
"command document too large");
Py_DECREF(DocumentTooLarge);
}
}
goto cmditerfail;
}
if (!(new_buffer = buffer_new())) {
PyErr_NoMemory();
goto cmditerfail;
}
/* New buffer including the current overflow document */
if (!buffer_write_bytes(new_buffer,
(const char*)buffer_get_buffer(buffer), lst_len_loc + 5) ||
!buffer_write_bytes(new_buffer, "0\x00", 2) ||
!buffer_write_bytes(new_buffer,
(const char*)buffer_get_buffer(buffer) + cur_doc_begin, cur_size)) {
buffer_free(new_buffer);
goto cmditerfail;
}
/*
* Roll the existing buffer back to the beginning
* of the last document encoded.
*/
buffer_update_position(buffer, sub_doc_begin);
if (!buffer_write_bytes(buffer, "\x00\x00", 2)) {
buffer_free(new_buffer);
goto cmditerfail;
}
result = _send_write_command(ctx, buffer, lst_len_loc,
cmd_len_loc, &errors, to_publish);
buffer_free(buffer);
buffer = new_buffer;
if (!result)
goto cmditerfail;
#if PY_MAJOR_VERSION >= 3
result = Py_BuildValue("NN",
PyLong_FromLong(idx_offset), result);
#else
result = Py_BuildValue("NN",
PyInt_FromLong(idx_offset), result);
#endif
if (!result)
goto cmditerfail;
if (PyList_Append(results, result) < 0) {
Py_DECREF(result);
goto cmditerfail;
}
Py_DECREF(result);
if (errors && ordered) {
destroy_codec_options(&options);
Py_DECREF(iterator);
buffer_free(buffer);
return results;
}
idx_offset += idx;
idx = 0;
Py_DECREF(to_publish);
if (!(to_publish = PyList_New(0))) {
goto cmditerfail;
}
}
if (PyList_Append(to_publish, doc) < 0) {
goto cmditerfail;
}
Py_CLEAR(doc);
idx += 1;
}
Py_DECREF(iterator);
if (PyErr_Occurred()) {
goto cmdfail;
}
if (empty) {
PyObject* InvalidOperation = _error("InvalidOperation");
if (InvalidOperation) {
PyErr_SetString(InvalidOperation, "cannot do an empty bulk write");
Py_DECREF(InvalidOperation);
}
goto cmdfail;
}
if (!buffer_write_bytes(buffer, "\x00\x00", 2))
goto cmdfail;
result = _send_write_command(ctx, buffer, lst_len_loc,
cmd_len_loc, &errors, to_publish);
if (!result)
goto cmdfail;
#if PY_MAJOR_VERSION >= 3
result = Py_BuildValue("NN", PyLong_FromLong(idx_offset), result);
#else
result = Py_BuildValue("NN", PyInt_FromLong(idx_offset), result);
#endif
if (!result)
goto cmdfail;
buffer_free(buffer);
if (PyList_Append(results, result) < 0) {
Py_DECREF(result);
goto cmdfail;
}
Py_DECREF(result);
Py_DECREF(to_publish);
destroy_codec_options(&options);
return results;
cmditerfail:
Py_XDECREF(doc);
Py_DECREF(iterator);
cmdfail:
destroy_codec_options(&options);
Py_DECREF(results);
Py_XDECREF(to_publish);
buffer_free(buffer);
return NULL;
}
static PyObject* _cbson_do_batched_insert(PyObject* self, PyObject* args) {
struct module_state *state = GETSTATE(self);
/* NOTE just using a random number as the request_id */
int request_id = rand();
int send_safe, flags = 0;
int length_location, message_length;
int collection_name_length;
char* collection_name = NULL;
PyObject* docs;
PyObject* doc;
PyObject* iterator;
PyObject* ctx;
PyObject* last_error_args;
PyObject* result;
PyObject* max_bson_size_obj;
PyObject* max_message_size_obj;
PyObject* to_publish = NULL;
unsigned char check_keys;
unsigned char safe;
unsigned char continue_on_error;
codec_options_t options;
unsigned char empty = 1;
long max_bson_size;
long max_message_size;
buffer_t buffer;
PyObject *exc_type = NULL, *exc_value = NULL, *exc_trace = NULL;
if (!PyArg_ParseTuple(args, "et#ObbObO&O",
"utf-8",
&collection_name,
&collection_name_length,
&docs, &check_keys, &safe,
&last_error_args,
&continue_on_error,
convert_codec_options, &options,
&ctx)) {
return NULL;
}
if (continue_on_error) {
flags += 1;
}
/*
* If we are doing unacknowledged writes *and* continue_on_error
* is True it's pointless (and slower) to send GLE.
*/
send_safe = (safe || !continue_on_error);
max_bson_size_obj = PyObject_GetAttrString(ctx, "max_bson_size");
#if PY_MAJOR_VERSION >= 3
max_bson_size = PyLong_AsLong(max_bson_size_obj);
#else
max_bson_size = PyInt_AsLong(max_bson_size_obj);
#endif
Py_XDECREF(max_bson_size_obj);
if (max_bson_size == -1) {
destroy_codec_options(&options);
PyMem_Free(collection_name);
return NULL;
}
max_message_size_obj = PyObject_GetAttrString(ctx, "max_message_size");
#if PY_MAJOR_VERSION >= 3
max_message_size = PyLong_AsLong(max_message_size_obj);
#else
max_message_size = PyInt_AsLong(max_message_size_obj);
#endif
Py_XDECREF(max_message_size_obj);
if (max_message_size == -1) {
destroy_codec_options(&options);
PyMem_Free(collection_name);
return NULL;
}
buffer = buffer_new();
if (!buffer) {
destroy_codec_options(&options);
PyErr_NoMemory();
PyMem_Free(collection_name);
return NULL;
}
length_location = init_insert_buffer(buffer,
request_id,
flags,
collection_name,
collection_name_length);
if (length_location == -1) {
goto insertfail;
}
if (!(to_publish = PyList_New(0))) {
goto insertfail;
}
iterator = PyObject_GetIter(docs);
if (iterator == NULL) {
PyObject* InvalidOperation = _error("InvalidOperation");
if (InvalidOperation) {
PyErr_SetString(InvalidOperation, "input is not iterable");
Py_DECREF(InvalidOperation);
}
goto insertfail;
}
while ((doc = PyIter_Next(iterator)) != NULL) {
int before = buffer_get_position(buffer);
int cur_size;
if (!write_dict(state->_cbson, buffer, doc, check_keys,
&options, 1)) {
goto iterfail;
}
cur_size = buffer_get_position(buffer) - before;
if (cur_size > max_bson_size) {
/* If we've encoded anything send it before raising. */
if (!empty) {
buffer_update_position(buffer, before);
message_length = buffer_get_position(buffer) - length_location;
memcpy(buffer_get_buffer(buffer) + length_location,
&message_length, 4);
result = _send_insert(self, ctx, last_error_args, buffer,
collection_name, collection_name_length,
request_id, send_safe, &options,
to_publish);
if (!result)
goto iterfail;
Py_DECREF(result);
}
_set_document_too_large(cur_size, max_bson_size);
goto iterfail;
}
empty = 0;
/* We have enough data, send this batch. */
if (buffer_get_position(buffer) > max_message_size) {
int new_request_id = rand();
int message_start;
buffer_t new_buffer = buffer_new();
if (!new_buffer) {
PyErr_NoMemory();
goto iterfail;
}
message_start = init_insert_buffer(new_buffer,
new_request_id,
flags,
collection_name,
collection_name_length);
if (message_start == -1) {
buffer_free(new_buffer);
goto iterfail;
}
/* Copy the overflow encoded document into the new buffer. */
if (!buffer_write_bytes(new_buffer,
(const char*)buffer_get_buffer(buffer) + before, cur_size)) {
buffer_free(new_buffer);
goto iterfail;
}
/* Roll back to the beginning of this document. */
buffer_update_position(buffer, before);
message_length = buffer_get_position(buffer) - length_location;
memcpy(buffer_get_buffer(buffer) + length_location, &message_length, 4);
result = _send_insert(self, ctx, last_error_args, buffer,
collection_name, collection_name_length,
request_id, send_safe, &options, to_publish);
buffer_free(buffer);
buffer = new_buffer;
request_id = new_request_id;
length_location = message_start;
Py_DECREF(to_publish);
if (!(to_publish = PyList_New(0))) {
goto insertfail;
}
if (!result) {
PyObject *etype = NULL, *evalue = NULL, *etrace = NULL;
PyObject* OperationFailure;
PyErr_Fetch(&etype, &evalue, &etrace);
OperationFailure = _error("OperationFailure");
if (OperationFailure) {
if (PyErr_GivenExceptionMatches(etype, OperationFailure)) {
if (!safe || continue_on_error) {
Py_DECREF(OperationFailure);
if (!safe) {
/* We're doing unacknowledged writes and
* continue_on_error is False. Just return. */
Py_DECREF(etype);
Py_XDECREF(evalue);
Py_XDECREF(etrace);
Py_DECREF(to_publish);
Py_DECREF(iterator);
Py_DECREF(doc);
buffer_free(buffer);
PyMem_Free(collection_name);
Py_RETURN_NONE;
}
/* continue_on_error is True, store the error
* details to re-raise after the final batch */
Py_XDECREF(exc_type);
Py_XDECREF(exc_value);
Py_XDECREF(exc_trace);
exc_type = etype;
exc_value = evalue;
exc_trace = etrace;
if (PyList_Append(to_publish, doc) < 0) {
goto iterfail;
}
Py_CLEAR(doc);
continue;
}
}
Py_DECREF(OperationFailure);
}
/* This isn't OperationFailure, we couldn't
* import OperationFailure, or we are doing
* acknowledged writes. Re-raise immediately. */
PyErr_Restore(etype, evalue, etrace);
goto iterfail;
} else {
Py_DECREF(result);
}
}
if (PyList_Append(to_publish, doc) < 0) {
goto iterfail;
}
Py_CLEAR(doc);
}
Py_DECREF(iterator);
if (PyErr_Occurred()) {
goto insertfail;
}
if (empty) {
PyObject* InvalidOperation = _error("InvalidOperation");
if (InvalidOperation) {
PyErr_SetString(InvalidOperation, "cannot do an empty bulk insert");
Py_DECREF(InvalidOperation);
}
goto insertfail;
}
message_length = buffer_get_position(buffer) - length_location;
memcpy(buffer_get_buffer(buffer) + length_location, &message_length, 4);
/* Send the last (or only) batch */
result = _send_insert(self, ctx, last_error_args, buffer,
collection_name, collection_name_length,
request_id, safe, &options, to_publish);
Py_DECREF(to_publish);
PyMem_Free(collection_name);
buffer_free(buffer);
if (!result) {
Py_XDECREF(exc_type);
Py_XDECREF(exc_value);
Py_XDECREF(exc_trace);
return NULL;
} else {
Py_DECREF(result);
}
if (exc_type) {
/* Re-raise any previously stored exception
* due to continue_on_error being True */
PyErr_Restore(exc_type, exc_value, exc_trace);
return NULL;
}
Py_RETURN_NONE;
iterfail:
Py_XDECREF(doc);
Py_DECREF(iterator);
insertfail:
Py_XDECREF(exc_type);
Py_XDECREF(exc_value);
Py_XDECREF(exc_trace);
Py_XDECREF(to_publish);
buffer_free(buffer);
PyMem_Free(collection_name);
return NULL;
}
static PyObject* _cbson_insert_message(PyObject* self, PyObject* args) {
/* Used by the Bulk API to insert into pre-2.6 servers. Collection.insert
* uses _cbson_do_batched_insert. */
struct module_state *state = GETSTATE(self);
/* NOTE just using a random number as the request_id */
int request_id = rand();
char* collection_name = NULL;
int collection_name_length;
PyObject* docs;
PyObject* doc;
PyObject* iterator;
int before, cur_size, max_size = 0;
int flags = 0;
unsigned char check_keys;
unsigned char safe;
unsigned char continue_on_error;
codec_options_t options;
PyObject* last_error_args;
buffer_t buffer;
int length_location, message_length;
PyObject* result;
if (!PyArg_ParseTuple(args, "et#ObbObO&",
"utf-8",
&collection_name,
&collection_name_length,
&docs, &check_keys, &safe,
&last_error_args,
&continue_on_error,
convert_codec_options, &options)) {
return NULL;
}
if (continue_on_error) {
flags += 1;
}
buffer = buffer_new();
if (!buffer) {
PyErr_NoMemory();
destroy_codec_options(&options);
PyMem_Free(collection_name);
return NULL;
}
length_location = init_insert_buffer(buffer,
request_id,
flags,
collection_name,
collection_name_length);
if (length_location == -1) {
destroy_codec_options(&options);
PyMem_Free(collection_name);
buffer_free(buffer);
return NULL;
}
iterator = PyObject_GetIter(docs);
if (iterator == NULL) {
PyObject* InvalidOperation = _error("InvalidOperation");
if (InvalidOperation) {
PyErr_SetString(InvalidOperation, "input is not iterable");
Py_DECREF(InvalidOperation);
}
destroy_codec_options(&options);
buffer_free(buffer);
PyMem_Free(collection_name);
return NULL;
}
while ((doc = PyIter_Next(iterator)) != NULL) {
before = buffer_get_position(buffer);
if (!write_dict(state->_cbson, buffer, doc, check_keys,
&options, 1)) {
Py_DECREF(doc);
Py_DECREF(iterator);
destroy_codec_options(&options);
buffer_free(buffer);
PyMem_Free(collection_name);
return NULL;
}
Py_DECREF(doc);
cur_size = buffer_get_position(buffer) - before;
max_size = (cur_size > max_size) ? cur_size : max_size;
}
Py_DECREF(iterator);
if (PyErr_Occurred()) {
destroy_codec_options(&options);
buffer_free(buffer);
PyMem_Free(collection_name);
return NULL;
}
if (!max_size) {
PyObject* InvalidOperation = _error("InvalidOperation");
if (InvalidOperation) {
PyErr_SetString(InvalidOperation, "cannot do an empty bulk insert");
Py_DECREF(InvalidOperation);
}
destroy_codec_options(&options);
buffer_free(buffer);
PyMem_Free(collection_name);
return NULL;
}
message_length = buffer_get_position(buffer) - length_location;
memcpy(buffer_get_buffer(buffer) + length_location, &message_length, 4);
if (safe) {
if (!add_last_error(self, buffer, request_id, collection_name,
collection_name_length, &options, last_error_args)) {
destroy_codec_options(&options);
buffer_free(buffer);
PyMem_Free(collection_name);
return NULL;
}
}
PyMem_Free(collection_name);
/* objectify buffer */
result = Py_BuildValue("i" BYTES_FORMAT_STRING "i", request_id,
buffer_get_buffer(buffer),
buffer_get_position(buffer),
max_size);
destroy_codec_options(&options);
buffer_free(buffer);
return result;
}
static int write_element_allow_id(VALUE key, VALUE value, VALUE extra, int allow_id) {
buffer_t buffer = (buffer_t)NUM2LL(rb_ary_entry(extra, 0));
VALUE check_keys = rb_ary_entry(extra, 1);
if (TYPE(key) == T_SYMBOL) {
// TODO better way to do this... ?
key = rb_str_new2(rb_id2name(SYM2ID(key)));
}
if (TYPE(key) != T_STRING) {
buffer_free(buffer);
rb_raise(rb_eTypeError, "keys must be strings or symbols");
}
if (!allow_id && strcmp("_id", RSTRING_PTR(key)) == 0) {
return ST_CONTINUE;
}
if (check_keys == Qtrue) {
int i;
if (RSTRING_LEN(key) > 0 && RSTRING_PTR(key)[0] == '$') {
buffer_free(buffer);
rb_raise(InvalidName, "key must not start with '$'");
}
for (i = 0; i < RSTRING_LEN(key); i++) {
if (RSTRING_PTR(key)[i] == '.') {
buffer_free(buffer);
rb_raise(InvalidName, "key must not contain '.'");
}
}
}
switch(TYPE(value)) {
case T_BIGNUM:
case T_FIXNUM:
{
if (rb_funcall(value, rb_intern(">"), 1, LL2NUM(9223372036854775807LL)) == Qtrue ||
rb_funcall(value, rb_intern("<"), 1, LL2NUM(-9223372036854775808LL)) == Qtrue) {
buffer_free(buffer);
rb_raise(rb_eRangeError, "MongoDB can only handle 8-byte ints");
}
if (rb_funcall(value, rb_intern(">"), 1, INT2NUM(2147483647L)) == Qtrue ||
rb_funcall(value, rb_intern("<"), 1, INT2NUM(-2147483648L)) == Qtrue) {
long long ll_value;
write_name_and_type(buffer, key, 0x12);
ll_value = NUM2LL(value);
SAFE_WRITE(buffer, (char*)&ll_value, 8);
} else {
int int_value;
write_name_and_type(buffer, key, 0x10);
int_value = NUM2LL(value);
SAFE_WRITE(buffer, (char*)&int_value, 4);
}
break;
}
case T_TRUE:
{
write_name_and_type(buffer, key, 0x08);
SAFE_WRITE(buffer, &one, 1);
break;
}
case T_FALSE:
{
write_name_and_type(buffer, key, 0x08);
SAFE_WRITE(buffer, &zero, 1);
break;
}
case T_FLOAT:
{
double d = NUM2DBL(value);
write_name_and_type(buffer, key, 0x01);
SAFE_WRITE(buffer, (char*)&d, 8);
break;
}
case T_NIL:
{
write_name_and_type(buffer, key, 0x0A);
break;
}
case T_HASH:
{
write_name_and_type(buffer, key, 0x03);
write_doc(buffer, value, check_keys);
break;
}
case T_ARRAY:
{
buffer_position length_location, start_position, obj_length;
int items, i;
VALUE* values;
write_name_and_type(buffer, key, 0x04);
start_position = buffer_get_position(buffer);
// save space for length
length_location = buffer_save_space(buffer, 4);
if (length_location == -1) {
rb_raise(rb_eNoMemError, "failed to allocate memory in buffer.c");
}
items = RARRAY_LEN(value);
values = RARRAY_PTR(value);
for(i = 0; i < items; i++) {
char* name;
VALUE key;
INT2STRING(&name, i);
key = rb_str_new2(name);
write_element(key, values[i], pack_extra(buffer, check_keys));
free(name);
}
// write null byte and fill in length
SAFE_WRITE(buffer, &zero, 1);
obj_length = buffer_get_position(buffer) - start_position;
SAFE_WRITE_AT_POS(buffer, length_location, (const char*)&obj_length, 4);
break;
}
case T_STRING:
{
if (strcmp(rb_class2name(RBASIC(value)->klass),
"Mongo::Code") == 0) {
buffer_position length_location, start_position, total_length;
int length;
write_name_and_type(buffer, key, 0x0F);
start_position = buffer_get_position(buffer);
length_location = buffer_save_space(buffer, 4);
if (length_location == -1) {
rb_raise(rb_eNoMemError, "failed to allocate memory in buffer.c");
}
length = RSTRING_LEN(value) + 1;
SAFE_WRITE(buffer, (char*)&length, 4);
SAFE_WRITE(buffer, RSTRING_PTR(value), length - 1);
SAFE_WRITE(buffer, &zero, 1);
write_doc(buffer, rb_funcall(value, rb_intern("scope"), 0), Qfalse);
total_length = buffer_get_position(buffer) - start_position;
SAFE_WRITE_AT_POS(buffer, length_location, (const char*)&total_length, 4);
break;
} else {
int length;
write_name_and_type(buffer, key, 0x02);
value = TO_UTF8(value);
length = RSTRING_LEN(value) + 1;
SAFE_WRITE(buffer, (char*)&length, 4);
write_utf8(buffer, value);
SAFE_WRITE(buffer, &zero, 1);
break;
}
}
case T_SYMBOL:
{
const char* str_value = rb_id2name(SYM2ID(value));
int length = strlen(str_value) + 1;
write_name_and_type(buffer, key, 0x0E);
SAFE_WRITE(buffer, (char*)&length, 4);
SAFE_WRITE(buffer, str_value, length);
break;
}
case T_OBJECT:
{
// TODO there has to be a better way to do these checks...
const char* cls = rb_class2name(RBASIC(value)->klass);
if (strcmp(cls, "Mongo::Binary") == 0 ||
strcmp(cls, "ByteBuffer") == 0) {
const char subtype = strcmp(cls, "ByteBuffer") ?
(const char)FIX2INT(rb_funcall(value, rb_intern("subtype"), 0)) : 2;
VALUE string_data = rb_funcall(value, rb_intern("to_s"), 0);
int length = RSTRING_LEN(string_data);
write_name_and_type(buffer, key, 0x05);
if (subtype == 2) {
const int other_length = length + 4;
SAFE_WRITE(buffer, (const char*)&other_length, 4);
SAFE_WRITE(buffer, &subtype, 1);
}
SAFE_WRITE(buffer, (const char*)&length, 4);
if (subtype != 2) {
SAFE_WRITE(buffer, &subtype, 1);
}
SAFE_WRITE(buffer, RSTRING_PTR(string_data), length);
break;
}
if (strcmp(cls, "Mongo::ObjectID") == 0) {
VALUE as_array = rb_funcall(value, rb_intern("to_a"), 0);
int i;
write_name_and_type(buffer, key, 0x07);
for (i = 0; i < 12; i++) {
char byte = (char)FIX2INT(RARRAY_PTR(as_array)[i]);
SAFE_WRITE(buffer, &byte, 1);
}
break;
}
if (strcmp(cls, "Mongo::DBRef") == 0) {
buffer_position length_location, start_position, obj_length;
VALUE ns, oid;
write_name_and_type(buffer, key, 0x03);
start_position = buffer_get_position(buffer);
// save space for length
length_location = buffer_save_space(buffer, 4);
if (length_location == -1) {
rb_raise(rb_eNoMemError, "failed to allocate memory in buffer.c");
}
ns = rb_funcall(value, rb_intern("namespace"), 0);
write_element(rb_str_new2("$ref"), ns, pack_extra(buffer, Qfalse));
oid = rb_funcall(value, rb_intern("object_id"), 0);
write_element(rb_str_new2("$id"), oid, pack_extra(buffer, Qfalse));
// write null byte and fill in length
SAFE_WRITE(buffer, &zero, 1);
obj_length = buffer_get_position(buffer) - start_position;
SAFE_WRITE_AT_POS(buffer, length_location, (const char*)&obj_length, 4);
break;
}
}
case T_DATA:
{
// TODO again, is this really the only way to do this?
const char* cls = rb_class2name(RBASIC(value)->klass);
if (strcmp(cls, "Time") == 0) {
double t = NUM2DBL(rb_funcall(value, rb_intern("to_f"), 0));
long long time_since_epoch = (long long)round(t * 1000);
write_name_and_type(buffer, key, 0x09);
SAFE_WRITE(buffer, (const char*)&time_since_epoch, 8);
break;
}
}
case T_REGEXP:
{
int length = RREGEXP_SRC_LEN(value);
char* pattern = (char*)RREGEXP_SRC_PTR(value);
long flags = RREGEXP(value)->ptr->options;
VALUE has_extra;
write_name_and_type(buffer, key, 0x0B);
SAFE_WRITE(buffer, pattern, length);
SAFE_WRITE(buffer, &zero, 1);
if (flags & IGNORECASE) {
char ignorecase = 'i';
SAFE_WRITE(buffer, &ignorecase, 1);
}
if (flags & MULTILINE) {
char multiline = 'm';
SAFE_WRITE(buffer, &multiline, 1);
}
if (flags & EXTENDED) {
char extended = 'x';
SAFE_WRITE(buffer, &extended, 1);
}
has_extra = rb_funcall(value, rb_intern("respond_to?"), 1, rb_str_new2("extra_options_str"));
if (TYPE(has_extra) == T_TRUE) {
VALUE extra = rb_funcall(value, rb_intern("extra_options_str"), 0);
buffer_position old_position = buffer_get_position(buffer);
SAFE_WRITE(buffer, RSTRING_PTR(extra), RSTRING_LEN(extra));
qsort(buffer_get_buffer(buffer) + old_position, RSTRING_LEN(extra), sizeof(char), cmp_char);
}
SAFE_WRITE(buffer, &zero, 1);
break;
}
default:
{
buffer_free(buffer);
rb_raise(rb_eTypeError, "no c encoder for this type yet (%d)", TYPE(value));
break;
}
}
return ST_CONTINUE;
}