diff --git a/src/procrusql/__init__.py b/src/procrusql/__init__.py index a0994f7..b8d722f 100644 --- a/src/procrusql/__init__.py +++ b/src/procrusql/__init__.py @@ -48,11 +48,12 @@ class Node: self.order = order class HaveData(Node): - def __init__(self, name, depends, table, key, extra): + def __init__(self, name, depends, table, key, extra, schema="public"): super().__init__(name, depends) self.table = table self.key = key self.extra = extra + self.schema = schema def check(self): log_check.info("Checking %s", self.name) @@ -63,8 +64,9 @@ class HaveData(Node): ] key_check = sql.SQL(" and ").join(key_checks) q = sql.SQL( - "select * from {table} where {key_check}" + "select * from {schema}.{table} where {key_check}" ).format( + schema=sql.Identifier(self.schema), table=sql.Identifier(self.table), key_check=key_check ) @@ -78,8 +80,9 @@ class HaveData(Node): if self.result[0][c] != self.extra[c]: log_action.info("Updating %s: %s <- %s", key_values, c, self.extra[c]) q = sql.SQL( - "update {table} set {column}={placeholder} where {key_check}" + "update {schema}.{table} set {column}={placeholder} where {key_check}" ).format( + schema=sql.Identifier(self.schema), table=sql.Identifier(self.table), column=sql.Identifier(c), placeholder=sql.Placeholder(), @@ -96,8 +99,9 @@ class HaveData(Node): columns = list(self.key.keys()) + list(self.extra.keys()) values = key_values + extra_values q = sql.SQL( - "insert into {table}({columns}) values({placeholders}) returning *" + "insert into {schema}.{table}({columns}) values({placeholders}) returning *" ).format( + schema=sql.Identifier(self.schema), table=sql.Identifier(self.table), columns=sql.SQL(", ").join([sql.Identifier(x) for x in columns]), placeholders=sql.SQL(", ").join([sql.Placeholder() for x in columns]),