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)