* bulk_yearly_batch
* bulk_yearly_field
* bulk_yearly_time
+* bulk_direct - functionally identical to bulk
== HANDLER ARGUMENTS ==
sql statement(s) to execute after creating partition table. Usable
variables are the same as in part_template
+encoding:
+ name of destination encoding. handler replaces all invalid encoding symbols
+ and logs them as warnings
+
NB! londiste3 does not currently support table renaming and field mapping when
creating or coping initial data to destination table. --expect-sync and
--skip-truncate should be used and --create switch is to be avoided.
import sys
import datetime
import new
+import codecs
import skytools
from londiste.handler import BaseHandler
from skytools import quote_ident, quote_fqident, UsageError
', '.join('%%(%s)s' % arg for arg in PART_FUNC_ARGS))
-#----------------------------------------
+
+#------------------------------------------------------------------------------
# LOADERS
-#----------------------------------------
+#------------------------------------------------------------------------------
class BaseLoader:
else:
# fscking problems with long-lived temp tables
self.drop(curs)
-
+
def create_temp(self, curs):
""" check if temp table exists. Returns False if using existing temp
table and True if creating new
return True
def bulk_insert(self, curs, data, table = None):
- """Copy data to table. If table not provided, use temp table.
+ """Copy data to table. If table not provided, use temp table.
When re-using existing temp table, it is always truncated first and
analyzed after copy.
"""
_use_temp = table is None
# if table not specified use temp
if _use_temp:
- table = self.temp
+ table = self.temp
# truncate when re-using existing table
if not self.create_temp(curs):
- self.truncate(curs)
+ self.truncate(curs)
self.log.debug("bulk: COPY %d rows into %s" % (len(data), table))
skytools.magic_insert(curs, table, data, self.fields,
quoted_table = True)
LOADERS = {'direct': DirectLoader, 'bulk': BulkLoader}
-#----------------------------------------
+
+#------------------------------------------------------------------------------
# ROW HANDLERS
-#----------------------------------------
+#------------------------------------------------------------------------------
class RowHandler:
'keep_latest': KeepLatestRowHandler}
-#----------------------------------------
-# DISPATCHER
-#----------------------------------------
-class AttrDict(dict):
- """Dict with values accessible with attributes"""
- def __getattr__(self, name):
- return self[name]
- def __setattr__(self, name, value):
- self[name] = value
+#------------------------------------------------------------------------------
+# ENCODING VALIDATOR
+#------------------------------------------------------------------------------
+
+
+class EncodingValidator:
+ def __init__(self, log, encoding = 'utf-8', replacement = u'\ufffd'):
+ self.log = log
+ self.encoding = encoding
+ self.replacement = replacement
+ self.columns = None
+ self.error_count = 0
+ codecs.register_error("error_handler", self._error_handler)
+
+ def _error_handler(self, exc):
+ # process only UnicodeDecodeError
+ if not isinstance(exc, UnicodeDecodeError):
+ raise exc
+ # find starting position of line with error and log warning
+ _line_start = exc.object.rfind('\n', 0, exc.start) + 1
+ _col = self.columns[exc.object.count('\t', _line_start, exc.start)]
+ _msg = "replacing invalid %s sequence %r in column %s"%\
+ (self.encoding, exc.object[exc.start:exc.end], _col)
+ self.log.warning(_msg)
+ # increase error count
+ self.error_count += 1
+ # return replacement char and position to continue from
+ # NB! doesn't replace multiple symbols, so it's harder to break file
+ # structure like replace \t or \n
+ return self.replacement, exc.start + 1
+
+ def validate(self, data, columns):
+ self.columns = columns
+ self.error_count = 0
+ _unicode = data.decode(self.encoding, "error_handler")
+ # when no erros then return input data as is, else re-encode fixed data
+ if self.error_count == 0:
+ return data
+ else:
+ return _unicode.encode(self.encoding)
+
+ def validate_dict(self, data):
+ _cols, _vals = zip(*data.items())
+ _fixed = self.validate('\t'.join(_vals), _cols)
+ if self.error_count == 0:
+ return data
+ else:
+ return dict(zip(_cols, _fixed.split('\t')))
+
+
+
+#------------------------------------------------------------------------------
+# DISPATCHER
+#------------------------------------------------------------------------------
class Dispatcher(BaseHandler):
self.conf = self.get_config()
hdlr_cls = ROW_HANDLERS[self.conf.row_mode]
self.row_handler = hdlr_cls(self.log)
+ if self.conf.encoding:
+ self.encoding_validator = EncodingValidator(self.log,
+ self.conf.encoding)
+ else:
+ self.encoding_validator = None
def get_config(self):
"""Processes args dict"""
- conf = AttrDict()
+ conf = skytools.dbdict()
# set table mode
conf.table_mode = self.get_arg('table_mode', TABLE_MODES)
if conf.table_mode == 'part':
conf.field_map[tmp[0]] = tmp[0]
else:
conf.field_map[tmp[0]] = tmp[1]
+ # encoding validator
+ conf.encoding = self.args.get('encoding')
return conf
def get_arg(self, name, value_list, default = None):
if dst not in self.row_handler.table_map:
self.row_handler.add_table(dst, LOADERS[self.conf.load_mode],
self.pkeys, self.conf)
+ if self.encoding_validator:
+ data = self.encoding_validator.validate_dict(data)
self.row_handler.process(dst, op, data)
#BaseHandler.process_event(self, ev, sql_queue_func, arg)
exec_with_vals(self.conf.post_part)
self.log.info("Created table: %s" % dst)
+ def real_copy(self, tablename, src_curs, dst_curs, column_list, cond_list):
+ """do actual table copy and return tuple with number of bytes and rows
+ copyed
+ """
+ _src_cols = _dst_cols = column_list
+ _write_hook = None
+ condition = ' and '.join(cond_list)
+
+ if self.conf.skip_fields:
+ _src_cols = [col for col in column_list
+ if col not in self.conf.skip_fields]
+ _dst_cols = _src_cols
+ if self.conf.field_map:
+ _src_cols = [col for col in _src_cols if col in self.conf.field_map]
+ _dst_cols = [self.conf.field_map[col] for col in _src_cols]
+
+ if self.encoding_validator:
+ def _write_hook(obj, data):
+ return self.encoding_validator.validate(data, _src_cols)
+
+ return skytools.full_copy(tablename, src_curs, dst_curs, _src_cols,
+ condition, self.table_name, _dst_cols,
+ write_hook = _write_hook)
+
+
+
+#------------------------------------------------------------------------------
# register handler class
+#------------------------------------------------------------------------------
+
+
__londiste_handlers__ = [Dispatcher]
+
+
+#------------------------------------------------------------------------------
# helper function for creating dispachers with different default values
+#------------------------------------------------------------------------------
+
+
def handler(name):
def wrapper(func):
def _init_override(self, table_name, args, log):
return func
return wrapper
-def dupd(*p):
+
+def update(*p):
""" Update dicts given in params with its precessor param dict
in reverse order """
return reduce(lambda x, y: x.update(y) or x,
(p[i] for i in range(len(p)-1,-1,-1)), {})
+
+
+#------------------------------------------------------------------------------
+# build set of handlers with different default values for easier use
+#------------------------------------------------------------------------------
+
+
LOAD = { '': { 'load_mode': 'direct' },
'bulk': { 'load_mode': 'bulk' }
}
'row_mode': 'keep_latest',
}
-# build set of handlers with different default values for easier use
for load, load_dict in LOAD.items():
for period, period_dict in PERIOD.items():
for mode, mode_dict in MODE.items():
# define creator func to keep default dicts in separate context
def create_handler():
handler_name = '_'.join(p for p in (load, period, mode) if p)
- default = dupd(mode_dict, period_dict, load_dict, BASE)
+ default = update(mode_dict, period_dict, load_dict, BASE)
@handler(handler_name)
def handler_func(args):
- return dupd(args, default)
+ return update(args, default)
create_handler()
-# TODO: bulk & ignore handlers
+
+@handler('bulk_direct')
+def bulk_direct_handler(args):
+ return update(args, {'load_mode': 'bulk', 'table_mode': 'direct'})
class CopyPipe(object):
"Splits one big COPY to chunks."
- def __init__(self, dstcurs, tablename = None, limit = 512*1024, cancel_func=None, sql_from = None):
+ def __init__(self, dstcurs, tablename = None, limit = 512*1024,
+ sql_from = None):
self.tablename = tablename
self.sql_from = sql_from
self.dstcurs = dstcurs
self.buf = StringIO()
self.limit = limit
- self.cancel_func = None
+ #hook for new data, hook func should return new data
+ #def write_hook(obj, data):
+ # return data
+ self.write_hook = None
+ #hook for flush, hook func result is discarded
+ # def flush_hook(obj):
+ # return None
+ self.flush_hook = None
self.total_rows = 0
self.total_bytes = 0
def write(self, data):
"New data from psycopg"
+ if self.write_hook:
+ data = self.write_hook(self, data)
self.total_bytes += len(data)
self.total_rows += data.count("\n")
def flush(self):
"Send data out."
- if self.cancel_func:
- self.cancel_func()
+ if self.flush_hook:
+ self.flush_hook(self)
if self.buf.tell() <= 0:
return
self.buf.seek(0)
self.buf.truncate()
+
def full_copy(tablename, src_curs, dst_curs, column_list = [], condition = None,
- dst_tablename = None, dst_column_list = None):
+ dst_tablename = None, dst_column_list = None,
+ write_hook = None, flush_hook = None):
"""COPY table from one db to another."""
# default dst table and dst columns to source ones
sql_to = "COPY %s TO stdout" % src
sql_from = "COPY %s FROM stdin" % dst
buf = CopyPipe(dst_curs, sql_from = sql_from)
+ buf.write_hook = write_hook
+ buf.flush_hook = flush_hook
src_curs.copy_expert(sql_to, buf)
else:
if condition:
# regular psycopg copy_to generates invalid sql for subselect copy
raise Exception('copy_expert() is needed for conditional copy')
buf = CopyPipe(dst_curs, dst)
+ buf.write_hook = write_hook
+ buf.flush_hook = flush_hook
src_curs.copy_to(buf, src)
buf.flush()
col = skytools.quote_ident(new_k)
val = skytools.quote_literal(row[k])
whe_list.append("%s = %s" % (col, val))
- whe_str = " and ".join(whe_list)
+ whe_str = " and ".join(whe_list)
return "delete from only %s where %s;" % (skytools.quote_fqident(tbl), whe_str)
if __name__ == '__main__':
projects / skytools.git / commitdiff