Merge branch 'master' into volo/devexp/cleanup

snuba/datasets/plans/entity_validation.py

@@ -0,0 +1,88 @@
+from __future__ import annotations
+
+from typing import Union
+
+import sentry_sdk
+
+from snuba.datasets.entities.factory import get_entity
+from snuba.query import Query
+from snuba.query.composite import CompositeQuery
+from snuba.query.data_source.join import JoinClause
+from snuba.query.data_source.simple import Entity as QueryEntity
+from snuba.query.exceptions import InvalidQueryException, ValidationException
+from snuba.query.logical import Query as LogicalQuery
+from snuba.query.parser.validation.functions import FunctionCallsValidator
+from snuba.query.query_settings import QuerySettings
+from snuba.state import explain_meta
+
+EXPRESSION_VALIDATORS = [FunctionCallsValidator()]
+
+
+def _validate_query(query: Query) -> None:
+ """
+ Applies all the expression validators in one pass over the AST.
+ """
+
+ for exp in query.get_all_expressions():
+ for v in EXPRESSION_VALIDATORS:
+ v.validate(exp, query.get_from_clause())
+
+
+def _validate_entities_with_query(
+ query: Union[CompositeQuery[QueryEntity], LogicalQuery]
+) -> None:
+ """
+ Applies all validator defined on the entities in the query
+ """
+ if isinstance(query, LogicalQuery):
+ entity = get_entity(query.get_from_clause().key)
+ try:
+ for v in entity.get_validators():
+ v.validate(query)
+ except InvalidQueryException as e:
+ raise ValidationException(
+ f"validation failed for entity {query.get_from_clause().key.value}: {e}",
+ should_report=e.should_report,
+ )
+ else:
+ from_clause = query.get_from_clause()
+ if isinstance(from_clause, JoinClause):
+ alias_map = from_clause.get_alias_node_map()
+ for alias, node in alias_map.items():
+ assert isinstance(node.data_source, QueryEntity) # mypy
+ entity = get_entity(node.data_source.key)
+ try:
+ for v in entity.get_validators():
+ v.validate(query, alias)
+ except InvalidQueryException as e:
+ raise ValidationException(
+ f"validation failed for entity {node.data_source.key.value}: {e}",
+ should_report=e.should_report,
+ )
+
+
+VALIDATORS = [_validate_query, _validate_entities_with_query]
+
+
+def run_entity_validators(
+ query: Union[CompositeQuery[QueryEntity], LogicalQuery],
+ settings: QuerySettings | None = None,
+) -> None:
+ """
+ Main function for applying all validators associated with an entity
+ """
+ for validator_func in VALIDATORS:
+ description = getattr(validator_func, "__name__", "custom")
+ with sentry_sdk.start_span(op="validator", description=description):
+ if settings and settings.get_dry_run():
+ with explain_meta.with_query_differ(
+ "entity_validator", description, query
+ ):
+ validator_func(query)
+ else:
+ validator_func(query)
+
+ if isinstance(query, CompositeQuery):
+ from_clause = query.get_from_clause()
+ if isinstance(from_clause, (LogicalQuery, CompositeQuery)):
+ run_entity_validators(from_clause, settings)

snuba/pipeline/stages/query_processing.py

@@ -1,5 +1,6 @@
 from snuba.clickhouse.query import Query as ClickhouseQuery
 from snuba.datasets.plans.entity_processing import run_entity_processing_executor
+from snuba.datasets.plans.entity_validation import run_entity_validators
 from snuba.datasets.plans.storage_processing import (
  apply_storage_processors,
  build_best_plan,
@@ -22,6 +23,10 @@ class EntityProcessingStage(
  def _process_data(
  self, pipe_input: QueryPipelineData[Request]
  ) -> ClickhouseQuery | CompositeQuery[Table]:
+ # Execute entity validators on logical query
+ run_entity_validators(pipe_input.data.query, pipe_input.query_settings)
+
+ # Run entity processors on the query and transform logical query into a physical query
  if isinstance(pipe_input.data.query, LogicalQuery):
  return run_entity_processing_executor(
  pipe_input.data.query, pipe_input.query_settings

snuba/query/parser/validation/__init__.py

@@ -1,36 +0,0 @@
-from abc import ABC, abstractmethod
-from typing import Sequence
-
-from snuba.query import Query
-from snuba.query.data_source import DataSource
-from snuba.query.expressions import Expression
-
-
-class ExpressionValidator(ABC):
- """
- Validates an individual expression in a Snuba logical query.
- """
-
- @abstractmethod
- def validate(self, exp: Expression, data_source: DataSource) -> None:
- """
- If the expression is valid according to this validator it
- returns, otherwise it raises a subclass of
- snuba.query.parser.ValidationException
- """
- raise NotImplementedError
-
-
-from snuba.query.parser.validation.functions import FunctionCallsValidator
-
-validators: Sequence[ExpressionValidator] = [FunctionCallsValidator()]
-
-
-def validate_query(query: Query) -> None:
- """
- Applies all the expression validators in one pass over the AST.
- """
-
- for exp in query.get_all_expressions():
- for v in validators:
- v.validate(exp, query.get_from_clause())

snuba/query/parser/validation/functions.py

@@ -1,3 +1,4 @@
+from abc import ABC, abstractmethod
 from typing import List, Mapping
 
 from snuba.clickhouse.columns import Array, DateTime, String
@@ -10,7 +11,6 @@
 from snuba.query.data_source.visitor import DataSourceVisitor
 from snuba.query.exceptions import InvalidExpressionException
 from snuba.query.expressions import Expression, FunctionCall
-from snuba.query.parser.validation import ExpressionValidator
 from snuba.query.validation import FunctionCallValidator, InvalidFunctionCall
 from snuba.query.validation.functions import AllowedFunctionValidator
 from snuba.query.validation.signature import Any, Column, SignatureValidator
@@ -90,6 +90,21 @@ def visit_join_clause(self, node: JoinClause[QueryEntity]) -> List[QueryEntity]:
  return node.right_node.accept(self) + node.left_node.accept(self)
 
 
+class ExpressionValidator(ABC):
+ """
+ Validates an individual expression in a Snuba logical query.
+ """
+
+ @abstractmethod
+ def validate(self, exp: Expression, data_source: DataSource) -> None:
+ """
+ If the expression is valid according to this validator it
+ returns, otherwise it raises a subclass of
+ snuba.query.parser.ValidationException
+ """
+ raise NotImplementedError
+
+
 class FunctionCallsValidator(ExpressionValidator):
  """
  Applies all function validators on the provided expression.

snuba/query/snql/parser.py

@@ -72,7 +72,6 @@
  validate_aliases,
 )
 from snuba.query.parser.exceptions import ParsingException, PostProcessingError
-from snuba.query.parser.validation import validate_query
 from snuba.query.query_settings import QuerySettings
 from snuba.query.schema import POSITIVE_OPERATORS
 from snuba.query.snql.discover_entity_selection import select_discover_entity
@@ -1490,8 +1489,6 @@ def _post_process(
 
 VALIDATORS = [
  validate_identifiers_in_lambda,
- validate_query,
- validate_entities_with_query,
 ]
 
 

tests/clickhouse/query_dsl/test_project_id.py

@@ -5,9 +5,11 @@
 
 from snuba.clickhouse.query_dsl.accessors import get_object_ids_in_query_ast
 from snuba.datasets.factory import get_dataset
+from snuba.datasets.plans.entity_validation import run_entity_validators
 from snuba.datasets.plans.translator.query import identity_translate
+from snuba.query.exceptions import ValidationException
 from snuba.query.logical import Query
-from snuba.query.parser.exceptions import ParsingException
+from snuba.query.query_settings import HTTPQuerySettings
 from snuba.query.snql.parser import parse_snql_query
 
 test_cases: Sequence[Tuple[Mapping[str, Any], Optional[Set[int]]]] = [
@@ -176,17 +178,19 @@ def test_find_projects(
 ) -> None:
  events = get_dataset("events")
  if expected_projects is None:
- with pytest.raises(ParsingException):
+ with pytest.raises(ValidationException):
  request = json_to_snql(query_body, "events")
  request.validate()
  query = parse_snql_query(str(request.query), events)
  assert isinstance(query, Query)
+ run_entity_validators(query, HTTPQuerySettings())
  identity_translate(query)
  else:
  request = json_to_snql(query_body, "events")
  request.validate()
  query = parse_snql_query(str(request.query), events)
  assert isinstance(query, Query)
+ run_entity_validators(query, HTTPQuerySettings())
  translated_query = identity_translate(query)
  project_ids_ast = get_object_ids_in_query_ast(translated_query, "project_id")
  assert project_ids_ast == expected_projects

tests/pipeline/test_entity_processing_stage_composite.py

@@ -328,7 +328,7 @@
  ),
  SelectedExpression(
  "_snuba_right",
- Column("_snuba_right", "groups", "right_col"),
+ Column("_snuba_right", "groups", "status"),
  ),
  ],
  condition=binary_condition(
@@ -398,7 +398,7 @@
  ),
  SelectedExpression(
  "_snuba_right",
- Column("_snuba_right", None, "right_col"),
+ Column("_snuba_right", None, "status"),
  ),
  ],
  ),

tests/subscriptions/test_subscription.py

@@ -8,8 +8,7 @@
 from snuba.datasets.entities.factory import get_entity
 from snuba.datasets.entity_subscriptions.validators import InvalidSubscriptionError
 from snuba.datasets.factory import get_dataset
-from snuba.query.exceptions import InvalidQueryException
-from snuba.query.parser.exceptions import ParsingException
+from snuba.query.exceptions import InvalidQueryException, ValidationException
 from snuba.query.validation.validators import ColumnValidationMode
 from snuba.redis import RedisClientKey, get_redis_client
 from snuba.subscriptions.data import SubscriptionData
@@ -86,7 +85,7 @@ def test(self, subscription: SubscriptionData) -> None:
  def test_invalid_condition_column(self, subscription: SubscriptionData) -> None:
  override_entity_column_validator(EntityKey.EVENTS, ColumnValidationMode.ERROR)
  creator = SubscriptionCreator(self.dataset, EntityKey.EVENTS)
- with raises(ParsingException):
+ with raises(ValidationException):
  creator.create(
  subscription,
  self.timer,

tests/web/test_query_cache.py

@@ -10,8 +10,9 @@
 from snuba.datasets.entities.factory import get_entity
 from snuba.datasets.factory import get_dataset
 from snuba.query import SelectedExpression
+from snuba.query.conditions import in_condition
 from snuba.query.data_source.simple import Entity
-from snuba.query.expressions import Column
+from snuba.query.expressions import Column, Literal
 from snuba.query.logical import Query
 from snuba.query.query_settings import HTTPQuerySettings
 from snuba.request import Request
@@ -29,6 +30,7 @@ def run_query() -> None:
  selected_columns=[
  SelectedExpression("event_id", Column("_snuba_event_id", None, "event_id")),
  ],
+ condition=in_condition(Column(None, None, "project_id"), [Literal(None, 123)]),
  )
 
  query_settings = HTTPQuerySettings(referrer="asd")

tests/web/test_transform_names.py

@@ -12,6 +12,7 @@
 from snuba.datasets.factory import get_dataset
 from snuba.processor import InsertEvent
 from snuba.query import SelectedExpression
+from snuba.query.conditions import in_condition
 from snuba.query.data_source.simple import Entity
 from snuba.query.expressions import Column, FunctionCall, Literal
 from snuba.query.logical import Query
@@ -75,6 +76,7 @@ def test_transform_column_names() -> None:
  ),
  ),
  ],
+ condition=in_condition(Column(None, None, "project_id"), [Literal(None, 1)]),
  )
  query_settings = HTTPQuerySettings(referrer="asd")