From ca014d3fed7fa5bae15bd14bd7b4331cd3f6ae3b Mon Sep 17 00:00:00 2001
From: Dominik Moritz <domoritz@gmail.com>
Date: Wed, 26 Sep 2012 13:26:11 +0100
Subject: [PATCH] pipe sql into psql instead of passing it as parameter,
 refactor parsing of database urls

---
 ckan/lib/cli.py               | 38 +++++++++++++++++++++++------
 ckanext/datastore/commands.py | 46 +++++++++++++++++------------------
 2 files changed, 52 insertions(+), 32 deletions(-)

diff --git a/ckan/lib/cli.py b/ckan/lib/cli.py
index c20296b3517..8a77de79ae2 100644
--- a/ckan/lib/cli.py
+++ b/ckan/lib/cli.py
@@ -13,6 +13,35 @@
 #   i.e. do the imports in methods, after _load_config is called.
 #   Otherwise loggers get disabled.
 
+
+def parse_db_config(config_key='sqlalchemy.url'):
+    ''' Takes a config key for a database connection url and parses it into
+    a dictionary. Expects a url like:
+
+    'postgres://tester:pass@localhost/ckantest3'
+    '''
+    from pylons import config
+    url = config[config_key]
+    regex = [
+        '^\s*(?P<db_type>\w*)',
+        '://',
+        '(?P<db_user>[^:]*)',
+        ':?',
+        '(?P<db_pass>[^@]*)',
+        '@',
+        '(?P<db_host>[^/:]*)',
+        ':?',
+        '(?P<db_port>[^/]*)',
+        '/',
+        '(?P<db_name>[\w.-]*)'
+    ]
+    db_details_match = re.match(''.join(regex), url)
+    if not db_details_match:
+        raise Exception('Could not extract db details from url: %r' % url)
+    db_details = db_details_match.groupdict()
+    return db_details
+
+
 class MockTranslator(object):
     def gettext(self, value):
         return value
@@ -134,14 +163,7 @@ def command(self):
             sys.exit(1)
 
     def _get_db_config(self):
-        from pylons import config
-        url = config['sqlalchemy.url']
-        # e.g. 'postgres://tester:pass@localhost/ckantest3'
-        db_details_match = re.match('^\s*(?P<db_type>\w*)://(?P<db_user>[^:]*):?(?P<db_pass>[^@]*)@(?P<db_host>[^/:]*):?(?P<db_port>[^/]*)/(?P<db_name>[\w.-]*)', url)
-        if not db_details_match:
-            raise Exception('Could not extract db details from url: %r' % url)
-        db_details = db_details_match.groupdict()
-        return db_details
+        return parse_db_config()
 
     def _get_postgres_cmd(self, command):
         self.db_details = self._get_db_config()
diff --git a/ckanext/datastore/commands.py b/ckanext/datastore/commands.py
index 27da26d7e0f..3e2311f19d4 100644
--- a/ckanext/datastore/commands.py
+++ b/ckanext/datastore/commands.py
@@ -1,9 +1,12 @@
 import re
-from ckan.lib.cli import CkanCommand
-
+import sys
 import logging
+
+import ckan.lib.cli as cli
+
 log = logging.getLogger(__name__)
 
+
 read_only_user_sql = '''
 -- revoke permissions for the new user
 REVOKE CREATE ON SCHEMA public FROM PUBLIC;
@@ -34,7 +37,7 @@
 '''
 
 
-class SetupDatastoreCommand(CkanCommand):
+class SetupDatastoreCommand(cli.CkanCommand):
     '''Perform commands to set up the datastore.
     Make sure that the datastore urls are set properly before you run these commands.
 
@@ -68,9 +71,9 @@ def command(self):
         cmd = self.args[0]
         self._load_config()
 
-        self.db_write_url_parts = self._get_db_config('ckan.datastore.write_url')
-        self.db_read_url_parts = self._get_db_config('ckan.datastore.read_url')
-        self.db_ckan_url_parts = self._get_db_config('sqlalchemy.url')
+        self.db_write_url_parts = cli.parse_db_config('ckan.datastore.write_url')
+        self.db_read_url_parts = cli.parse_db_config('ckan.datastore.read_url')
+        self.db_ckan_url_parts = cli.parse_db_config('sqlalchemy.url')
 
         assert self.db_write_url_parts['db_name'] == self.db_read_url_parts['db_name'], "write and read db should be the same"
 
@@ -95,31 +98,26 @@ def command(self):
             log.error('Command "%s" not recognized' % (cmd,))
             return
 
-    def _get_db_config(self, name):
-        from pylons import config
-        url = config[name]
-        # e.g. 'postgres://tester:pass@localhost/ckantest3'
-        db_details_match = re.match('^\s*(?P<db_type>\w*)://(?P<db_user>[^:]*):?(?P<db_pass>[^@]*)@(?P<db_host>[^/:]*):?(?P<db_port>[^/]*)/(?P<db_name>[\w.-]*)', url)
-        if not db_details_match:
-            raise Exception('Could not extract db details from url: %r' % url)
-        db_details = db_details_match.groupdict()
-        return db_details
-
-    def _run_cmd(self, command_line):
+    def _run_cmd(self, command_line, inputstring=''):
         import subprocess
-        retcode = subprocess.call(command_line, shell=True)
-        if retcode != 0:
-            raise SystemError('Command exited with errorcode: %i' % retcode)
+        p = subprocess.Popen(
+            command_line, shell=True,
+            stdin=subprocess.PIPE,
+            stdout=subprocess.PIPE,
+            stderr=subprocess.PIPE)
+        stdout_value, stderr_value = p.communicate(input=inputstring)
+        if stderr_value:
+            print '\nAn error occured: {0}'.format(stderr_value)
+            sys.exit(1)
 
     def _run_sql(self, sql, as_sql_user, database='postgres'):
         if self.verbose:
             print "Executing: \n#####\n", sql, "\n####\nOn database:", database
         if not self.simulate:
-            self._run_cmd("psql --username='{username}' --dbname='{database}' -c \"{sql}\" -W".format(
+            self._run_cmd("psql --username='{username}' --dbname='{database}' -W".format(
                 username=as_sql_user,
-                database=database,
-                sql=sql.replace('"', '\\"')
-            ))
+                database=database
+            ), inputstring=sql)
 
     def create_db(self):
         sql = "create database {0}".format(self.db_write_url_parts['db_name'])