CASSANDRA-20021 Add cqlsh autocompletion for the identity mapping feature

pylib/cqlshlib/cql3handling.py

@@ -286,6 +286,7 @@ def dequote_value(cqlword):
                           | <createFunctionStatement>
                           | <createAggregateStatement>
                           | <createTriggerStatement>
+                          | <addIdentityStatement>
                           | <dropKeyspaceStatement>
                           | <dropColumnFamilyStatement>
                           | <dropIndexStatement>
@@ -294,6 +295,7 @@ def dequote_value(cqlword):
                           | <dropFunctionStatement>
                           | <dropAggregateStatement>
                           | <dropTriggerStatement>
+                          | <dropIdentityStatement>
                           | <alterTableStatement>
                           | <alterKeyspaceStatement>
                           | <alterUserTypeStatement>
@@ -1749,6 +1751,7 @@ def rolename_completer(ctxt, cass):
                              "ON" cf=<columnFamilyName>
                          ;
 '''
+
 explain_completion('createTriggerStatement', 'class', '\'fully qualified class name\'')
 
 
@@ -1765,6 +1768,14 @@ def drop_trigger_completer(ctxt, cass):
     return list(map(maybe_escape_name, names))
 
 
+syntax_rules += r'''
+<addIdentityStatement> ::= "ADD" "IDENTITY" ("IF" "NOT" "EXISTS")? <stringLiteral>
+                            "TO" "ROLE" <rolename>
+                            ;
+<dropIdentityStatement> ::= "DROP" "IDENTITY" ("IF" "EXISTS")? <stringLiteral>
+                            ;
+'''
+
 # END SYNTAX/COMPLETION RULE DEFINITIONS
 
 CqlRuleSet.append_rules(syntax_rules)

pylib/cqlshlib/test/test_cqlsh_completion.py

@@ -166,7 +166,7 @@ class TestCqlshCompletion(CqlshCompletionCase):
     cqlver = '3.1.6'
 
     def test_complete_on_empty_string(self):
-        self.trycompletions('', choices=('?', 'ALTER', 'BEGIN', 'CAPTURE', 'CONSISTENCY',
+        self.trycompletions('', choices=('?', 'ADD', 'ALTER', 'BEGIN', 'CAPTURE', 'CONSISTENCY',
                                          'COPY', 'CREATE', 'DEBUG', 'DELETE', 'DESC', 'DESCRIBE',
                                          'DROP', 'GRANT', 'HELP', 'INSERT', 'LIST', 'LOGIN', 'PAGING', 'REVOKE',
                                          'SELECT', 'SHOW', 'SOURCE', 'TRACING', 'ELAPSED', 'EXPAND', 'SERIAL', 'TRUNCATE',
@@ -613,7 +613,7 @@ def test_complete_in_string_literals(self):
     def test_complete_in_drop(self):
         self.trycompletions('DR', immediate='OP ')
         self.trycompletions('DROP ',
-                            choices=['AGGREGATE', 'COLUMNFAMILY', 'FUNCTION',
+                            choices=['AGGREGATE', 'COLUMNFAMILY', 'FUNCTION', 'IDENTITY',
                                      'INDEX', 'KEYSPACE', 'ROLE', 'TABLE',
                                      'TRIGGER', 'TYPE', 'USER', 'MATERIALIZED'])
 
@@ -1048,12 +1048,12 @@ def test_complete_in_use(self):
                                              'system_virtual_schema', 'system_cluster_metadata'])
 
     def test_complete_in_create_index(self):
-        self.trycompletions('CREATE I', immediate='NDEX ')
+        self.trycompletions('CREATE IN', immediate='DEX ')
         self.trycompletions('CREATE INDEX ', choices=['<new_index_name>', 'IF', 'ON'])
         self.trycompletions('CREATE INDEX example ', immediate='ON ')
 
     def test_complete_in_drop_index(self):
-        self.trycompletions('DROP I', immediate='NDEX ')
+        self.trycompletions('DROP IN', immediate='DEX ')
 
     def test_complete_in_alter_keyspace(self):
         self.trycompletions('ALTER KEY', 'SPACE ')
@@ -1184,9 +1184,15 @@ def test_complete_in_alter_type(self):
         self.trycompletions('ALTER TYPE IF EXISTS new_type ADD IF NOT EXISTS ', choices=['<new_field_name>'])
         self.trycompletions('ALTER TYPE IF EXISTS new_type RENAME ', choices=['IF', '<quotedName>', '<identifier>'])
 
+    # USER checks
+    def test_complete_in_create_user(self):
+        self.trycompletions('CREATE USER ', choices=['<username>', 'IF'])
+        self.trycompletions('CREATE USER IF ', immediate='NOT EXISTS ')
+
     def test_complete_in_alter_user(self):
         self.trycompletions('ALTER USER ', choices=['<identifier>', 'IF', '<pgStringLiteral>', '<quotedStringLiteral>'])
 
+    # ROLE checks
     def test_complete_in_create_role(self):
         self.trycompletions('CREATE ROLE ', choices=['<rolename>', 'IF'])
         self.trycompletions('CREATE ROLE IF ', immediate='NOT EXISTS ')
@@ -1205,18 +1211,45 @@ def test_complete_in_alter_role(self):
         self.trycompletions('ALTER ROLE foo WITH ACCESS TO ', choices=['ALL', 'DATACENTERS'])
         self.trycompletions('ALTER ROLE foo WITH ACCESS FROM ', choices=['ALL', 'CIDRS'])
 
-    def test_complete_in_create_user(self):
-        self.trycompletions('CREATE USER ', choices=['<username>', 'IF'])
-        self.trycompletions('CREATE USER IF ', immediate='NOT EXISTS ')
 
     def test_complete_in_drop_role(self):
         self.trycompletions('DROP ROLE ', choices=['<identifier>', 'IF', '<quotedName>'])
 
+
+    # IDENTITY checks
+    def test_complete_in_add_identity(self):
+        self.trycompletions('ADD ID', immediate='ENTITY ')
+        self.trycompletions('ADD IDENTITY IF ', immediate='NOT EXISTS ')
+        self.trycompletions('ADD IDENTITY IF NOT ', immediate='EXISTS ')
+        self.trycompletions('ADD IDENTITY ',
+            choices=['<pgStringLiteral>', '<quotedStringLiteral>', 'IF'])
+        self.trycompletions("ADD IDENTITY '[email protected]' TO R", immediate='OLE ')
+        self.trycompletions("ADD IDENTITY '[email protected]' ", immediate='TO ROLE ')
+        self.trycompletions("ADD IDENTITY '[email protected]' TO ROLE ",
+                            choices=['<identifier>', '<quotedName>'], other_choices_ok=True)
+        self.trycompletions("ADD IDENTITY '[email protected]' TO ROLE data_engineer ",
+                                    choices=[';'])
+        self.trycompletions("ADD IDENTITY IF NOT EXISTS '[email protected]' TO ROLE data_engineer ",
+                            choices=[';'])
+
+    def test_complete_in_drop_identity(self):
+        self.trycompletions('DROP IDENTITY ',
+                            choices=['<pgStringLiteral>', '<quotedStringLiteral>', 'IF'])
+        self.trycompletions('DROP IDENTITY IF ', immediate='EXISTS ')
+        self.trycompletions('DROP IDENTITY IF EXISTS ',
+                            choices=['<pgStringLiteral>', '<quotedStringLiteral>'])
+        self.trycompletions("DROP IDENTITY '[email protected]' ", choices=[';'])
+        self.trycompletions("DROP IDENTITY IF EXISTS '[email protected]' ", choices=[';'])
+
     def test_complete_in_list(self):
         self.trycompletions('LIST ',
                             choices=['ALL', 'AUTHORIZE', 'DESCRIBE', 'EXECUTE', 'ROLES', 'USERS', 'ALTER',
                                      'CREATE', 'DROP', 'MODIFY', 'SELECT', 'UNMASK', 'SELECT_MASKED', 'SUPERUSERS'])
 
+    def test_complete_drop_end(self):
+        self.trycompletions("DROP USER '[email protected]' ", choices=[';'])
+        self.trycompletions("DROP USER IF EXISTS '[email protected]' ", choices=[';'])
+
     # Non-CQL Shell Commands
 
     def test_complete_in_capture(self):