Adding support for GraphQL

Closes #265

pyapacheatlas/core/client.py

@@ -4,6 +4,7 @@
 from .discovery.purview import PurviewDiscoveryClient
 from .typedef import BaseTypeDef, TypeCategory
 from .msgraph import MsGraphClient
+from .graphql import GraphQLClient
 from .entity import AtlasClassification, AtlasEntity
 from ..auth.base import AtlasAuthBase
 import logging
@@ -1576,6 +1577,10 @@ def __init__(self, account_name, authentication=None, **kwargs):
         self.discovery = PurviewDiscoveryClient(
             f"https://{account_name.lower()}.purview.azure.com/catalog/api",
             authentication, requests_args=requests_args)
+        self.graphql = GraphQLClient(
+            endpoint_url = f"https://{account_name.lower()}.purview.azure.com/datamap/api/graphql",
+            authentication=authentication
+        )
         super().__init__(endpoint_url, authentication,
                          glossary=glossary, requests_args=requests_args,
                          **kwargs)

pyapacheatlas/core/graphql.py

@@ -0,0 +1,83 @@
+from .. import __version__
+import json
+from typing import Union, List
+from json import JSONDecodeError
+import requests
+from .util import AtlasResponse
+
+
+class GraphQLException(BaseException):
+    pass
+
+
+class GraphQLClient():
+    """
+    Implements the GraphQL connection in Microsoft Purview.
+
+    Do not instantiate this directly, instead, it is used by PurviewClient.
+    This is not available to that AtlasClient.
+    """
+
+    def __init__(self, endpoint_url, authentication, **kwargs):
+        super().__init__()
+        self.endpoint_url = endpoint_url
+        self.authentication = authentication
+        self._requests_args = kwargs.get("requests_args", {})
+        self._USER_AGENT = {"User-Agent": "pyapacheatlas/{0} {1}".format(
+            __version__, requests.utils.default_headers().get("User-Agent"))}
+
+    def _generate_request_headers(self, include: dict = {}, exclude: List[str] = []):
+        auth = {} if self.authentication is None else self.authentication.get_authentication_headers()
+
+        if include:
+            auth.update(include)
+        if exclude:
+            for key in exclude:
+                if key in auth:
+                    auth.pop(key)
+        return dict(**auth, **self._USER_AGENT)
+
+    def _post_http(self, url: str, params: dict = None,
+                   json: Union[list, dict] = None, files: dict = None,
+                   **kwargs) -> AtlasResponse:
+        """
+        :kwargs dict headers_include:Additional headers to include.
+        :kwargs List[str] headers_include:Additional headers to include.
+        """
+        extra_args = {}
+        if json:
+            extra_args["json"] = json
+        if params:
+            extra_args["params"] = params
+        if files:
+            extra_args["files"] = files
+        response_args = {}
+        if "responseNotJson" in kwargs:
+            response_args["responseNotJson"] = kwargs["responseNotJson"]
+        return AtlasResponse(
+            requests.post(
+                url,
+                headers=self._generate_request_headers(kwargs.get(
+                    "headers_include"), kwargs.get("headers_exclude")),
+                **extra_args,
+                **self._requests_args
+            ),
+            **response_args
+        )
+
+    def query(self, query:str):
+        """
+        Execute a GraphQL query on the Purview Datamap.
+
+        :param str query:
+
+        :return: The GraphQL response object
+        :rtype: dict
+
+        [Reference](https://learn.microsoft.com/en-us/purview/tutorial-graphql-api)
+        """
+        graphQL_response = self._post_http(
+            self.endpoint_url,
+            json={"query":query}
+        )
+        return graphQL_response.body

samples/CRUD/read_graphql_basic.py

@@ -0,0 +1,39 @@
+import json
+import os
+
+from azure.identity import AzureCliCredential
+from pyapacheatlas.core import PurviewClient
+
+if __name__ == "__main__":
+    """
+    This sample demonstrates how to use the Purview GraphQL endpoint.
+
+    For more sample queries see: https://learn.microsoft.com/en-us/purview/tutorial-graphql-api
+    """
+    cred = AzureCliCredential()
+    client = PurviewClient(
+            account_name=os.environ.get("PURVIEW_NAME", "InsertDefaultPurviewAccountName"),
+            authentication=cred
+        )
+
+    graphql_query = """
+    query {
+        entities(where: { guid: ["a8bd4ed2-3c2a-4745-b7bb-68f6f6f60000"] }) {
+            guid
+            createTime
+            updateTime
+            typeName
+            attributes
+            name
+            qualifiedName
+            description
+        }
+    }
+    """
+
+    # Call the query endpoint with your GraphQL query
+    resp = client.graphql.query(graphql_query)
+
+    # Iterate over the results inside of the data and entities objects
+    for entity in resp["data"]["entities"]:
+        print(json.dumps(entity))