Permalink
Browse files

add sql_shell command to dataduct CLI, allowing users to connect to m…

…ysql/redshift using creds provided by dataduct configs
  • Loading branch information...
1 parent c9b984b commit ff52414155301b2b0775ab9f9b292b10838f434d @cliu587 cliu587 committed Nov 17, 2015
Showing with 133 additions and 0 deletions.
  1. +30 −0 bin/dataduct
  2. +39 −0 dataduct/data_access/open_shell.py
  3. +28 −0 dataduct/utils/cli.py
  4. +36 −0 docs/commands.rst
View
@@ -17,6 +17,7 @@ CREATE = 'create'
VALIDATE = 'validate'
ACTIVATE = 'activate'
VISUALIZE = 'visualize'
+SQL_SHELL = 'sql_shell'
CONFIG = 'config'
CONFIG_TO_S3 = 'sync_to_s3'
@@ -27,6 +28,9 @@ DROP = 'drop'
GRANT = 'grant'
RECREATE = 'recreate'
+REDSHIFT = 'redshift'
+MYSQL = 'mysql'
+
def initialize_etl_objects(pipeline_definitions, time_delta=None,
frequency_override=None, backfill=False):
@@ -284,6 +288,30 @@ def main():
help='Visualize the database er-diagram',
)
+ # SQL Shell parser
+ sql_shell_parser = subparsers.add_parser(
+ SQL_SHELL,
+ formatter_class=formatter_class,
+ add_help=False,
+ parents=[help_parser]
+ )
+ sql_shell_subparser = sql_shell_parser.add_subparsers(
+ dest='database_type',
+ help='Database Type (MySQL or Redshift)',
+ )
+ sql_shell_subparser.add_parser(
+ REDSHIFT,
+ formatter_class=formatter_class,
+ parents=[mode_parser, help_parser],
+ add_help=False
+ )
+ sql_shell_subparser.add_parser(
+ MYSQL,
+ formatter_class=formatter_class,
+ parents=[mode_parser, help_parser, host_alias_parser],
+ add_help=False
+ )
+
# Check if autocomplete is possible
try:
import argcomplete
@@ -309,6 +337,8 @@ def main():
pipeline_actions(frequency_override=frequency_override, **arg_vars)
elif args.command == DATABASE:
database_actions(**arg_vars)
+ elif args.command == SQL_SHELL:
+ open_sql_shell(**arg_vars)
else:
raise ValueError('Unknown argument provided, use dataduct -h')
@@ -0,0 +1,39 @@
+import os
+from ..utils.hook import hook
+
+import logging
+logger = logging.getLogger(__name__)
+
+@hook('connect_to_redshift')
+def open_psql_shell(redshift_creds, **kwargs):
+ print 'redshift_creds: {}'.format(redshift_creds)
+ command = [
+ "psql",
+ "-h", redshift_creds["HOST"],
+ "-p", str(redshift_creds["PORT"]),
+ "-U", redshift_creds["USERNAME"],
+ "-d", redshift_creds["DATABASE_NAME"],
+ "-vPROMPT1=%[%033[0m%]" + redshift_creds["CLUSTER_ID"] + "%R%[%033[0m%]%# ",
+ "-vPROMPT2=%[%033[0m%]" + redshift_creds["CLUSTER_ID"] + "%R%[%033[0m%]%# ",
+ ]
+ env = dict(os.environ)
+ env['PGPASSWORD'] = redshift_creds["PASSWORD"]
+ logger.info("Running command: {}".format(' '.join(command)))
+ os.execvpe(command[0], command, env=env)
+
+
+@hook('connect_to_mysql')
+def open_mysql_shell(sql_creds, **kwargs):
+ command = [
+ "mysql",
+ "-h", sql_creds["HOST"],
+ "-u", sql_creds["USERNAME"],
+ "--default-character-set=utf8"
+ ]
+ if sql_creds.get("DATABASE"):
+ command.extend(["-D", sql_creds["DATABASE"]])
+
+ env = dict(os.environ)
+ env['MYSQL_PWD'] = sql_creds["PASSWORD"]
+ logger.info("Running command: {}".format(' '.join(command)))
+ os.execvpe(command[0], command, env=env)
View
@@ -71,6 +71,23 @@ def __call__(self, parser, namespace, values, option_string=None):
parser.exit()
+def open_sql_shell(database_type, host_alias=None, **kwargs):
+ """Opens a sql shell for MySQL or Redshift
+ """
+
+ # late import because we need the Singleton config to be loaded in
+ # the dataduct main
+ from dataduct.data_access import open_shell
+ from dataduct.config import Config
+ config = Config()
+ if database_type == 'redshift':
+ open_shell.open_psql_shell()
+ else:
+ assert config.mysql.get(host_alias), \
+ 'host_alias "{}" does not exist in config'.format(host_alias)
+ open_shell.open_mysql_shell(sql_creds=config.mysql[host_alias])
+
+
# Change the width of the output format
formatter_class = lambda prog: RawTextHelpFormatter(prog, max_help_position=50)
@@ -204,3 +221,14 @@ def __call__(self, parser, namespace, values, option_string=None):
dest='s3_path',
help='S3 Path',
)
+
+# database parser
+host_alias_help = 'MySQL Host Alias'
+host_alias_parser = ArgumentParser(
+ description=host_alias_help,
+ add_help=False,
+)
+host_alias_parser.add_argument(
+ dest='host_alias',
+ help='MySQL Host Alias'
+)
View
@@ -146,3 +146,39 @@ Arguments:
- ``-h, --help``: Show help message and exit.
- ``-m MODE, --mode MODE``: Mode or config variables to use. e.g. ``-m production``
- ``filename``: The filename for saving the config file.
+
+SQL Shell Commands
+----------------------
+
+Commands used to connect to either MySQL or Redshift via the terminal.
+
+MySQL
+^^^^^^^^^^
+Connects to a MySQL database using a host alias.
+
+Usage:
+
+::
+
+ dataduct sql_shell mysql [-h] [-m MODE] host_alias
+
+Arguments:
+
+- ``-h, --help``: Show help message and exit.
+- ``-m MODE, --mode MODE``: Mode or config variables to use. e.g. ``-m production``
+- ``host_alias``: The host alias of the database to connect to.
+
+Redshift
+^^^^^^^^^^
+Connects to the Redshift database specified in Dataduct configs.
+
+Usage:
+
+::
+
+ dataduct sql_shell redshift [-h] [-m MODE]
+
+Arguments:
+
+- ``-h, --help``: Show help message and exit.
+- ``-m MODE, --mode MODE``: Mode or config variables to use. e.g. ``-m production``

0 comments on commit ff52414

Please sign in to comment.