From dcda3db140789be63e9231056669fe529e3c5566 Mon Sep 17 00:00:00 2001 From: Jeremy Carbaugh Date: Tue, 1 Sep 2009 17:45:05 -0400 Subject: [PATCH] add replace option to sqliteemitter and add better kwargs to dictreader on csvsource --- saucebrush/emitters.py | 15 ++++++++++++--- saucebrush/sources.py | 4 ++-- 2 files changed, 14 insertions(+), 5 deletions(-) diff --git a/saucebrush/emitters.py b/saucebrush/emitters.py index 9c516f1..636cd08 100644 --- a/saucebrush/emitters.py +++ b/saucebrush/emitters.py @@ -85,24 +85,33 @@ class SqliteEmitter(Emitter): as a third parameter to SqliteEmitter.) """ - def __init__(self, dbname, table_name, fieldnames=None): + def __init__(self, dbname, table_name, fieldnames=None, replace=False, quiet=False): super(SqliteEmitter, self).__init__() import sqlite3 self._conn = sqlite3.connect(dbname) self._cursor = self._conn.cursor() self._table_name = table_name + self._replace = replace + self._quiet = quiet if fieldnames: create = "CREATE TABLE IF NOT EXISTS %s (%s)" % (table_name, ', '.join([' '.join((field, 'TEXT')) for field in fieldnames])) self._cursor.execute(create) def emit_record(self, record): + import sqlite3 # input should be escaped with ? if data isn't trusted qmarks = ','.join(('?',) * len(record)) - insert = 'INSERT INTO %s (%s) VALUES (%s)' % (self._table_name, + insert = 'INSERT OR REPLACE' if self._replace else 'INSERT' + insert = '%s INTO %s (%s) VALUES (%s)' % (insert, self._table_name, ','.join(record.keys()), qmarks) - self._cursor.execute(insert, record.values()) + try: + self._cursor.execute(insert, record.values()) + except sqlite3.IntegrityError, ie: + if not self._quiet: + raise ie + self.reject_record(record, ie.message) def done(self): self._conn.commit() diff --git a/saucebrush/sources.py b/saucebrush/sources.py index c00a3e1..cc1330d 100644 --- a/saucebrush/sources.py +++ b/saucebrush/sources.py @@ -22,9 +22,9 @@ class CSVSource(object): ignoring the first row (presumed to be column names). """ - def __init__(self, csvfile, fieldnames=None, skiprows=0, delimiter=','): + def __init__(self, csvfile, fieldnames=None, skiprows=0, **kwargs): import csv - self._dictreader = csv.DictReader(csvfile, fieldnames, delimiter=delimiter) + self._dictreader = csv.DictReader(csvfile, fieldnames, **kwargs) for _ in xrange(skiprows): self._dictreader.next()