def store_using_tmp(cursor, tmp_table, table, column_names, modified): """ Store the data using the PostgreSQL specific COPY FROM command and a temporary table. The temporary table is joined against the target table to make sure only new records are inserted. """ all_column_names = ['entity_id', 'timestamp'] all_column_names.extend(column_names) tmp_column_names = ", ".join('tmp."{0}"'.format(name) for name in all_column_names) dest_column_names = ", ".join('"{0}"'.format(name) for name in all_column_names) insert_query = ( "INSERT INTO {table} ({dest_columns}) " "SELECT {tmp_columns} FROM {tmp_table} AS tmp " "LEFT JOIN {table} ON " "tmp.\"timestamp\" = {table}.\"timestamp\" " "AND tmp.entity_id = {table}.entity_id " "WHERE {table}.entity_id IS NULL").format( table=table.render(), dest_columns=dest_column_names, tmp_columns=tmp_column_names, tmp_table=tmp_table.render()) try: cursor.execute(insert_query) except psycopg2.Error as exc: raise translate_postgresql_exception(exc)
def store_batch_insert(cursor, table, datapackage, modified): column_names = ["entity_id", "timestamp", "modified"] column_names.extend(datapackage.trend_names) dest_column_names = ",".join('"{0}"'.format(column_name) for column_name in column_names) parameters = ", ".join(["%s"] * len(column_names)) query = ( "INSERT INTO {0} ({1}) " "VALUES ({2})").format(table.render(), dest_column_names, parameters) rows = [(entity_id, datapackage.timestamp, modified) + tuple(values) for entity_id, values in datapackage.rows] logging.debug(cursor.mogrify(query, first(rows))) try: cursor.executemany(query, rows) except psycopg2.DatabaseError as exc: m = str(exc) if m.find("violates check constraint") > -1: print(cursor.mogrify(query, first(rows))) print(m) raise translate_postgresql_exception(exc)
def store_copy_from(cursor, table, datapackage, modified): """ Store the data using the PostgreSQL specific COPY FROM command :param conn: DBAPI2 database connection :param table: Name of table, including schema :param datapackage: A DataPackage object """ copy_from_file = create_copy_from_file(datapackage.timestamp, modified, datapackage.rows) copy_from_query = create_copy_from_query(table, datapackage.trend_names) logging.debug(copy_from_query) try: try: cursor.copy_expert(copy_from_query, copy_from_file) except Exception as exc: logging.debug(exc) raise exc except psycopg2.DatabaseError as exc: if exc.pgcode is None and str(exc).find("no COPY in progress") != -1: # Might happen after database connection loss raise NoCopyInProgress() else: raise translate_postgresql_exception(exc)
def execute(self, cursor, args=None): """Execute wrapped query with provided cursor and arguments.""" try: cursor.execute(self.sql, args) except psycopg2.DatabaseError as exc: raise translate_postgresql_exception(exc) return cursor
def execute(self, cursor): args = (self.type.id, self.timestamp) try: cursor.callproc("materialization.materialize", args) except psycopg2.ProgrammingError as exc: raise translate_postgresql_exception(exc) return cursor.fetchone()
def test_translate_postgresql_exception(): """ The translated exception should be NoSuchTable """ query = 'SELECT 1 FROM "non-existing-table"' with closing(connect()) as conn: with closing(conn.cursor()) as cursor: try: cursor.execute(query) except psycopg2.DatabaseError as exc: raise translate_postgresql_exception(exc)
def create_column(cursor, table, column_name, data_type): """ Create a new column with matching datatype for the specified trend. """ table_ref = table.render() query = ("ALTER TABLE {0} " 'ADD COLUMN "{1}" {2}').format(table_ref, column_name, data_type) try: cursor.execute(query) except psycopg2.ProgrammingError as exc: raise translate_postgresql_exception(exc)
def create(self, cursor): query = ( "SELECT trend.create_partition(trendstore, %s) " "FROM trend.trendstore " "WHERE id = %s") args = self.index, self.trendstore.id try: cursor.execute(query, args) except psycopg2.IntegrityError: raise DuplicateTable() except psycopg2.ProgrammingError as exc: raise translate_postgresql_exception(exc)
def store_using_update(cursor, tmp_table, table, column_names, modified): set_columns = ", ".join("\"{0}\"={1}.\"{0}\"".format(name, tmp_table.render()) for name in column_names) update_query = ( 'UPDATE {0} SET modified=greatest(%s, {0}.modified), {1} ' 'FROM {2} ' 'WHERE {0}.entity_id={2}.entity_id AND {0}."timestamp"={2}."timestamp"' ).format(table.render(), set_columns, tmp_table.render()) args = (modified, ) try: cursor.execute(update_query, args) except psycopg2.DatabaseError as exc: raise translate_postgresql_exception(exc)
def delete_by_entityselection(cursor, table, timestamp, entityselection): """ Delete rows from table for a specific timestamp and entity_ids in entityselection """ entityselection.create_temp_table(cursor, "entity_filter") delete_query = ( "DELETE FROM {} d " "USING entity_filter f " "WHERE d.timestamp = %s AND f.entity_id = d.entity_id" ).format(table.render()) args = (timestamp,) logging.debug(cursor.mogrify(delete_query, args)) try: cursor.execute(delete_query, args) except psycopg2.DatabaseError as exc: raise translate_postgresql_exception(exc)