Merge pull request kensho-technologies#166 from kensho-technologies/add-sql-filters

Adding support for Filter blocks to SQL backend

Author: jmeulemans

graphql_compiler/compiler/emit_sql.py

@@ -2,9 +2,13 @@
 """Transform a SqlNode tree into an executable SQLAlchemy query."""
 from collections import namedtuple
 
-from sqlalchemy import select
+from sqlalchemy import Column, bindparam, select
+from sqlalchemy.sql import expression as sql_expressions
+from sqlalchemy.sql.elements import BindParameter, and_
 
 from . import sql_context_helpers
+from ..compiler import expressions
+from ..compiler.ir_lowering_sql import constants
 
 
 # The compilation context holds state that changes during compilation as the tree is traversed
@@ -20,6 +24,12 @@
     # renamed status. This tuple is used to construct the query outputs, and track when a name
     # changes due to collapsing into a CTE.
     'query_path_to_output_fields',
+    # 'query_path_to_filters': Dict[Tuple[str, ...], List[Filter]], mapping from each query_path
+    # to the Filter blocks that apply to that query path
+    'query_path_to_filters',
+    # 'query_path_to_node': Dict[Tuple[str, ...], SqlNode], mapping from each
+    # query_path to the SqlNode located at that query_path.
+    'query_path_to_node',
     # 'compiler_metadata': CompilerMetadata, SQLAlchemy metadata about Table objects, and
     # further backend specific configuration.
     'compiler_metadata',
@@ -40,6 +50,8 @@ def emit_code_from_ir(sql_query_tree, compiler_metadata):
         query_path_to_selectable=dict(),
         query_path_to_location_info=sql_query_tree.query_path_to_location_info,
         query_path_to_output_fields=sql_query_tree.query_path_to_output_fields,
+        query_path_to_filters=sql_query_tree.query_path_to_filters,
+        query_path_to_node=sql_query_tree.query_path_to_node,
         compiler_metadata=compiler_metadata,
     )
 
@@ -88,27 +100,187 @@ def _create_query(node, context):
     Returns:
         Selectable, selectable of the generated query.
     """
-    output_columns = _get_output_columns(node, context)
+    visited_nodes = [node]
+    output_columns = _get_output_columns(visited_nodes, context)
+    filters = _get_filters(visited_nodes, context)
     selectable = sql_context_helpers.get_node_selectable(node, context)
-    query = select(output_columns).select_from(selectable)
+    query = select(output_columns).select_from(selectable).where(and_(*filters))
     return query
 
 
-def _get_output_columns(node, context):
-    """Get the output columns required by the query.
+def _get_output_columns(nodes, context):
+    """Get the output columns for a list of SqlNodes.
 
     Args:
-        node: SqlNode, the current node.
+        nodes: List[SqlNode], the nodes to get output columns from.
         context: CompilationContext, global compilation state and metadata.
 
     Returns:
         List[Column], list of SqlAlchemy Columns to output for this query.
     """
-    sql_outputs = context.query_path_to_output_fields[node.query_path]
     columns = []
-    for sql_output in sql_outputs:
-        field_name = sql_output.field_name
-        column = sql_context_helpers.get_column(field_name, node, context)
-        column = column.label(sql_output.output_name)
-        columns.append(column)
+    for node in nodes:
+        for sql_output in sql_context_helpers.get_outputs(node, context):
+            field_name = sql_output.field_name
+            column = sql_context_helpers.get_column(field_name, node, context)
+            column = column.label(sql_output.output_name)
+            columns.append(column)
     return columns
+
+
+def _get_filters(nodes, context):
+    """Get filters to apply to a list of SqlNodes.
+
+    Args:
+        nodes: List[SqlNode], the SqlNodes to get filters for.
+        context: CompilationContext, global compilation state and metadata.
+
+    Returns:
+        List[Expression], list of SQLAlchemy expressions.
+    """
+    filters = []
+    for node in nodes:
+        for filter_block in sql_context_helpers.get_filters(node, context):
+            filter_sql_expression = _transform_filter_to_sql(filter_block, node, context)
+            filters.append(filter_sql_expression)
+    return filters
+
+
+def _transform_filter_to_sql(filter_block, node, context):
+    """Transform a Filter block to its corresponding SQLAlchemy expression.
+
+    Args:
+        filter_block: Filter, the Filter block to transform.
+        node: SqlNode, the node Filter block applies to.
+        context: CompilationContext, global compilation state and metadata.
+
+    Returns:
+        Expression, SQLAlchemy expression equivalent to the Filter.predicate expression.
+    """
+    expression = filter_block.predicate
+    return _expression_to_sql(expression, node, context)
+
+
+def _expression_to_sql(expression, node, context):
+    """Recursively transform a Filter block predicate to its SQLAlchemy expression representation.
+
+    Args:
+        expression: expression, the compiler expression to transform.
+        node: SqlNode, the SqlNode the expression applies to.
+        context: CompilationContext, global compilation state and metadata.
+
+    Returns:
+        Expression, SQLAlchemy Expression equivalent to the passed compiler expression.
+    """
+    _expression_transformers = {
+        expressions.LocalField: _transform_local_field_to_expression,
+        expressions.Variable: _transform_variable_to_expression,
+        expressions.Literal: _transform_literal_to_expression,
+        expressions.BinaryComposition: _transform_binary_composition_to_expression,
+    }
+    expression_type = type(expression)
+    if expression_type not in _expression_transformers:
+        raise NotImplementedError(
+            u'Unsupported compiler expression "{}" of type "{}" cannot be converted to SQL '
+            u'expression.'.format(expression, type(expression)))
+    return _expression_transformers[expression_type](expression, node, context)
+
+
+def _transform_binary_composition_to_expression(expression, node, context):
+    """Transform a BinaryComposition compiler expression into a SQLAlchemy expression.
+
+    Recursively calls _expression_to_sql to convert its left and right sub-expressions.
+
+    Args:
+        expression: expression, BinaryComposition compiler expression.
+        node: SqlNode, the SqlNode the expression applies to.
+        context: CompilationContext, global compilation state and metadata.
+
+    Returns:
+        Expression, SQLAlchemy expression.
+    """
+    if expression.operator not in constants.SUPPORTED_OPERATORS:
+        raise NotImplementedError(
+            u'Filter operation "{}" is not supported by the SQL backend.'.format(
+                expression.operator))
+    sql_operator = constants.SUPPORTED_OPERATORS[expression.operator]
+    left = _expression_to_sql(expression.left, node, context)
+    right = _expression_to_sql(expression.right, node, context)
+    if sql_operator.cardinality == constants.CARDINALITY_UNARY:
+        left, right = _get_column_and_bindparam(left, right, sql_operator)
+        clause = getattr(left, sql_operator.name)(right)
+        return clause
+    elif sql_operator.cardinality == constants.CARDINALITY_BINARY:
+        clause = getattr(sql_expressions, sql_operator.name)(left, right)
+        return clause
+    elif sql_operator.cardinality == constants.CARDINALITY_LIST_VALUED:
+        left, right = _get_column_and_bindparam(left, right, sql_operator)
+        # ensure that SQLAlchemy treats the right bind parameter as list valued
+        right.expanding = True
+        clause = getattr(left, sql_operator.name)(right)
+        return clause
+    raise AssertionError(u'Unreachable, operator cardinality {} for compiler expression {} is '
+                         u'unknown'.format(sql_operator.cardinality, expression))
+
+
+def _get_column_and_bindparam(left, right, operator):
+    """Return left and right expressions in (Column, BindParameter) order."""
+    if not isinstance(left, Column):
+        left, right = right, left
+    if not isinstance(left, Column):
+        raise AssertionError(
+            u'SQLAlchemy operator {} expects Column as left side the of expression, got {} '
+            u'of type {} instead.'.format(operator, left, type(left)))
+    if not isinstance(right, BindParameter):
+        raise AssertionError(
+            u'SQLAlchemy operator {} expects BindParameter as the right side of the expression, '
+            u'got {} of type {} instead.'.format(operator, right, type(right)))
+    return left, right
+
+
+def _transform_literal_to_expression(expression, node, context):
+    """Transform a Literal compiler expression into its SQLAlchemy expression representation.
+
+    Args:
+        expression: expression, Literal compiler expression.
+        node: SqlNode, the SqlNode the expression applies to.
+        context: CompilationContext, global compilation state and metadata.
+
+    Returns:
+        Expression, SQLAlchemy expression.
+    """
+    return expression.value
+
+
+def _transform_variable_to_expression(expression, node, context):
+    """Transform a Variable compiler expression into its SQLAlchemy expression representation.
+
+    Args:
+        expression: expression, Variable compiler expression.
+        node: SqlNode, the SqlNode the expression applies to.
+        context: CompilationContext, global compilation state and metadata.
+
+    Returns:
+        Expression, SQLAlchemy expression.
+    """
+    variable_name = expression.variable_name
+    if not variable_name.startswith(u'$'):
+        raise AssertionError(u'Unexpectedly received variable name {} that is not '
+                             u'prefixed with "$"'.format(variable_name))
+    return bindparam(variable_name[1:])
+
+
+def _transform_local_field_to_expression(expression, node, context):
+    """Transform a LocalField compiler expression into its SQLAlchemy expression representation.
+
+    Args:
+        expression: expression, LocalField compiler expression.
+        node: SqlNode, the SqlNode the expression applies to.
+        context: CompilationContext, global compilation state and metadata.
+
+    Returns:
+        Expression, SQLAlchemy expression.
+    """
+    column_name = expression.field_name
+    column = sql_context_helpers.get_column(column_name, node, context)
+    return column

graphql_compiler/compiler/ir_lowering_sql/__init__.py

@@ -4,6 +4,7 @@
 
 from .sql_tree import SqlNode, SqlQueryTree
 from .. import blocks
+from ...compiler import expressions
 from ...compiler.helpers import Location
 from ..ir_lowering_sql import constants
 from ..metadata import LocationInfo
@@ -45,10 +46,14 @@ def lower_ir(ir_blocks, query_metadata_table, type_equivalence_hints=None):
     query_path_to_location_info = _map_query_path_to_location_info(query_metadata_table)
     query_path_to_output_fields = _map_query_path_to_outputs(
         construct_result, query_path_to_location_info)
+    block_index_to_location = _map_block_index_to_location(ir_blocks)
+
+    # perform lowering steps
+    ir_blocks = lower_unary_transformations(ir_blocks)
 
     # iteratively construct SqlTree
     query_path_to_node = {}
-    block_index_to_location = _map_block_index_to_location(ir_blocks)
+    query_path_to_filters = {}
     tree_root = None
     for index, block in enumerate(ir_blocks):
         if isinstance(block, constants.SKIPPABLE_BLOCK_TYPES):
@@ -64,12 +69,15 @@ def lower_ir(ir_blocks, query_metadata_table, type_equivalence_hints=None):
                         block, tree_root, ir_blocks, query_metadata_table))
             tree_root = SqlNode(block=block, query_path=query_path)
             query_path_to_node[query_path] = tree_root
+        elif isinstance(block, blocks.Filter):
+            query_path_to_filters.setdefault(query_path, []).append(block)
         else:
             raise AssertionError(
                 u'Unsupported block {} unexpectedly passed validation for IR blocks '
                 u'{} with query metadata table {} .'.format(block, ir_blocks, query_metadata_table))
 
-    return SqlQueryTree(tree_root, query_path_to_location_info, query_path_to_output_fields)
+    return SqlQueryTree(tree_root, query_path_to_location_info, query_path_to_output_fields,
+                        query_path_to_filters, query_path_to_node)
 
 
 def _validate_all_blocks_supported(ir_blocks, query_metadata_table):
@@ -223,3 +231,21 @@ def _map_block_index_to_location(ir_blocks):
                 block_index_to_location[ix] = ir_block.location
             current_block_ixs = []
     return block_index_to_location
+
+
+def lower_unary_transformations(ir_blocks):
+    """Raise exception if any unary transformation block encountered."""
+    def visitor_fn(expression):
+        """Raise error if current expression is a UnaryTransformation."""
+        if not isinstance(expression, expressions.UnaryTransformation):
+            return expression
+        raise NotImplementedError(
+            u'UnaryTransformation expression "{}" encountered with IR blocks {} is unsupported by '
+            u'the SQL backend.'.format(expression, ir_blocks)
+        )
+
+    new_ir_blocks = [
+        block.visit_and_update_expressions(visitor_fn)
+        for block in ir_blocks
+    ]
+    return new_ir_blocks

graphql_compiler/compiler/ir_lowering_sql/constants.py

@@ -20,13 +20,42 @@
 
 SUPPORTED_BLOCK_TYPES = (
     blocks.QueryRoot,
+    blocks.Filter,
 )
 
 SUPPORTED_OUTPUT_EXPRESSION_TYPES = (
     expressions.OutputContextField,
 )
 
 
+Operator = namedtuple('Operator', ('name', 'cardinality'))
+
+
+CARDINALITY_UNARY = 'UNARY'
+CARDINALITY_BINARY = 'BINARY'
+CARDINALITY_LIST_VALUED = 'LIST_VALUED'
+
+
+# The mapping supplied for SUPPORTED_OPERATORS allows for programmatic resolution of expressions
+# to their SQLAlchemy equivalents. As a concrete example, when converting the GraphQL filter
+# column_name @filter(op_name: "=", value: ["$variable_name"])
+# the corresponding python uses SQLAlchemy operation `__eq__` in a call like
+# getattr(Column(column_name), '__eq__')(BindParameter('variable_name')
+# which programattically generates the equivalent of the desired SQLAlchemy expression
+# Column('column_name') == BindParameter('variable_name')
+SUPPORTED_OPERATORS = {
+    u'contains': Operator(u'in_', CARDINALITY_LIST_VALUED),
+    u'&&': Operator(u'and_', CARDINALITY_BINARY),
+    u'||': Operator(u'or_', CARDINALITY_BINARY),
+    u'=': Operator(u'__eq__', CARDINALITY_UNARY),
+    u'<': Operator(u'__lt__', CARDINALITY_UNARY),
+    u'>': Operator(u'__gt__', CARDINALITY_UNARY),
+    u'<=': Operator(u'__le__', CARDINALITY_UNARY),
+    u'>=': Operator(u'__ge__', CARDINALITY_UNARY),
+    u'has_substring': Operator(u'contains', CARDINALITY_UNARY),
+}
+
+
 class SqlBackend(object):
 
     def __init__(self, backend):

graphql_compiler/compiler/ir_lowering_sql/sql_tree.py

@@ -2,11 +2,14 @@
 
 
 class SqlQueryTree(object):
-    def __init__(self, root, query_path_to_location_info, query_path_to_output_fields):
+    def __init__(self, root, query_path_to_location_info,
+                 query_path_to_output_fields, query_path_to_filters, query_path_to_node):
         """Wrap a SqlNode root with additional location_info metadata."""
         self.root = root
         self.query_path_to_location_info = query_path_to_location_info
         self.query_path_to_output_fields = query_path_to_output_fields
+        self.query_path_to_filters = query_path_to_filters
+        self.query_path_to_node = query_path_to_node
 
 
 class SqlNode(object):