Add special commands and tests (dbcli#155)

abhisheksinha89 · web-flow · commit 48b23a08ac65 · 2018-02-02T14:16:31.000-08:00
* Special commands create source files

* Special commands add help special command

* Add list databases command

* Add list schemas special command

* Add list tables

* Add list views and list indexes

* Add special command tests

* Update tests to use dynamically generated resource names
diff --git a/build.py b/build.py
@@ -97,7 +97,7 @@ def unit_test():
     utility.exec_command(
         'pytest --cov mssqlcli tests/test_mssqlcliclient.py tests/test_main.py tests/test_fuzzy_completion.py '
         'tests/test_rowlimit.py tests/test_sqlcompletion.py tests/test_prioritization.py mssqlcli/jsonrpc/contracts/tests '
-        'tests/test_telemetry.py',
+        'tests/test_telemetry.py tests/test_special.py',
         utility.ROOT_DIR,
         continue_on_error=False)
 
diff --git a/mssqlcli/mssqlcliclient.py b/mssqlcli/mssqlcliclient.py
@@ -11,6 +11,7 @@
 from mssqlcli.jsonrpc.contracts import connectionservice, queryexecutestringservice as queryservice
 from mssqlcli.packages.parseutils.meta import ForeignKey
 from mssqlcli.sqltoolsclient import SqlToolsClient
+from packages import special
 
 logger = logging.getLogger(u'mssqlcli.mssqlcliclient')
 time_wait_if_no_response = 0.05
@@ -121,19 +122,25 @@ def connect(self):
             return self.owner_uri
 
     def execute_multi_statement_single_batch(self, query):
-        # Remove spaces, EOL and semi-colons from end
-        query = query.strip()
-        if not query:
-            yield None, None, None, query, False
-        else:
-            for sql in sqlparse.split(query):
-                # Remove spaces, EOL and semi-colons from end
-                sql = sql.strip().rstrip(';')
-                if not sql:
-                    yield None, None, None, sql, False
-                    continue
-                for rows, columns, status, statement, is_error in self.execute_single_batch_query(sql):
-                    yield rows, columns, status, statement, is_error
+        # Try to run first as special command
+        try:
+            for rows, columns, status, statement, is_error in special.execute(self, query):
+                yield rows, columns, status, statement, is_error
+        except special.CommandNotFound:
+            # Execute as normal sql
+            # Remove spaces, EOL and semi-colons from end
+            query = query.strip()
+            if not query:
+                yield None, None, None, query, False
+            else:
+                for sql in sqlparse.split(query):
+                    # Remove spaces, EOL and semi-colons from end
+                    sql = sql.strip().rstrip(';')
+                    if not sql:
+                        yield None, None, None, sql, False
+                        continue
+                    for rows, columns, status, statement, is_error in self.execute_single_batch_query(sql):
+                        yield rows, columns, status, statement, is_error
 
     def execute_single_batch_query(self, query):
         if not self.is_connected:
diff --git a/mssqlcli/packages/special/__init__.py b/mssqlcli/packages/special/__init__.py
@@ -0,0 +1,11 @@
+__all__ = []
+
+
+def export(defn):
+    """Decorator to explicitly mark functions that are exposed in a lib."""
+    globals()[defn.__name__] = defn
+    __all__.append(defn.__name__)
+    return defn
+
+from . import main
+from . import commands
diff --git a/mssqlcli/packages/special/commands.py b/mssqlcli/packages/special/commands.py
@@ -0,0 +1,95 @@
+import logging
+from .main import special_command, RAW_QUERY, PARSED_QUERY, NO_QUERY
+
+logger = logging.getLogger('mssqlcli.commands')
+
+
+@special_command('\\l', '\\l[+] [pattern]', 'List databases.', aliases=('\\list',))
+def list_databases(mssqlcliclient, pattern, verbose):
+    base_query = u'select {0} from sys.databases'
+    if verbose:
+        base_query = base_query.format('name, create_date, compatibility_level, collation_name')
+    else:
+        base_query = base_query.format('name')
+    if pattern:
+        base_query += " where name like '%{0}%'".format(pattern)
+
+    return mssqlcliclient.execute_multi_statement_single_batch(base_query)
+
+
+@special_command('\\dn', '\\dn[+] [pattern]', 'List schemas.')
+def list_schemas(mssqlcliclient, pattern, verbose):
+    base_query = u'select {0} from sys.schemas'
+    if verbose:
+        base_query = base_query.format('name, schema_id, principal_id')
+    else:
+        base_query = base_query.format('name')
+    if pattern:
+        base_query += " where name like '%{0}%'".format(pattern)
+
+    return mssqlcliclient.execute_multi_statement_single_batch(base_query)
+
+
+@special_command('\\dt', '\\dt[+] [pattern]', 'List tables.')
+def list_tables(mssqlcliclient, pattern, verbose):
+    base_query = u'select {0} from INFORMATION_SCHEMA.TABLES WHERE TABLE_TYPE=\'BASE TABLE\''
+    if verbose:
+        base_query = base_query.format('*')
+    else:
+        base_query = base_query.format('table_schema, table_name')
+    if pattern:
+        base_query += "and table_name like '%{0}%'".format(pattern)
+
+    return mssqlcliclient.execute_multi_statement_single_batch(base_query)
+
+
+@special_command('\\dv', '\\dv[+] [pattern]', 'List views.')
+def list_tables(mssqlcliclient, pattern, verbose):
+    base_query = u'select {0} from INFORMATION_SCHEMA.TABLES WHERE TABLE_TYPE=\'VIEW\''
+    if verbose:
+        base_query = base_query.format('table_catalog as catalog, table_schema as schema_name, '
+                                       'table_name as view_name')
+    else:
+        base_query = base_query.format('table_schema as schema_name, table_name as view_name')
+    if pattern:
+        base_query += "and table_name like '%{0}%'".format(pattern)
+
+    return mssqlcliclient.execute_multi_statement_single_batch(base_query)
+
+
+@special_command('\\di', '\\di[+] [pattern]', 'List indexes.')
+def list_indexes(mssqlcliclient, pattern, verbose):
+    base_query = '''
+SELECT
+     TableName = t.name,
+     IndexName = ind.name,
+     ColumnName = col.name {verbose}
+FROM
+     sys.indexes ind
+INNER JOIN
+     sys.index_columns ic ON  ind.object_id = ic.object_id and ind.index_id = ic.index_id
+INNER JOIN
+     sys.columns col ON ic.object_id = col.object_id and ic.column_id = col.column_id
+INNER JOIN
+     sys.tables t ON ind.object_id = t.object_id
+WHERE
+     ind.is_primary_key = 0
+     AND ind.is_unique = 0
+     AND ind.is_unique_constraint = 0
+     AND t.is_ms_shipped = 0
+     AND ind.name like '%{pattern}%'
+ORDER BY
+     t.name, ind.name, ind.index_id, ic.index_column_id;
+    '''
+
+    if verbose:
+        base_query = base_query.format(verbose=',IndexId = ind.index_id, ColumnId = ic.index_column_id',
+                                       pattern='{pattern}')
+    else:
+        base_query = base_query.format(verbose='', pattern='{pattern}')
+    if pattern:
+        base_query = base_query.format(pattern=pattern)
+    else:
+        base_query = base_query.format(pattern='')
+
+    return mssqlcliclient.execute_multi_statement_single_batch(base_query)
diff --git a/mssqlcli/packages/special/main.py b/mssqlcli/packages/special/main.py
@@ -0,0 +1,90 @@
+import logging
+from collections import namedtuple
+
+from . import export
+
+logger = logging.getLogger('mssqlcli.special')
+
+NO_QUERY = 0
+PARSED_QUERY = 1
+RAW_QUERY = 2
+
+SpecialCommand = namedtuple('SpecialCommand',
+                            ['handler', 'command', 'shortcut', 'description', 'arg_type', 'hidden',
+                                'case_sensitive'])
+
+COMMANDS = {}
+
+
+@export
+class CommandNotFound(Exception):
+    pass
+
+
+@export
+def parse_special_command(sql):
+    command, _, arg = sql.partition(' ')
+    verbose = '+' in command
+    command = command.strip().replace('+', '')
+    return (command, verbose, arg.strip())
+
+
+@export
+def special_command(command, shortcut, description, arg_type=PARSED_QUERY,
+                    hidden=False, case_sensitive=False, aliases=()):
+    def wrapper(wrapped):
+        register_special_command(wrapped, command, shortcut, description,
+                                 arg_type, hidden, case_sensitive, aliases)
+        return wrapped
+    return wrapper
+
+
+@export
+def register_special_command(handler, command, shortcut, description,
+                             arg_type=PARSED_QUERY, hidden=False, case_sensitive=False, aliases=()):
+    cmd = command.lower() if not case_sensitive else command
+    COMMANDS[cmd] = SpecialCommand(handler, command, shortcut, description,
+                                   arg_type, hidden, case_sensitive)
+    for alias in aliases:
+        cmd = alias.lower() if not case_sensitive else alias
+        COMMANDS[cmd] = SpecialCommand(handler, command, shortcut, description,
+                                       arg_type, case_sensitive=case_sensitive,
+                                       hidden=True)
+
+
+@export
+def execute(mssqlcliclient, sql):
+    """Execute a special command and return the results. If the special command
+    is not supported a KeyError will be raised.
+    """
+    command, verbose, pattern = parse_special_command(sql)
+
+    if (command not in COMMANDS) and (command.lower() not in COMMANDS):
+        raise CommandNotFound
+
+    try:
+        special_cmd = COMMANDS[command]
+    except KeyError:
+        special_cmd = COMMANDS[command.lower()]
+        if special_cmd.case_sensitive:
+            raise CommandNotFound('Command not found: %s' % command)
+
+    logger.debug(u'Executing special command {0} with argument {1}.'.format(command, pattern))
+
+    if special_cmd.arg_type == NO_QUERY:
+        return special_cmd.handler()
+    elif special_cmd.arg_type == PARSED_QUERY:
+        return special_cmd.handler(mssqlcliclient, pattern=pattern, verbose=verbose)
+    elif special_cmd.arg_type == RAW_QUERY:
+        return special_cmd.handler(mssqlcliclient, query=sql)
+
+
+@special_command('help', '\\?', 'Show this help.', arg_type=NO_QUERY, aliases=('\\?', '?'))
+def show_help():  # All the parameters are ignored.
+    headers = ['Command', 'Shortcut', 'Description']
+    result = []
+
+    for _, value in sorted(COMMANDS.items()):
+        if not value.hidden:
+            result.append((value.command, value.shortcut, value.description))
+    return [(result, headers, None, None, False)]
diff --git a/tests/test_special.py b/tests/test_special.py
@@ -0,0 +1,95 @@
+import unittest
+import uuid
+from mssqlutils import create_mssql_cli_client, shutdown
+from mssqlcli.packages.special.main import special_command, execute, NO_QUERY
+
+
+# All tests require a live connection to a SQL Server database
+class SpecialCommandsTests(unittest.TestCase):
+    session_guid = str(uuid.uuid4().hex)
+    table1 = 'mssql_cli_table1_{0}'.format(session_guid)
+    table2 = 'mssql_cli_table2_{0}'.format(session_guid)
+    view = 'mssql_cli_view_{0}'.format(session_guid)
+    database = 'mssql_cli_db_{0}'.format(session_guid)
+    schema = 'mssql_cli_schema_{0}'.format(session_guid)
+    index = 'mssql_cli_index_{0}'.format(session_guid)
+
+    @classmethod
+    def setUpClass(cls):
+        try:
+            # create the database objects to test upon
+            client = create_mssql_cli_client()
+            list(client.execute_single_batch_query('CREATE DATABASE {0};'.format(cls.database)))
+            list(client.execute_single_batch_query('CREATE TABLE {0} (a int, b varchar(25));'.format(cls.table1)))
+            list(client.execute_single_batch_query('CREATE TABLE {0} (x int, y varchar(25), z bit);'.format(cls.table2)))
+            list(client.execute_single_batch_query('CREATE VIEW {0} as SELECT a from {1};'.format(cls.view, cls.table1)))
+            list(client.execute_single_batch_query('CREATE SCHEMA {0};'.format(cls.schema)))
+            list(client.execute_single_batch_query('CREATE INDEX {0} ON {1} (x);'.format(cls.index, cls.table2)))
+        finally:
+            shutdown(client)
+
+    @classmethod
+    def tearDownClass(cls):
+        try:
+            # delete the database objects created
+            client = create_mssql_cli_client()
+            list(client.execute_single_batch_query('DROP DATABASE {0};'.format(cls.database)))
+            list(client.execute_single_batch_query('DROP INDEX {0} ON {1};'.format(cls.index, cls.table2)))
+            list(client.execute_single_batch_query('DROP TABLE {0};'.format(cls.table1)))
+            list(client.execute_single_batch_query('DROP TABLE {0};'.format(cls.table2)))
+            list(client.execute_single_batch_query('DROP VIEW {0} IF EXISTS;'.format(cls.view)))
+            list(client.execute_single_batch_query('DROP SCHEMA {0};'.format(cls.schema)))
+        finally:
+            shutdown(client)
+
+    def test_list_tables_command(self):
+        self.command('\\dt', self.table1, min_rows_expected=2, rows_expected_pattern_query=1, cols_expected=2,
+                     cols_expected_verbose=4)
+
+    def test_list_views_command(self):
+        self.command('\\dv', self.view, min_rows_expected=1, rows_expected_pattern_query=1, cols_expected=2,
+                     cols_expected_verbose=3)
+
+    def test_list_schemas_command(self):
+        self.command('\\dn', self.schema, min_rows_expected=1, rows_expected_pattern_query=1, cols_expected=1,
+                     cols_expected_verbose=3)
+
+    def test_list_indices_command(self):
+        self.command('\\di', self.index, min_rows_expected=1, rows_expected_pattern_query=1, cols_expected=3,
+                     cols_expected_verbose=5)
+
+    def test_list_databases_command(self):
+        self.command('\\l', self.database, min_rows_expected=1, rows_expected_pattern_query=1, cols_expected=1,
+                     cols_expected_verbose=4)
+
+    def test_add_new_special_command(self):
+        @special_command('\\empty', '\\empty[+]', 'returns an empty list', arg_type=NO_QUERY)
+        def empty_list_special_command():
+            return []
+
+        ret = execute(None, '\\empty')
+        self.assertTrue(len(ret) == 0)
+
+    def command(self, command, pattern, min_rows_expected, rows_expected_pattern_query,
+                cols_expected, cols_expected_verbose):
+        try:
+            client = create_mssql_cli_client()
+
+            for rows, col, message, query, is_error in \
+                    client.execute_multi_statement_single_batch(command):
+                self.assertTrue(len(rows) >= min_rows_expected)
+                self.assertTrue(len(col) == cols_expected)
+
+            # execute with pattern and verbose
+            command = command + "+ " + pattern
+            for rows, col, message, query, is_error in \
+                    client.execute_multi_statement_single_batch(command):
+                self.assertTrue(len(rows) == rows_expected_pattern_query)
+                self.assertTrue(len(col) == cols_expected_verbose)
+        finally:
+            shutdown(client)
+
+
+
+if __name__ == u'__main__':
+    unittest.main()
