Repository: kensho-technologies/graphql-compiler (language: python)

Issue:
[Query Planning Refactor]: SimpleExecute
This is part of a larger refactor to streamline the way we generate query plans. This is the simplest form of query plan where the given query is executed against the database.

Future work includes adding paginated queries and cross database queries.

There are some outstanding TODOs mainly around how to best utilize provider_id to look up both the corresponding schema and corresponding database query execution function.


Discussion / hints:

Fix patch:
diff --git a/graphql_compiler/query_planning/__init__.py b/graphql_compiler/query_planning/__init__.py
new file mode 100644
index 000000000..267b30955
--- /dev/null
+++ b/graphql_compiler/query_planning/__init__.py
@@ -0,0 +1,1 @@
+# Copyright 2021-present Kensho Technologies, LLC.
diff --git a/graphql_compiler/schema_transformation/make_query_plan.py b/graphql_compiler/query_planning/make_query_plan.py
similarity index 99%
rename from graphql_compiler/schema_transformation/make_query_plan.py
rename to graphql_compiler/query_planning/make_query_plan.py
index 6547b08f2..5de9128fe 100644
--- a/graphql_compiler/schema_transformation/make_query_plan.py
+++ b/graphql_compiler/query_planning/make_query_plan.py
@@ -21,7 +21,7 @@
 from ..ast_manipulation import get_only_query_definition
 from ..exceptions import GraphQLValidationError
 from ..schema import FilterDirective, OutputDirective
-from .split_query import AstType, SubQueryNode
+from ..schema_transformation.split_query import AstType, SubQueryNode
 
 
 @dataclass
@@ -66,47 +66,6 @@ class QueryPlanDescriptor:
     output_join_descriptors: List[OutputJoinDescriptor]
 
 
-def make_query_plan(
-    root_sub_query_node: SubQueryNode, intermediate_output_names: FrozenSet[str]
-) -> QueryPlanDescriptor:
-    """Return a QueryPlanDescriptor, whose query ASTs have @filters added.
-
-    For each parent of parent and child SubQueryNodes, a new @filter directive will be added
-    in the child AST. It will be added on the field whose @output directive has the out_name
-    equal to the child's out name as specified in the QueryConnection. The newly added @filter
-    will be a 'in_collection' type filter, and the name of the local variable is guaranteed to
-    be the same as the out_name of the @output on the parent.
-
-    ASTs contained in the input node and its children nodes will not be modified.
-
-    Args:
-        root_sub_query_node: representing the base of a query split into pieces
-                             that we want to turn into a query plan.
-        intermediate_output_names: names of outputs to be removed at the end.
-
-    Returns:
-        QueryPlanDescriptor containing a tree of SubQueryPlans that wrap around each individual
-        query AST, the set of intermediate output names that are to be removed at the end, and
-        information on which outputs are to be connect to which in what manner.
-    """
-    output_join_descriptors: List[OutputJoinDescriptor] = []
-
-    root_sub_query_plan = SubQueryPlan(
-        query_ast=root_sub_query_node.query_ast,
-        schema_id=root_sub_query_node.schema_id,
-        parent_query_plan=None,
-        child_query_plans=[],
-    )
-
-    _make_query_plan_recursive(root_sub_query_node, root_sub_query_plan, output_join_descriptors)
-
-    return QueryPlanDescriptor(
-        root_sub_query_plan=root_sub_query_plan,
-        intermediate_output_names=intermediate_output_names,
-        output_join_descriptors=output_join_descriptors,
-    )
-
-
 def _make_query_plan_recursive(
     sub_query_node: SubQueryNode,
     sub_query_plan: SubQueryPlan,
@@ -291,6 +250,66 @@ def _get_in_collection_filter_directive(input_filter_name: str) -> DirectiveNode
     )
 
 
+def _get_plan_and_depth_in_dfs_order(query_plan: SubQueryPlan) -> List[Tuple[SubQueryPlan, int]]:
+    """Return a list of topologically sorted (query plan, depth) tuples."""
+
+    def _get_plan_and_depth_in_dfs_order_helper(query_plan, depth):
+        plan_and_depth_in_dfs_order = [(query_plan, depth)]
+        for child_query_plan in query_plan.child_query_plans:
+            plan_and_depth_in_dfs_order.extend(
+                _get_plan_and_depth_in_dfs_order_helper(child_query_plan, depth + 1)
+            )
+        return plan_and_depth_in_dfs_order
+
+    return _get_plan_and_depth_in_dfs_order_helper(query_plan, 0)
+
+
+######
+# Public API
+######
+
+
+def make_query_plan(
+    root_sub_query_node: SubQueryNode, intermediate_output_names: FrozenSet[str]
+) -> QueryPlanDescriptor:
+    """Return a QueryPlanDescriptor, whose query ASTs have @filters added.
+
+    For each parent of parent and child SubQueryNodes, a new @filter directive will be added
+    in the child AST. It will be added on the field whose @output directive has the out_name
+    equal to the child's out name as specified in the QueryConnection. The newly added @filter
+    will be a 'in_collection' type filter, and the name of the local variable is guaranteed to
+    be the same as the out_name of the @output on the parent.
+
+    ASTs contained in the input node and its children nodes will not be modified.
+
+    Args:
+        root_sub_query_node: representing the base of a query split into pieces
+                             that we want to turn into a query plan.
+        intermediate_output_names: names of outputs to be removed at the end.
+
+    Returns:
+        QueryPlanDescriptor containing a tree of SubQueryPlans that wrap around each individual
+        query AST, the set of intermediate output names that are to be removed at the end, and
+        information on which outputs are to be connect to which in what manner.
+    """
+    output_join_descriptors: List[OutputJoinDescriptor] = []
+
+    root_sub_query_plan = SubQueryPlan(
+        query_ast=root_sub_query_node.query_ast,
+        schema_id=root_sub_query_node.schema_id,
+        parent_query_plan=None,
+        child_query_plans=[],
+    )
+
+    _make_query_plan_recursive(root_sub_query_node, root_sub_query_plan, output_join_descriptors)
+
+    return QueryPlanDescriptor(
+        root_sub_query_plan=root_sub_query_plan,
+        intermediate_output_names=intermediate_output_names,
+        output_join_descriptors=output_join_descriptors,
+    )
+
+
 def print_query_plan(query_plan_descriptor: QueryPlanDescriptor, indentation_depth: int = 4) -> str:
     """Return a string describing query plan."""
     query_plan_strings = [""]
@@ -311,17 +330,3 @@ def print_query_plan(query_plan_descriptor: QueryPlanDescriptor, indentation_dep
     query_plan_strings.append(str(query_plan_descriptor.intermediate_output_names) + "\n")
 
     return "".join(query_plan_strings)
-
-
-def _get_plan_and_depth_in_dfs_order(query_plan: SubQueryPlan) -> List[Tuple[SubQueryPlan, int]]:
-    """Return a list of topologically sorted (query plan, depth) tuples."""
-
-    def _get_plan_and_depth_in_dfs_order_helper(query_plan, depth):
-        plan_and_depth_in_dfs_order = [(query_plan, depth)]
-        for child_query_plan in query_plan.child_query_plans:
-            plan_and_depth_in_dfs_order.extend(
-                _get_plan_and_depth_in_dfs_order_helper(child_query_plan, depth + 1)
-            )
-        return plan_and_depth_in_dfs_order
-
-    return _get_plan_and_depth_in_dfs_order_helper(query_plan, 0)
diff --git a/mypy.ini b/mypy.ini
index b93fb909b..6d34e87db 100644
--- a/mypy.ini
+++ b/mypy.ini
@@ -160,6 +160,10 @@ disallow_untyped_calls = False
 [mypy-graphql_compiler.query_pagination.query_parameterizer.*]
 disallow_untyped_calls = False
 
+[mypy-graphql_compiler.query_planning.*]
+disallow_untyped_calls = False
+disallow_untyped_defs = False
+
 [mypy-graphql_compiler.schema_generation.graphql_schema.*]
 check_untyped_defs = False
 disallow_incomplete_defs = False
@@ -189,9 +193,6 @@ check_untyped_defs = False
 [mypy-graphql_compiler.schema_transformation.*]
 disallow_untyped_calls = False
 
-[mypy-graphql_compiler.schema_transformation.make_query_plan.*]
-disallow_untyped_defs = False
-
 [mypy-graphql_compiler.schema_transformation.split_query.*]
 disallow_incomplete_defs = False
 disallow_untyped_defs = False


Test patch:
diff --git a/graphql_compiler/tests/schema_transformation_tests/test_make_query_plan.py b/graphql_compiler/tests/schema_transformation_tests/test_make_query_plan.py
index eb825fc90..058589dc4 100644
--- a/graphql_compiler/tests/schema_transformation_tests/test_make_query_plan.py
+++ b/graphql_compiler/tests/schema_transformation_tests/test_make_query_plan.py
@@ -4,7 +4,7 @@
 
 from graphql import parse, print_ast
 
-from ...schema_transformation.make_query_plan import make_query_plan
+from ...query_planning.make_query_plan import make_query_plan
 from ...schema_transformation.split_query import split_query
 from .example_schema import basic_merged_schema