Added bcp_dump, bcp_restore, dump_file_mssql

6 files changed, 92 insertions, 101 deletions

diff --git a/sqlplain/Makefile b/sqlplain/Makefile
index 26ae78c..cc97ed5 100644
--- a/sqlplain/Makefile
+++ b/sqlplain/Makefile
@@ -1,5 +1,5 @@
 count:
-	wc -l __init__.py automatize.py configurator.py connection.py \
+	wc -l __init__.py configurator.py connection.py \
 	mssql_support.py postgres_support.py sqlite_support.py \
 	mssql_util.py postgres_util.py sqlite_util.py \
         sql_support.py uri.py util.py table.py runtransac.py memoize.py 
diff --git a/sqlplain/automatize.py b/sqlplain/automatize.py
deleted file mode 100644
index 0aded29..0000000
--- a/sqlplain/automatize.py
+++ /dev/null
@@ -1,79 +0,0 @@
-import os, sys, subprocess, re
-from sqlplain.configurator import configurator
-from sqlplain.util import create_db, create_schema
-from sqlplain.namedtuple import namedtuple
-
-VERSION = re.compile(r'(\d[\d\.-]+)')
-Chunk = namedtuple('Chunk', 'version fname code')
-
-try:
-    CalledProcessError = subprocess.CalledProcessError
-except AttributeError: # Python < 2.5
-    class CalledProcessError(Exception): pass
-    
-def getoutput(commandlist):
-    po = subprocess.Popen(commandlist, stdout=subprocess.PIPE)
-    out, err = po.communicate()
-    if po.returncode or err:
-        raise CalledProcessError('%s [return code %d]' % (err, po.returncode))
-    return out
-    
-def collect(directory, exts):
-    '''
-    Read the files with a given set of extensions from a directory
-    and returns them ordered by version number.
-    '''
-    sql = []
-    for fname in os.listdir(directory):
-        if fname.endswith(exts) and not fname.startswith('_'):
-            version = VERSION.search(fname)
-            if version:
-                code = file(os.path.join(directory, fname)).read()
-                sql.append(Chunk(version, fname, code))
-    return sorted(sql)
-
-def make_db(alias=None, uri=None, dir=None):
-    if alias is not None and uri is None:
-        uri = configurator.uri[alias]
-    if alias is not None and dir is None:
-        dir = configurator.dir[alias]
-    db = create_db(uri, drop=True)
-    chunks = collect(dir, ('.sql', '.py'))
-    for chunk in chunks:
-        if chunk.fname.endswith('.sql'):
-            db.execute(chunk.code)
-        elif chunk.fname.endswith('.py'):
-            exec chunk.code in {}
-
-def populate_db(db, dir=None):
-    uri = db.uri
-    if dir is None:
-        alias = configurator.alias[uri]
-        dir = configurator.dir[alias]
-    else:
-        raise TypeError('Please provide a scriptdir')
-    chunks = collect(dir, ('.sql', '.py'))
-    for chunk in chunks:
-        if chunk.fname.endswith('.sql'):
-            db.execute(chunk.code)
-        elif chunk.fname.endswith('.py'):
-            exec chunk.code in {}
-        
-def make_schema(alias=None, schema=None, uri=None, dir=None):
-    if alias is not None and uri is None:
-        uri = configurator.uri[alias]
-    if alias is not None and dir is None:
-        dir = configurator.dir[alias]
-    db = lazyconnect(uri)
-    schema = create_schema(db, schema, drop=True)
-    chunks = collect(dir, ('.sql', '.py'))
-    for chunk in chunks:
-        if chunk.fname.endswith('.sql'):
-            db.execute(chunk.code)
-        elif chunk.fname.endswith('.py'):
-            exec chunk.code in {}
-        
-if __name__ == '__main__':
-    makedb('utest')
-    makedb('autest')
-    makedb('ftest')
diff --git a/sqlplain/mssql_util.py b/sqlplain/mssql_util.py
index da6b865..7859d13 100644
--- a/sqlplain/mssql_util.py
+++ b/sqlplain/mssql_util.py
@@ -1,8 +1,28 @@
-from sqlplain.util import openclose
+import sys
+from sqlplain.util import openclose, getoutput
+
+BCP = ['freebcp', 'bcp'][sys.platform == 'win32']
 
 def get_kfields_mssql(conn, table):
     return [x.COLUMN_NAME for x in conn.execute('sp_pkeys %s' % table)]
 
+def dump_file_mssql(conn, fname, table_or_query, sep='\t', null='\N'):
+    """
+    Examples:
+    >> dump_file(conn, 'client.csv', 'client')
+    >> dump_file(conn, 'client.csv',
+      'select * from %(database)s..client' % conn.uri) 
+    """
+    uri = conn.uri
+    if table_or_query.lstrip().lower().startswith('select'):
+        out = 'queryout'
+    elif not '.' in table_or_query:
+        out = 'out'
+        table_or_query = '.'.join([uri['database'], '', table_or_query])
+    cmd = [BCP, table_or_query, out, fname, '-S', uri['host'],
+           '-U', uri['user'], '-P', uri['password'], '-c',  '-t', sep]
+    return getoutput(cmd)
+
 def insert_file_mssql(uri, fname, table, sep='\t'):
     return conn.execute(
         'BULK INSERT %s FROM ? SEPARATOR ?' % table, (fname, sep))
@@ -30,3 +50,16 @@ def drop_db_mssql(uri):
 def create_db_mssql(uri):
     openclose(uri.copy(database='master'),
               'CREATE DATABASE %(database)s' % uri)
+
+def bcp_dump(uri, tablename, filename):
+    "Save a table into a bcp binary file"
+    cmd = [BCP, '%s..%s' % (uri['database'], tablename), 'out', filename,
+           '-S', uri['host'], '-U', uri['user'], '-P', uri['password'], '-n']
+    return getoutput(cmd)
+
+def bcp_restore(uri, filename, tablename):
+    "Copy from a bcp binary file into a table"
+    cmd = [BCP, '%s..%s' % (uri['database'], tablename), 'in', filename,
+           '-S', uri['host'], '-U', uri['user'], '-P', uri['password'], '-n']
+    return getoutput(cmd)
+    
diff --git a/sqlplain/postgres_util.py b/sqlplain/postgres_util.py
index 1ae0536..a97e5a0 100644
--- a/sqlplain/postgres_util.py
+++ b/sqlplain/postgres_util.py
@@ -1,5 +1,4 @@
-from sqlplain.util import openclose
-from sqlplain.automatize import getoutput
+from sqlplain.util import openclose, getoutput
 
 GET_PKEYS = '''\
 SELECT attname FROM pg_attribute
@@ -34,7 +33,7 @@ def insert_file_postgres(conn, fname, table, sep=',', null='\N'):
     return conn.execute(templ % table, (fname, sep, null))
 
 def dump_file_postgres(conn, fname, query, sep='\t', null='\N'):
-    sql = "COPY (%s) TO ? WITH DELIMITER ? NULL ?" % query
+    sql = "COPY %s TO ? WITH DELIMITER ? NULL ?" % query
     return conn.execute(sql, (fname, sep, null))
 
 def exists_db_postgres(uri):
@@ -45,10 +44,12 @@ def exists_db_postgres(uri):
             return True
     return False
 
-def get_schema_postgres(uri, objectname):
-    cmd = ['pg_dump', '-s',
-           '-t', objectname,
-           '-h', uri['host'],
-           '-U', uri['user'],
-           '-d', uri['database']]
+def pg_dump(uri, *args):
+    cmd = ['pg_dump', '-h', uri['host'], '-U', uri['user'],
+           '-d', uri['database']] + list(args)
+    return getoutput(cmd)
+
+def pg_restore(uri, *args):
+    cmd = ['pg_restore', '-h', uri['host'], '-U', uri['user'],
+           '-d', uri['database']] + list(args)
     return getoutput(cmd)
diff --git a/sqlplain/tests/test_bcp.py b/sqlplain/tests/test_bcp.py
new file mode 100644
index 0000000..ba7d149
--- /dev/null
+++ b/sqlplain/tests/test_bcp.py
@@ -0,0 +1,8 @@
+from sqlplain import lazyconnect
+from sqlplain.mssql_util import *
+
+srs = lazyconnect('dbserver2')
+
+bcp_dump(srs.uri, 'client', '/tmp/client.bcpdump')
+
+bcp_restore(srs.uri, '/tmp/client.bcpdump', 'client')
diff --git a/sqlplain/util.py b/sqlplain/util.py
index 9d337a4..3f5886b 100644
--- a/sqlplain/util.py
+++ b/sqlplain/util.py
@@ -2,12 +2,36 @@
 Notice: create_db and drop_db are not transactional.
 """
 
-import os, sys, re
+import os, sys, re, subprocess
 from sqlplain.uri import URI
 from sqlplain import lazyconnect, do
 from sqlplain.connection import Transaction
 from sqlplain.namedtuple import namedtuple
-
+try:
+    CalledProcessError = subprocess.CalledProcessError
+except AttributeError: # Python < 2.5
+    class CalledProcessError(Exception):
+        def __init__(self, returncode, cmd):
+            self.returncode = returncode
+            self.cmd =cmd
+    
+def getoutput(commandlist):
+    'Returns the output of a system command or raise a CalledProcessError'
+    po = subprocess.Popen(commandlist, stdout=subprocess.PIPE)
+    out, err = po.communicate()
+    if po.returncode or err:
+        if err:
+            sys.stderr.write(err)
+            sys.stderr.flush()
+        cmd_str = ''
+        for cmd in commandlist:
+            if re.search(r'\s', cmd):
+                cmd_str += '"%s" ' % cmd
+            else:
+                cmd_str += cmd + " "
+        raise CalledProcessError(po.returncode, cmd_str)
+    return out
+    
 VERSION = re.compile(r'(\d[\d\.-]+)')
 Chunk = namedtuple('Chunk', 'version fname code')
 
@@ -59,6 +83,13 @@ def _collect(directory, exts):
                 sql.append(Chunk(version, fname, code))
     return sorted(sql)
 
+def runscripts(db, scriptdir, exts):
+    for chunk in _collect(scriptdir, exts):
+        if chunk.fname.endswith('.sql'):
+            db.executescript(chunk.code)
+        elif chunk.fname.endswith('.py'):
+            exec chunk.code in {}
+
 def create_db(uri, force=False, scriptdir=None, **kw):
     """
     Create the database specified by uri. If the database exists already
@@ -77,12 +108,7 @@ def create_db(uri, force=False, scriptdir=None, **kw):
     db = lazyconnect(uri, **kw)
     scriptdir = uri.scriptdir or scriptdir
     if scriptdir:
-        chunks = _collect(scriptdir, ('.sql', '.py'))
-        for chunk in chunks:
-            if chunk.fname.endswith('.sql'):
-                db.executescript(chunk.code)
-            elif chunk.fname.endswith('.py'):
-                exec chunk.code in {}
+        runscripts(db, scriptdir, ('.sql', '.py'))
     return db
 
 def create_table(conn, tname, body, force=False):
@@ -142,11 +168,11 @@ def insert_rows(conn, tname, rows):
         n += conn.execute(templ, row)
     return 
     
-def insert_file(conn, fname, tname, sep=','):
+def insert_file(conn, fname, tname, sep):
     "Bulk insert a CSV file into a table"""
     return _call('insert_file', conn, fname, tname, sep)
 
-def dump_file(conn, fname, query, sep=',', null='\N'):
+def dump_file(conn, fname, query, sep, null='\N'):
     "Save a query on a CSV file"
     return _call('dump_file', conn, fname, query, sep, null)
     
@@ -194,7 +220,7 @@ exists_schema = do("SELECT nspname FROM pg_namespace WHERE nspname=?")
 def drop_schema(db, schema):
     db.execute('DROP SCHEMA %s CASCADE' % schema)
 
-def create_schema(db, schema, force=False):
+def create_schema(db, schema, force=False, schema_dir=None):
     """
     Create the specified schema. If the schema exists already
     an error is raised, unless force is True: in that case the schema
@@ -204,3 +230,5 @@ def create_schema(db, schema, force=False):
         drop_schema(db, schema)
     db.execute('CREATE SCHEMA %s' % schema)
     set_schema(db, schema)
+    if schema_dir:
+        runscripts(db, schema_dir, ('.sql', '.py'))