[CLOUD-1366]: add support for impersonating users (#150)

stardog/admin.py

@@ -26,6 +26,7 @@ def __init__(
         username: object = None,
         password: object = None,
         auth: object = None,
+        run_as: str = None,
     ) -> None:
         """Initializes an admin connection to a Stardog server.
 
@@ -36,16 +37,19 @@ def __init__(
             Defaults to `admin`
           password (str, optional): Password to use in the connection.
             Defaults to `admin`
-        auth (requests.auth.AuthBase, optional): requests Authentication object.
+          auth (requests.auth.AuthBase, optional): requests Authentication object.
             Defaults to `None`
+          run_as (str, optional): the User to impersonate
 
         auth and username/password should not be used together.  If the are the value
         of `auth` will take precedent.
         Examples:
           >>> admin = Admin(endpoint='http://localhost:9999',
                             username='admin', password='admin')
         """
-        self.client = client.Client(endpoint, None, username, password, auth=auth)
+        self.client = client.Client(
+            endpoint, None, username, password, auth=auth, run_as=run_as
+        )
         # ensure the server is alive and at the specified location
         self.alive()
 

stardog/connection.py

@@ -25,6 +25,7 @@ def __init__(
         password=None,
         auth=None,
         session=None,
+        run_as=None,
     ):
         """Initializes a connection to a Stardog database.
 
@@ -38,13 +39,20 @@ def __init__(
             Defaults to `None`
           session (requests.session.Session, optional): requests Session object.
             Defaults to `None`
+          run_as (str, optional): the user to impersonate
 
         Examples:
           >>> conn = Connection('db', endpoint='http://localhost:9999',
                                 username='admin', password='admin')
         """
         self.client = client.Client(
-            endpoint, database, username, password, auth=auth, session=session
+            endpoint,
+            database,
+            username,
+            password,
+            auth=auth,
+            session=session,
+            run_as=run_as,
         )
         self.transaction = None
 

stardog/http/client.py

@@ -18,6 +18,7 @@ def __init__(
         password=None,
         session=None,
         auth=None,
+        run_as=None,
     ):
         self.url = endpoint if endpoint else self.DEFAULT_ENDPOINT
 
@@ -45,6 +46,9 @@ def __init__(
             )
         self.session.auth = auth
 
+        if run_as:
+            self.session.headers.update({"SD-Run-As": run_as})
+
     def post(self, path, **kwargs):
         return self.__wrap(self.session.post(self.url + path, **kwargs))
 

test/test_admin_basic.py

@@ -88,6 +88,28 @@ def music_options(self):
         }
 
 
+class TestUserImpersonation(TestStardog):
+    def test_impersonating_user_databases_visibility(self, conn_string, user, db):
+
+        with admin.Admin(
+            endpoint=conn_string["endpoint"],
+        ) as admin_user:
+            databases_admin_can_see = [db.name for db in admin_user.databases()]
+        with admin.Admin(
+            endpoint=conn_string["endpoint"],
+            run_as=user.name,
+        ) as admin_impersonating_user:
+            databases_impersonated_user_can_see = [
+                db.name for db in admin_impersonating_user.databases()
+            ]
+        assert len(databases_impersonated_user_can_see) == 0
+
+        # for cluster tests in Circle, catalog is disabled so the exact number of dbs
+        # varies (2 for single node, 1 for cluster since catalog isn't created)
+        assert len(databases_admin_can_see) > 0
+        assert db.name in databases_admin_can_see
+
+
 class TestUsers(TestStardog):
     def test_user_creation(self, admin, user):
         assert len(admin.users()) == len(default_users) + 1

test/test_connection.py

@@ -1,6 +1,6 @@
 import pytest
 
-from stardog import admin, connection, content, content_types, exceptions
+from stardog import connection, content, content_types, exceptions
 
 
 def starwars_contents() -> list:
@@ -100,6 +100,28 @@ def test_export(db, conn):
     assert b"default_obj" not in named_export
 
 
+def test_user_impersonation(conn_string, db, user):
+    with connection.Connection(
+        endpoint=conn_string["endpoint"], database=db.name
+    ) as admin_regular_conn:
+        # add some data to query
+        admin_regular_conn.begin()
+        admin_regular_conn.clear()
+        admin_regular_conn.add(content.File("test/data/starwars.ttl"))
+        admin_regular_conn.commit()
+
+        # confirm admin can read the data
+        q = admin_regular_conn.select('select * {?s :name "Luke Skywalker"}')
+        assert len(q["results"]["bindings"]) == 1
+
+    # attempting to query database as the user the admin is impersonating should return an unauthorized error
+    with connection.Connection(
+        endpoint=conn_string["endpoint"], database=db.name, run_as=user.name
+    ) as admin_impersonating_user:
+        with pytest.raises(exceptions.StardogException, match="403"):
+            admin_impersonating_user.select('select * {?s :name "Luke Skywalker"}')
+
+
 @pytest.mark.dbname("pystardog-test-database")
 @pytest.mark.conn_dbname("pystardog-test-database")
 def test_queries(db, conn):