Merge pull request #342 from PrefectHQ/graphql

Add utility for building GQL queries from Python dicts/sequences

src/prefect/utilities/graphql.py

@@ -0,0 +1,71 @@
+import textwrap
+from typing import Any
+
+
+def lowercase_first_letter(s: str) -> str:
+    """
+    Given a string, returns that string with a lowercase first letter
+    """
+    if s:
+        return s[0].lower() + s[1:]
+    return s
+
+
+class GQLObject:
+    """
+    Helper object for building GraphQL queries.
+    """
+
+    def __init__(self, name: str = None, _arguments: str = None):
+        self.__name = name or lowercase_first_letter(type(self).__name__)
+        self.__arguments = _arguments
+
+    def __call__(self, arguments: str) -> "GQLObject":
+        return type(self)(name=self.__name, _arguments=arguments)
+
+    def __repr__(self) -> str:
+        return '<GQL: "{name}">'.format(self.__name)
+
+    def __str__(self) -> str:
+        if not self.__arguments:
+            return self.__name
+        else:
+            return "{name}({arguments})".format(
+                name=self.__name, arguments=self.__arguments
+            )
+
+
+def parse_graphql(document: Any) -> str:
+    delimiter = "    "
+    parsed = _parse_graphql_inner(document, delimiter=delimiter)
+    parsed = parsed.replace(delimiter + "}", "}")
+    parsed = textwrap.dedent(parsed).strip()
+    return parsed
+
+
+def _parse_graphql_inner(document: Any, delimiter: str) -> str:
+    """
+    Inner loop function of for `parse_graphql`.
+    """
+    if isinstance(document, (tuple, list, set)):
+        return "\n".join(
+            [_parse_graphql_inner(item, delimiter=delimiter) for item in document]
+        )
+    elif isinstance(document, dict):
+        result = [
+            "{key} {{\n{value}\n}}".format(
+                key=_parse_graphql_inner(key, delimiter=delimiter),
+                value=_parse_graphql_inner(value, delimiter=delimiter),
+            )
+            for key, value in document.items()
+        ]
+        return _parse_graphql_inner(result, delimiter=delimiter)
+    elif isinstance(document, type) and issubclass(document, GQLObject):
+        raise TypeError(
+            'It looks like you included a `GQLObject` class ("{name}") '
+            "in your document. Did you mean to use an instance of that type?".format(
+                name=document.__name__
+            )
+        )
+    else:
+        return str(document).replace("\n", "\n" + delimiter)

tests/utilities/test_graphql.py

@@ -0,0 +1,181 @@
+from textwrap import dedent
+import pytest
+from prefect.utilities.graphql import GQLObject, parse_graphql
+from collections import OrderedDict
+
+
+class Account(GQLObject):
+    id = "id"
+    name = "name"
+
+
+class User(GQLObject):
+    id = "id"
+    name = "name"
+    account = Account("account")
+
+
+# avoid circular assignment since User isn't available when Account is created
+Account.users = User("users")
+
+
+class Query(GQLObject):
+    users = User("users")
+    accounts = Account("accounts")
+
+
+class Mutation(GQLObject):
+    createUser = "createUser"
+    createAccount = "createAccount"
+
+
+def verify(query, expected):
+    assert parse_graphql(query) == dedent(expected).strip()
+
+
+def test_default_gqlo_name_is_lowercase():
+    assert str(Account()) == "account"
+
+
+def test_parse_graphql_dedents_and_strips():
+    query = """
+
+
+
+        hi
+            there
+
+
+    """
+    assert parse_graphql(query) == "hi\n    there"
+
+
+def test_string_query_1():
+    verify(
+        query={"query": {"users": ["id", "name"]}},
+        expected="""
+            query {
+                users {
+                    id
+                    name
+                }
+            }
+        """,
+    )
+
+
+def test_string_query_2():
+    verify(
+        query={"query": {"users": [{"id(arg1: 1, arg2: 2)": ["result"]}, "name"]}},
+        expected="""
+            query {
+                users {
+                    id(arg1: 1, arg2: 2) {
+                        result
+                    }
+                    name
+                }
+            }
+        """,
+    )
+
+
+def test_string_query_3():
+
+    # do this to ensure field order on Python < 3.6
+    inner = OrderedDict()
+    inner["users"] = ["id", "name"]
+    inner["accounts"] = ["id", "name"]
+
+    verify(
+        query={"query": inner},
+        expected="""
+            query {
+                users {
+                    id
+                    name
+                }
+                accounts {
+                    id
+                    name
+                }
+            }
+        """,
+    )
+
+
+def test_gqlo_1():
+    verify(
+        query={"query": {Query.accounts: [Account.id, Account.name]}},
+        expected="""
+            query {
+                accounts {
+                    id
+                    name
+                }
+            }
+            """,
+    )
+
+
+def test_gqlo_2():
+
+    # do this to ensure field order on Python < 3.6
+    inner = OrderedDict()
+    inner[Query.accounts] = [Account.id, Account.name]
+    inner[Query.users] = [User.id, User.name]
+
+    verify(
+        query={"query": inner},
+        expected="""
+            query {
+                accounts {
+                    id
+                    name
+                }
+                users {
+                    id
+                    name
+                }
+            }
+            """,
+    )
+
+
+def test_gqlo_is_callable_for_arguments():
+    verify(
+        query={"query": {Query.accounts("where: {id: 5}"): [Account.id, Account.name]}},
+        expected="""
+            query {
+                accounts(where: {id: 5}) {
+                    id
+                    name
+                }
+            }
+            """,
+    )
+
+
+def test_nested_gqlo():
+    verify(
+        query={
+            "query": {
+                Query.accounts: {
+                    Query.accounts.users: {
+                        Query.accounts.users.account: Query.accounts.users.account.id
+                    }
+                }
+            }
+        },
+        expected="""
+            query {
+                accounts {
+                    users {
+                        account {
+                            id
+                        }
+                    }
+                }
+            }
+        """,
+    )