def _write(self, atoms, key_value_pairs, data, id): ext_tables = key_value_pairs.pop("external_tables", {}) Database._write(self, atoms, key_value_pairs, data) mtime = now() encode = self.encode blob = self.blob if not isinstance(atoms, AtomsRow): row = AtomsRow(atoms) row.ctime = mtime row.user = os.getenv('USER') else: row = atoms # Extract the external tables from AtomsRow names = self._get_external_table_names() for name in names: new_table = row.get(name, {}) if new_table: ext_tables[name] = new_table if not id and not key_value_pairs and not ext_tables: key_value_pairs = row.key_value_pairs for k, v in ext_tables.items(): dtype = self._guess_type(v) self._create_table_if_not_exists(k, dtype) constraints = row._constraints if constraints: if isinstance(constraints, list): constraints = encode(constraints) else: constraints = None values = (row.unique_id, row.ctime, mtime, row.user, blob(row.numbers), blob(row.positions), blob(row.cell), int(np.dot(row.pbc, [1, 2, 4])), blob(row.get('initial_magmoms')), blob(row.get('initial_charges')), blob(row.get('masses')), blob(row.get('tags')), blob(row.get('momenta')), constraints) if 'calculator' in row: values += (row.calculator, encode(row.calculator_parameters)) else: values += (None, None) if not data: data = row._data with self.managed_connection() as con: if not isinstance(data, (str, bytes)): data = encode(data, binary=self.version >= 9) values += (row.get('energy'), row.get('free_energy'), blob(row.get('forces')), blob(row.get('stress')), blob(row.get('dipole')), blob(row.get('magmoms')), row.get('magmom'), blob(row.get('charges')), encode(key_value_pairs), data, len(row.numbers), float_if_not_none(row.get('fmax')), float_if_not_none(row.get('smax')), float_if_not_none(row.get('volume')), float(row.mass), float(row.charge)) cur = con.cursor() if id is None: q = self.default + ', ' + ', '.join('?' * len(values)) cur.execute('INSERT INTO systems VALUES ({})'.format(q), values) id = self.get_last_id(cur) else: self._delete(cur, [id], [ 'keys', 'text_key_values', 'number_key_values', 'species' ]) q = ', '.join(name + '=?' for name in self.columnnames[1:]) cur.execute('UPDATE systems SET {} WHERE id=?'.format(q), values + (id, )) count = row.count_atoms() if count: species = [(atomic_numbers[symbol], n, id) for symbol, n in count.items()] cur.executemany('INSERT INTO species VALUES (?, ?, ?)', species) text_key_values = [] number_key_values = [] for key, value in key_value_pairs.items(): if isinstance(value, (numbers.Real, np.bool_)): number_key_values.append([key, float(value), id]) else: assert isinstance(value, str) text_key_values.append([key, value, id]) cur.executemany('INSERT INTO text_key_values VALUES (?, ?, ?)', text_key_values) cur.executemany('INSERT INTO number_key_values VALUES (?, ?, ?)', number_key_values) cur.executemany('INSERT INTO keys VALUES (?, ?)', [(key, id) for key in key_value_pairs]) # Insert entries in the valid tables for tabname in ext_tables.keys(): entries = ext_tables[tabname] entries['id'] = id self._insert_in_external_table(cur, name=tabname, entries=ext_tables[tabname]) return id
def _write(self, atoms, key_value_pairs, data, id): ext_tables = key_value_pairs.pop("external_tables", {}) Database._write(self, atoms, key_value_pairs, data) encode = self.encode con = self.connection or self._connect() self._initialize(con) cur = con.cursor() mtime = now() blob = self.blob text_key_values = [] number_key_values = [] if not isinstance(atoms, AtomsRow): row = AtomsRow(atoms) row.ctime = mtime row.user = os.getenv('USER') else: row = atoms # Extract the external tables from AtomsRow names = self._get_external_table_names(db_con=con) for name in names: new_table = row.get(name, {}) if new_table: ext_tables[name] = new_table if id: self._delete( cur, [id], ['keys', 'text_key_values', 'number_key_values', 'species']) else: if not key_value_pairs: key_value_pairs = row.key_value_pairs constraints = row._constraints if constraints: if isinstance(constraints, list): constraints = encode(constraints) else: constraints = None values = (row.unique_id, row.ctime, mtime, row.user, blob(row.numbers), blob(row.positions), blob(row.cell), int(np.dot(row.pbc, [1, 2, 4])), blob(row.get('initial_magmoms')), blob(row.get('initial_charges')), blob(row.get('masses')), blob(row.get('tags')), blob(row.get('momenta')), constraints) if 'calculator' in row: values += (row.calculator, encode(row.calculator_parameters)) else: values += (None, None) if not data: data = row._data if not isinstance(data, basestring): data = encode(data) values += (row.get('energy'), row.get('free_energy'), blob(row.get('forces')), blob(row.get('stress')), blob(row.get('dipole')), blob(row.get('magmoms')), row.get('magmom'), blob(row.get('charges')), encode(key_value_pairs), data, len(row.numbers), float_if_not_none(row.get('fmax')), float_if_not_none(row.get('smax')), float_if_not_none(row.get('volume')), float(row.mass), float(row.charge)) if id is None: q = self.default + ', ' + ', '.join('?' * len(values)) cur.execute('INSERT INTO systems VALUES ({})'.format(q), values) id = self.get_last_id(cur) else: q = ', '.join(name + '=?' for name in self.columnnames[1:]) cur.execute('UPDATE systems SET {} WHERE id=?'.format(q), values + (id, )) count = row.count_atoms() if count: species = [(atomic_numbers[symbol], n, id) for symbol, n in count.items()] cur.executemany('INSERT INTO species VALUES (?, ?, ?)', species) text_key_values = [] number_key_values = [] for key, value in key_value_pairs.items(): if isinstance(value, (numbers.Real, np.bool_)): number_key_values.append([key, float(value), id]) else: assert isinstance(value, basestring) text_key_values.append([key, value, id]) cur.executemany('INSERT INTO text_key_values VALUES (?, ?, ?)', text_key_values) cur.executemany('INSERT INTO number_key_values VALUES (?, ?, ?)', number_key_values) cur.executemany('INSERT INTO keys VALUES (?, ?)', [(key, id) for key in key_value_pairs]) # Update external tables valid_entries = [] for k, v in ext_tables.items(): try: # Guess the type of the value dtype = self._guess_type(v) self._create_table_if_not_exists(k, dtype, db_con=con) v["id"] = id valid_entries.append(k) except ValueError as exc: # Close the connection without committing if self.connection is None: con.close() # Raise error again raise ValueError(exc) # Insert entries in the valid tables for tabname in valid_entries: try: self._insert_in_external_table(cur, name=tabname, entries=ext_tables[tabname]) except ValueError as exc: # Close the connection without committing if self.connection is None: con.close() # Raise the error again raise ValueError(exc) if self.connection is None: con.commit() con.close() return id
def _write(self, atoms, key_value_pairs, data, id): Database._write(self, atoms, key_value_pairs, data) encode = self.encode con = self.connection or self._connect() self._initialize(con) cur = con.cursor() mtime = now() blob = self.blob text_key_values = [] number_key_values = [] if not isinstance(atoms, AtomsRow): row = AtomsRow(atoms) row.ctime = mtime row.user = os.getenv('USER') else: row = atoms if id: self._delete(cur, [id], ['keys', 'text_key_values', 'number_key_values', 'species']) else: if not key_value_pairs: key_value_pairs = row.key_value_pairs constraints = row._constraints if constraints: if isinstance(constraints, list): constraints = encode(constraints) else: constraints = None values = (row.unique_id, row.ctime, mtime, row.user, blob(row.numbers), blob(row.positions), blob(row.cell), int(np.dot(row.pbc, [1, 2, 4])), blob(row.get('initial_magmoms')), blob(row.get('initial_charges')), blob(row.get('masses')), blob(row.get('tags')), blob(row.get('momenta')), constraints) if 'calculator' in row: values += (row.calculator, encode(row.calculator_parameters)) else: values += (None, None) if not data: data = row._data if not isinstance(data, basestring): data = encode(data) values += (row.get('energy'), row.get('free_energy'), blob(row.get('forces')), blob(row.get('stress')), blob(row.get('dipole')), blob(row.get('magmoms')), row.get('magmom'), blob(row.get('charges')), encode(key_value_pairs), data, len(row.numbers), float_if_not_none(row.get('fmax')), float_if_not_none(row.get('smax')), float_if_not_none(row.get('volume')), float(row.mass), float(row.charge)) if id is None: q = self.default + ', ' + ', '.join('?' * len(values)) cur.execute('INSERT INTO systems VALUES ({})'.format(q), values) id = self.get_last_id(cur) else: q = ', '.join(name + '=?' for name in self.columnnames[1:]) cur.execute('UPDATE systems SET {} WHERE id=?'.format(q), values + (id,)) count = row.count_atoms() if count: species = [(atomic_numbers[symbol], n, id) for symbol, n in count.items()] cur.executemany('INSERT INTO species VALUES (?, ?, ?)', species) text_key_values = [] number_key_values = [] for key, value in key_value_pairs.items(): if isinstance(value, (numbers.Real, np.bool_)): number_key_values.append([key, float(value), id]) else: assert isinstance(value, basestring) text_key_values.append([key, value, id]) cur.executemany('INSERT INTO text_key_values VALUES (?, ?, ?)', text_key_values) cur.executemany('INSERT INTO number_key_values VALUES (?, ?, ?)', number_key_values) cur.executemany('INSERT INTO keys VALUES (?, ?)', [(key, id) for key in key_value_pairs]) if self.connection is None: con.commit() con.close() return id
def _write(self, atoms, key_value_pairs, data): Database._write(self, atoms, key_value_pairs, data) con = self.connection or self._connect() self._initialize(con) cur = con.cursor() id = None if not isinstance(atoms, AtomsRow): row = AtomsRow(atoms) row.ctime = mtime = now() row.user = os.getenv('USER') else: row = atoms cur.execute('SELECT id FROM systems WHERE unique_id=?', (row.unique_id,)) results = cur.fetchall() if results: id = results[0][0] self._delete(cur, [id], ['keys', 'text_key_values', 'number_key_values']) mtime = now() constraints = row._constraints if constraints: if isinstance(constraints, list): constraints = encode(constraints) else: constraints = None values = (row.unique_id, row.ctime, mtime, row.user, blob(row.numbers), blob(row.positions), blob(row.cell), int(np.dot(row.pbc, [1, 2, 4])), blob(row.get('initial_magmoms')), blob(row.get('initial_charges')), blob(row.get('masses')), blob(row.get('tags')), blob(row.get('momenta')), constraints) if 'calculator' in row: if not isinstance(row.calculator_parameters, basestring): row.calculator_parameters = encode(row.calculator_parameters) values += (row.calculator, row.calculator_parameters) else: values += (None, None) if key_value_pairs is None: key_value_pairs = row.key_value_pairs if not data: data = row._data if not isinstance(data, basestring): data = encode(data) values += (row.get('energy'), row.get('free_energy'), blob(row.get('forces')), blob(row.get('stress')), blob(row.get('dipole')), blob(row.get('magmoms')), row.get('magmom'), blob(row.get('charges')), encode(key_value_pairs), data, len(row.numbers), float_if_not_none(row.get('fmax')), float_if_not_none(row.get('smax')), float_if_not_none(row.get('volume')), float(row.mass), float(row.charge)) if id is None: q = self.default + ', ' + ', '.join('?' * len(values)) cur.execute('INSERT INTO systems VALUES ({0})'.format(q), values) else: q = ', '.join(line.split()[0].lstrip() + '=?' for line in init_statements[0].splitlines()[2:]) cur.execute('UPDATE systems SET {0} WHERE id=?'.format(q), values + (id,)) if id is None: id = self.get_last_id(cur) count = row.count_atoms() if count: species = [(atomic_numbers[symbol], n, id) for symbol, n in count.items()] cur.executemany('INSERT INTO species VALUES (?, ?, ?)', species) text_key_values = [] number_key_values = [] for key, value in key_value_pairs.items(): if isinstance(value, (float, int)): number_key_values.append([key, float(value), id]) else: assert isinstance(value, basestring) text_key_values.append([key, value, id]) cur.executemany('INSERT INTO text_key_values VALUES (?, ?, ?)', text_key_values) cur.executemany('INSERT INTO number_key_values VALUES (?, ?, ?)', number_key_values) cur.executemany('INSERT INTO keys VALUES (?, ?)', [(key, id) for key in key_value_pairs]) if self.connection is None: con.commit() con.close() return id