allow replace option for PostgresDbOutput

stetl/outputs/dboutput.py

@@ -7,6 +7,7 @@
 from stetl.output import Output
 from stetl.util import Util
 from stetl.packet import FORMAT
+from stetl.component import Config
 from stetl.postgis import PostGIS
 
 log = Util.get_log('dboutput')
@@ -31,7 +32,43 @@ class PostgresDbOutput(DbOutput):
 
     consumes=FORMAT.string
     """
-
+    # Start attribute config meta
+    @Config(ptype=str, required=True, default=None)
+    def database(self):
+        """
+        Database name.
+        """
+        pass
+
+    @Config(ptype=str, required=False, default=None)
+    def user(self):
+        """
+        DB User name.
+        """
+        pass
+
+    @Config(ptype=str, required=False, default=None)
+    def password(self):
+        """
+        DB Password for user.
+        """
+        pass
+
+    @Config(ptype=str, required=False, default=None)
+    def host(self):
+        """
+        Hostname for DB.
+        """
+        pass
+
+    @Config(ptype=str, required=False, default='public')
+    def schema(self):
+        """
+        Postgres schema name for DB.
+        """
+        pass
+
+    # End attribute config meta
     def __init__(self, configdict, section):
         DbOutput.__init__(self, configdict, section, consumes=FORMAT.string)
 
@@ -50,15 +87,41 @@ class PostgresInsertOutput(PostgresDbOutput):
     Output by inserting single record into Postgres database.
     Input is a record (Python dic structure) or a Python list of dicts (records).
     Creates an INSERT for Postgres to insert each single record.
+    When the "replace" parameter is True, any existing record keyed by "key" is
+    attempted to be deleted first.
+
+    NB a constraint is that each record needs to contain all values as
+    an INSERT query is built once for the columns in the first record.
 
     consumes=FORMAT.record
     """
+    # Start attribute config meta
+    @Config(ptype=str, required=False, default='public')
+    def table(self):
+        """
+        Table for inserts.
+        """
+        pass
+
+    @Config(ptype=bool, required=False, default=False)
+    def replace(self):
+        """
+        Replace record if exists?
+        """
+        pass
+
+    @Config(ptype=str, required=False, default=None)
+    def key(self):
+        """
+        The key column name of the table, required when replacing records.
+        """
+        pass
+    # End attribute config meta
 
     def __init__(self, configdict, section, consumes=FORMAT.record):
         DbOutput.__init__(self, configdict, section, consumes=[FORMAT.record_array, FORMAT.record])
         self.query = None
         self.db = None
-        self.key = self.cfg.get('key')
 
     def init(self):
         # Connect only once to DB
@@ -79,6 +142,16 @@ def create_query(self, record):
         log.info('query is %s', query)
         return query
 
+    def insert(self, record):
+        if self.replace and self.key and self.key in record:
+            # Try to delete (replace option)
+            del_query = "DELETE FROM %s WHERE %s = '%s'" % (self.cfg.get('table'), self.key, record[self.key])
+            self.db.execute(del_query)
+
+        # Do insert with values from the record dict
+        self.db.execute(self.query, record.values())
+        self.db.commit(close=False)
+
     def write(self, packet):
         # Deal with empty or zero-length data structures (list or dict)
         if packet.data is None or len(packet.data) == 0:
@@ -101,18 +174,16 @@ def write(self, packet):
         # Check if record is single (dict) or array (list of dict)
         if type(record) is dict:
             # Do insert with values from the single record
-            self.db.execute(self.query, record.values())
-            self.db.commit(close=False)
+            self.insert(record)
 
             # log.info('committed record key=%s' % record[self.key])
 
         elif type(record) is list:
-                # Multiple records in list
-                for rec in record:
-                    # Do insert with values from the record
-                    self.db.execute(self.query, rec.values())
-                    self.db.commit(close=False)
+            # Multiple records in list
+            for rec in record:
+                # Do insert with values from the record
+                self.insert(rec)
 
-                log.info('committed %d records' % len(record))
+            log.info('committed %d records' % len(record))
 
         return packet