Merge pull request #19 from tristeng/master

Added in some new functions...

.gitignore

@@ -1,3 +1,4 @@
 *.pyc
 .coverage
 htmlcov/
+.idea

pyaccumulo/__init__.py

@@ -233,3 +233,74 @@ def create_batch_writer(self, table, max_memory=10*1024, latency_ms=30*1000, tim
     def delete_rows(self, table, srow, erow):
         self.client.deleteRows(self.login, table, srow, erow)
 
+    def attach_iterator(self, table, setting, scopes):
+        self.client.attachIterator(self.login, table, setting, scopes)
+
+    def remove_iterator(self, table, iterator, scopes):
+        self.client.removeIterator(self.login, table, iterator, scopes)
+
+    def following_key(self, key, part):
+        return self.client.getFollowing(key, part)
+
+    def get_max_row(self, table, auths=None, srow=None, sinclude=None, erow=None, einclude=None):
+        return self.client.getMaxRow(self.login, table, auths, srow, sinclude, erow, einclude)
+
+    def add_mutations_and_flush(self, table, muts):
+        """
+        Add mutations to a table without the need to create and manage a batch writer.
+        """
+        if not isinstance(muts, list) and not isinstance(muts, tuple):
+            muts = [muts]
+        cells = {}
+        for mut in muts:
+            cells.setdefault(mut.row, []).extend(mut.updates)
+        self.client.updateAndFlush(self.login, table, cells)
+
+    def create_user(self, user, password):
+        self.client.createLocalUser(self.login, user, password)
+
+    def drop_user(self, user):
+        self.client.dropLocalUser(self.login, user)
+
+    def list_users(self):
+        return self.client.listLocalUsers(self.login)
+
+    def set_user_authorizations(self, user, auths):
+        self.client.changeUserAuthorizations(self.login, user, auths)
+
+    def get_user_authorizations(self, user):
+        return self.client.getUserAuthorizations(self.login, user)
+
+    def grant_system_permission(self, user, perm):
+        self.client.grantSystemPermission(self.login, user, perm)
+
+    def revoke_system_permission(self, user, perm):
+        self.client.revokeSystemPermission(self.login, user, perm)
+
+    def has_system_permission(self, user, perm):
+        return self.client.hasSystemPermission(self.login, user, perm)
+
+    def grant_table_permission(self, user, table, perm):
+        self.client.grantTablePermission(self.login, user, table, perm)
+
+    def revoke_table_permission(self, user, table, perm):
+        self.client.revokeTablePermission(self.login, user, table, perm)
+
+    def has_table_permission(self, user, table, perm):
+        return self.client.hasTablePermission(self.login, user, table, perm)
+
+    def add_splits(self, table, splits):
+        self.client.addSplits(self.login, table, splits)
+
+    def add_constraint(self, table, class_name):
+        return self.client.addConstraint(self.login, table, class_name)
+
+    def list_constraints(self, table):
+        return self.client.listConstraints(self.login, table)
+
+    def remove_constraint(self, table, constraint):
+        """
+        :param table: table name
+        :param constraint: the constraint number as returned by list_constraints
+        """
+        self.client.removeConstraint(self.login, table, constraint)

tests/test_core.py

@@ -19,8 +19,16 @@
 import pyaccumulo
 from pyaccumulo import *
 from mock import Mock
+from pyaccumulo.proxy.ttypes import IteratorScope, PartialKey, SystemPermission, TablePermission
+
 
 class AccumuloTest(unittest.TestCase):
+    def _get_mock_connection(self):
+        conn = Accumulo(_connect=False)
+        conn.client = Mock()
+        conn.login = "Login"
+        return conn
+
     def test_get_scan_columns(self):
         self.assertEquals(None, pyaccumulo._get_scan_columns([]))
         self.assertEquals([ ScanColumn(colFamily="cf", colQualifier="cq") ], pyaccumulo._get_scan_columns([ ["cf", "cq"]]))
@@ -31,9 +39,7 @@ def test_following_array(self):
         self.assertEquals(None, pyaccumulo.following_array(None))
 
     def test_list_tables(self):
-        conn = Accumulo(_connect=False)
-        conn.client = Mock()
-        conn.login = "Login"
+        conn = self._get_mock_connection()
         conn.client.listTables = Mock()
         conn.client.listTables.return_value = set(["t1", "t2", "t3"])
 
@@ -43,9 +49,7 @@ def test_list_tables(self):
 
 
     def test_table_exists(self):
-        conn = Accumulo(_connect=False)
-        conn.client = Mock()
-        conn.login = "Login"
+        conn = self._get_mock_connection()
         conn.client.tableExists = Mock()
         conn.client.tableExists.return_value = True
 
@@ -55,27 +59,21 @@ def test_table_exists(self):
 
 
     def test_create_table(self):
-        conn = Accumulo(_connect=False)
-        conn.client = Mock()
-        conn.login = "Login"
+        conn = self._get_mock_connection()
         conn.client.createTable = Mock()
         conn.client.createTable.return_value = True
 
         conn.create_table("mytable")
         conn.client.createTable.assert_called_with("Login", "mytable", True, TimeType.MILLIS)        
 
     def test_delete_table(self):
-        conn = Accumulo(_connect=False)
-        conn.client = Mock()
-        conn.login = "Login"
+        conn = self._get_mock_connection()
         conn.client.deleteTable = Mock()
         conn.delete_table("mytable")
         conn.client.deleteTable.assert_called_with("Login", "mytable")
 
     def test_rename_table(self):
-        conn = Accumulo(_connect=False)
-        conn.client = Mock()
-        conn.login = "Login"
+        conn = self._get_mock_connection()
         conn.client.renameTable = Mock()
         conn.rename_table("mytable", "newtable")
         conn.client.renameTable.assert_called_with("Login", "mytable", "newtable")
@@ -157,6 +155,142 @@ def test_close(self):
         conn.close()
         conn.transport.close.assert_called_with()
 
+    def test_attach_iterator(self):
+        conn = self._get_mock_connection()
+        conn.client.attachIterator = Mock()
+        iter_setting = IteratorSetting(name="i1")
+        scopes = {IteratorScope.MAJC, IteratorScope.MINC, IteratorScope.SCAN}
+        conn.attach_iterator("mytable", iter_setting, scopes)
+        conn.client.attachIterator.assert_called_with("Login", "mytable", iter_setting, scopes)
+
+    def test_remove_iterator(self):
+        conn = self._get_mock_connection()
+        conn.client.removeIterator = Mock()
+        scopes = {IteratorScope.MAJC, IteratorScope.MINC, IteratorScope.SCAN}
+        conn.remove_iterator("mytable", "i1", scopes)
+        conn.client.removeIterator.assert_called_with("Login", "mytable", "i1", scopes)
+
+    def test_following_key(self):
+        conn = self._get_mock_connection()
+        key = Key(row="aabbcc", colFamily="cf", colQualifier="cq")
+        conn.client.getFollowing = Mock(return_value=following_key(key))
+        fkey = conn.following_key(key, PartialKey.ROW_COLFAM_COLQUAL_COLVIS)
+        conn.client.getFollowing.assert_called_with(key, PartialKey.ROW_COLFAM_COLQUAL_COLVIS)
+        self.assertEqual(fkey, Key(row="aabbcc", colFamily="cf", colQualifier="cq\0"))
+
+    def test_get_max_row(self):
+        conn = self._get_mock_connection()
+        conn.client.getMaxRow = Mock(return_value="aabe")
+        row = conn.get_max_row("mytable", auths=["auth1", "auth2"], srow="aabb", sinclude=True, erow="bbaa", einclude=False)
+        conn.client.getMaxRow.assert_called_with("Login", "mytable", ["auth1", "auth2"], "aabb", True, "bbaa", False)
+        self.assertEqual(row, "aabe")
+
+    def test_add_mutations_and_flush(self):
+        conn = self._get_mock_connection()
+        conn.client.updateAndFlush = Mock()
+        muts = [Mutation("aabb"), Mutation("bbaa")]
+        conn.add_mutations_and_flush("mytable", muts)
+        conn.client.updateAndFlush.assert_called_with("Login", "mytable", {'aabb': [], 'bbaa': []})
+        conn.add_mutations_and_flush("mytable", Mutation("ccdd"))
+        conn.client.updateAndFlush.assert_called_with("Login", "mytable", {'ccdd': []})
+
+    def test_create_user(self):
+        conn = self._get_mock_connection()
+        conn.createLocalUser = Mock()
+        conn.create_user("user", "password")
+        conn.client.createLocalUser.assert_called_with("Login", "user", "password")
+
+    def test_drop_user(self):
+        conn = self._get_mock_connection()
+        conn.client.dropLocalUser = Mock()
+        conn.drop_user("user")
+        conn.client.dropLocalUser.assert_called_with("Login", "user")
+
+    def test_list_users(self):
+        conn = self._get_mock_connection()
+        conn.client.listLocalUsers = Mock(return_value=["user1", "user2"])
+        users = conn.list_users()
+        conn.client.listLocalUsers.assert_called_with("Login")
+        self.assertEqual(users, ["user1", "user2"])
+
+    def test_set_user_authorizations(self):
+        conn = self._get_mock_connection()
+        conn.client.changeUserAuthorizations = Mock()
+        conn.set_user_authorizations("user", ["auth1", "auth2"])
+        conn.client.changeUserAuthorizations.assert_called_with("Login", "user", ["auth1", "auth2"])
+
+    def test_get_user_authorizations(self):
+        conn = self._get_mock_connection()
+        conn.client.getUserAuthorizations = Mock(return_value=["auth1", "auth2"])
+        auths = conn.get_user_authorizations("user1")
+        conn.client.getUserAuthorizations.assert_called_with("Login", "user1")
+        self.assertEqual(auths, ["auth1", "auth2"])
+
+    def test_grant_system_permission(self):
+        conn = self._get_mock_connection()
+        conn.client.grantSystemPermission = Mock()
+        conn.grant_system_permission("user1", SystemPermission.CREATE_TABLE)
+        conn.client.grantSystemPermission.assert_called_with("Login", "user1", SystemPermission.CREATE_TABLE)
+
+    def test_revoke_system_permission(self):
+        conn = self._get_mock_connection()
+        conn.client.revokeSystemPermission = Mock()
+        conn.revoke_system_permission("user1", SystemPermission.DROP_USER)
+        conn.client.revokeSystemPermission.assert_called_with("Login", "user1", SystemPermission.DROP_USER)
+
+    def test_has_system_permission(self):
+        conn = self._get_mock_connection()
+        conn.client.hasSystemPermission = Mock(return_value=True)
+        has_perm = conn.has_system_permission("user1", SystemPermission.SYSTEM)
+        conn.client.hasSystemPermission.assert_called_with("Login", "user1", SystemPermission.SYSTEM)
+        self.assertTrue(has_perm)
+
+    def test_grant_table_permission(self):
+        conn = self._get_mock_connection()
+        conn.client.grantTablePermission = Mock()
+        conn.grant_table_permission("user1", "mytable", TablePermission.BULK_IMPORT)
+        conn.client.grantTablePermission.assert_called_with("Login", "user1", "mytable", TablePermission.BULK_IMPORT)
+
+    def test_revoke_table_permission(self):
+        conn = self._get_mock_connection()
+        conn.client.revokeTablePermission = Mock()
+        conn.revoke_table_permission("user1", "mytable", TablePermission.DROP_TABLE)
+        conn.client.revokeTablePermission.assert_called_with("Login", "user1", "mytable", TablePermission.DROP_TABLE)
+
+    def test_has_table_permission(self):
+        conn = self._get_mock_connection()
+        conn.client.hasTablePermission = Mock(return_value=False)
+        has_perm = conn.has_table_permission("user1", "mytable", TablePermission.GRANT)
+        conn.client.hasTablePermission.assert_called_with("Login", "user1", "mytable", TablePermission.GRANT)
+        self.assertFalse(has_perm)
+
+    def test_add_splits(self):
+        conn = self._get_mock_connection()
+        conn.client.addSplits = Mock()
+        conn.add_splits("mytable", ["aa", "bb", "cc"])
+        conn.client.addSplits.assert_called_with("Login", "mytable", ["aa", "bb", "cc"])
+
+    def test_add_constraints(self):
+        conn = self._get_mock_connection()
+        conn.client.addConstraint = Mock(return_value=True)
+        ret_val = conn.add_constraint("mytable", "org.apache.accumulo.examples.simple.constraints.MaxMutationSize")
+        conn.client.addConstraint.assert_called_with(
+            "Login", "mytable", "org.apache.accumulo.examples.simple.constraints.MaxMutationSize")
+        self.assertTrue(ret_val)
+
+    def test_list_constraints(self):
+        conn = self._get_mock_connection()
+        conn.client.listConstraints = Mock(return_value={"c1": 1, "c2": 1})
+        constraints = conn.list_constraints("mytable")
+        conn.client.listConstraints.assert_called_with("Login", "mytable")
+        self.assertEqual(constraints, {"c1": 1, "c2": 1})
+
+    def test_remove_constraint(self):
+        conn = self._get_mock_connection()
+        conn.client.removeConstraint = Mock()
+        conn.remove_constraint("mytable", 1)
+        conn.client.removeConstraint.assert_called_with("Login", "mytable", 1)
+
 class RangeTest(unittest.TestCase):
     def test_to_range(self):