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" "USER" ("IF" "NOT" "EXISTS")? <stringLiteral>
+                            "TO" "ROLE" <rolename>
+                            ;
+<dropIdentityStatement> ::= "DROP" "USER" ("IF" "EXISTS")? <stringLiteral>
+                            ;
+'''
+
 # END SYNTAX/COMPLETION RULE DEFINITIONS
 
 CqlRuleSet.append_rules(syntax_rules)

pylib/cqlshlib/test/test_identity_mappings_completion.py

@@ -0,0 +1,72 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from .test_cqlsh_completion import CqlshCompletionCase
+
+
+class TestIdentityMappingsCompletion(CqlshCompletionCase):
+
+    """
+        Autcomplete Tests for ADD and DROP Identity Mappings for @cql3handling.py
+    """
+
+    def add_test_user_autocomplete(self):
+        self.trycompletions('ADD U', immediate='SER ')
+
+
+    def add_test_exists_autocomplete(self):
+        self.trycompletions('ADD USER IF ', immediate='NOT EXISTS ')
+        self.trycompletions('ADD USER IF NOT ', immediate='EXISTS ')
+
+    def test_expect_str_literal_add_autocomplete(self):
+        self.trycompletions('ADD USER ',
+            choices=['<pgStringLiteral>', '<quotedStringLiteral>', 'IF'])
+
+    def add_test_to_autocomplete(self):
+        self.trycompletions("ADD USER '[email protected]' ", immediate='TO ')
+        self.trycompletions("ADD USER '[email protected]' T", immediate='O ')
+
+    def add_test_role_autocomplete(self):
+        self.trycompletions("ADD USER '[email protected]' TO ", immediate='ROLE ')
+        self.trycompletions("ADD USER '[email protected]' TO R", immediate='OLE ')
+
+    def add_test_rolename_autocomplete(self):
+        self.trycompletions("ADD USER '[email protected]' TO ROLE ",
+                            choices=['<identifier>', '<quotedName>'],
+                            other_choices_ok=True)
+
+    def add_test_complete_user_full_statement(self):
+        self.trycompletions("ADD USER IF NOT EXISTS '[email protected]' TO ROLE data_engineer ",
+                            choices=[';'])
+
+    def drop_test_autocomplete(self):
+        self.trycompletions('DROP ',
+                            choices=['AGGREGATE', 'COLUMNFAMILY', 'FUNCTION',
+                                     'INDEX', 'KEYSPACE', 'ROLE', 'TABLE',
+                                     'TRIGGER', 'TYPE', 'USER', 'MATERIALIZED'])
+
+    def drop_test_user_autcomplete(self):
+        self.trycompletions('DROP USER ',
+                            choices=['<pgStringLiteral>', '<quotedStringLiteral>', 'IF'])
+
+    def drop_test_exists(self):
+        self.trycompletions('DROP USER IF ', immediate='EXISTS ')
+        self.trycompletions('DROP USER IF EXISTS ',
+                            choices=['<pgStringLiteral>', '<quotedStringLiteral>'])
+
+    def test_complete_drop_user_statement_end(self):
+        self.trycompletions("DROP USER '[email protected]' ", choices=[';'])
+        self.trycompletions("DROP USER IF EXISTS '[email protected]' ", choices=[';'])