def create(columns, values, get_objects=False, get_primary_keys=False): """Create a large number of objects efficiently. :param columns: The Storm columns to insert values into. Must be from a single class. :param values: A list of lists of values for the columns. :param get_objects: Return the created objects. :param get_primary_keys: Return the created primary keys. :return: A list of the created objects if get_created, otherwise None. """ # Flatten Reference faux-columns into their primary keys. db_cols = list(chain.from_iterable(map(dbify_column, columns))) clses = set(col.cls for col in db_cols) if len(clses) != 1: raise ValueError( "The Storm columns to insert values into must be from a single " "class.") if get_objects and get_primary_keys: raise ValueError( "get_objects and get_primary_keys are mutually exclusive.") if len(values) == 0: return [] if (get_objects or get_primary_keys) else None [cls] = clses primary_key = get_cls_info(cls).primary_key # Mangle our value list into compilable values. Normal columns just # get passed through the variable factory, while References get # squashed into primary key variables. db_values = [ list( chain.from_iterable( dbify_value(col, val) for col, val in zip(columns, value))) for value in values ] if get_objects or get_primary_keys: result = IStore(cls).execute( Returning( Insert(db_cols, values=db_values, primary_columns=primary_key))) keys = map(itemgetter(0), result) if len(primary_key) == 1 else result if get_objects: return load(cls, keys) else: return list(keys) else: IStore(cls).execute(Insert(db_cols, values=db_values)) return None
def test_execute_insert_returning_without_variables(self): """Without primary_variables, the RETURNING system won't be used.""" column1 = Column("id1", "returning_test") insert = Insert({}, primary_columns=(column1, )) self.connection.execute(insert) result = self.connection.execute("SELECT * FROM returning_test") self.assertEquals(result.get_one(), (123, 456))
def test_execute_insert_returning_without_columns(self): """Without primary_columns, the RETURNING system won't be used.""" column1 = Column("id1", "returning_test") variable1 = IntVariable() insert = Insert({column1: 123}, primary_variables=(variable1, )) self.connection.execute(insert) self.assertFalse(variable1.is_defined()) result = self.connection.execute("SELECT * FROM returning_test") self.assertEquals(result.get_one(), (123, 456))
def test_wb_execute_insert_returning_not_used_with_old_postgres(self): """Shouldn't try to use RETURNING with PostgreSQL < 8.2.""" column1 = Column("id1", "returning_test") column2 = Column("id2", "returning_test") variable1 = IntVariable() variable2 = IntVariable() insert = Insert({}, primary_columns=(column1, column2), primary_variables=(variable1, variable2)) self.database._version = 80109 self.connection.execute(insert) self.assertFalse(variable1.is_defined()) self.assertFalse(variable2.is_defined()) result = self.connection.execute("SELECT * FROM returning_test") self.assertEquals(result.get_one(), (123, 456))
def test_execute_insert_returning(self): if self.database._version < 80200: return # Can't run this test with old PostgreSQL versions. column1 = Column("id1", "returning_test") column2 = Column("id2", "returning_test") variable1 = IntVariable() variable2 = IntVariable() insert = Insert({}, primary_columns=(column1, column2), primary_variables=(variable1, variable2)) self.connection.execute(insert) self.assertTrue(variable1.is_defined()) self.assertTrue(variable2.is_defined()) self.assertEquals(variable1.get(), 123) self.assertEquals(variable2.get(), 456) result = self.connection.execute("SELECT * FROM returning_test") self.assertEquals(result.get_one(), (123, 456))
def test_execute_insert_auto_increment_primary_key(self): id_column = Column("id", "test") id_variable = IntVariable() title_column = Column("title", "test") title_variable = UnicodeVariable(u"testing") # This is not part of the table. It is just used to show that # only one primary key variable is set from the insert ID. dummy_column = Column("dummy", "test") dummy_variable = IntVariable() insert = Insert({title_column: title_variable}, primary_columns=(id_column, dummy_column), primary_variables=(id_variable, dummy_variable)) self.connection.execute(insert) self.assertTrue(id_variable.is_defined()) self.assertFalse(dummy_variable.is_defined()) # The newly inserted row should have the maximum id value for # the table. result = self.connection.execute("SELECT MAX(id) FROM test") self.assertEqual(result.get_one()[0], id_variable.get())
def test_returning_column_context(self): column2 = TrackContext() insert = Insert({column1: elem1}, table1, primary_columns=column2) compile(Returning(insert)) self.assertEquals(column2.context, COLUMN)