def test_entity_data_model() -> None:
entity_data_model = EntityColumnSet(columns=[
Column("event_id", FixedString(32)),
Column("project_id", UInt(64)),
WildcardColumn("tags", String()),
WildcardColumn("contexts", String()),
])
event_id_col = entity_data_model.get("event_id")
assert event_id_col is not None
assert event_id_col.name == "event_id"
assert event_id_col.type == FixedString(32)
assert entity_data_model.get("tags[asdf]") == FlattenedColumn(
None, "tags[asdf]", String())
assert entity_data_model.get("asdf") is None
assert entity_data_model.get("tags[asd f]") is None
assert entity_data_model.get("asdf[gkrurrtsjhfkjgh]") is None
assert entity_data_model == EntityColumnSet(columns=[
Column("event_id", FixedString(32)),
Column("project_id", UInt(64)),
WildcardColumn("tags", String()),
WildcardColumn("contexts", String()),
])
def test_apply_quota(
enabled: int,
referrer: str,
config_to_set: str,
expected_quota: Optional[ResourceQuota],
) -> None:
state.set_config(ENABLED_CONFIG, enabled)
state.set_config(config_to_set, 5)
query = Query(
QueryEntity(EntityKey.EVENTS, EntityColumnSet([])),
selected_columns=[
SelectedExpression("column2", Column(None, None, "column2"))
],
condition=binary_condition(
ConditionFunctions.EQ,
Column("_snuba_project_id", None, "project_id"),
Literal(None, 1),
),
)
settings = HTTPQuerySettings()
settings.referrer = referrer
ResourceQuotaProcessor("project_id").process_query(query, settings)
assert settings.get_resource_quota() == expected_quota
def __init__(
self,
*,
storages: Sequence[Storage],
query_pipeline_builder: QueryPipelineBuilder[ClickhouseQueryPlan],
abstract_column_set: ColumnSet,
join_relationships: Mapping[str, JoinRelationship],
writable_storage: Optional[WritableTableStorage],
validators: Optional[Sequence[QueryValidator]],
required_time_column: Optional[str],
validate_data_model: ColumnValidationMode = ColumnValidationMode.
DO_NOTHING,
) -> None:
self.__storages = storages
self.__query_pipeline_builder = query_pipeline_builder
self.__writable_storage = writable_storage
# Eventually, the EntityColumnSet should be passed in
# For now, just convert it so we have the right
# type from here on
self.__data_model = EntityColumnSet(abstract_column_set.columns)
self.__join_relationships = join_relationships
self.required_time_column = required_time_column
columns_exist_validator = EntityContainsColumnsValidator(
self.__data_model, validation_mode=validate_data_model)
self.__validators = ([
*validators, columns_exist_validator
] if validators is not None else [columns_exist_validator])
def test_project_rate_limit_processor(unprocessed: Expression, project_id: int) -> None:
query = Query(
QueryEntity(EntityKey.EVENTS, EntityColumnSet([])),
selected_columns=[SelectedExpression("column2", Column(None, None, "column2"))],
condition=unprocessed,
)
settings = HTTPQuerySettings()
num_before = len(settings.get_rate_limit_params())
ProjectRateLimiterProcessor("project_id").process_query(query, settings)
assert len(settings.get_rate_limit_params()) == num_before + 1
rate_limiter = settings.get_rate_limit_params()[-1]
assert rate_limiter.rate_limit_name == PROJECT_RATE_LIMIT_NAME
assert rate_limiter.bucket == str(project_id)
assert rate_limiter.per_second_limit == 1000
assert rate_limiter.concurrent_limit == 1000
def test_org_rate_limit_processor_overridden(unprocessed: Expression,
org_id: int) -> None:
query = Query(
QueryEntity(EntityKey.EVENTS, EntityColumnSet([])),
selected_columns=[
SelectedExpression("column2", Column(None, None, "column2"))
],
condition=unprocessed,
)
settings = HTTPQuerySettings()
state.set_config(f"org_per_second_limit_{org_id}", 5)
state.set_config(f"org_concurrent_limit_{org_id}", 10)
num_before = len(settings.get_rate_limit_params())
OrganizationRateLimiterProcessor("org_id").process_query(query, settings)
assert len(settings.get_rate_limit_params()) == num_before + 1
rate_limiter = settings.get_rate_limit_params()[-1]
assert rate_limiter.rate_limit_name == ORGANIZATION_RATE_LIMIT_NAME
assert rate_limiter.bucket == str(org_id)
assert rate_limiter.per_second_limit == 5
assert rate_limiter.concurrent_limit == 10
def test_selector_function(
time_condition: Expression,
beginning_of_time: Optional[datetime],
exec_both: float,
trust_secondary: float,
expected_value: Tuple[str, List[str]],
) -> None:
query = Query(
QueryEntity(EntityKey.EVENTS, EntityColumnSet([])),
selected_columns=[
SelectedExpression("column2", Column(None, None, "column2"))
],
condition=time_condition,
)
set_config("rollout_upgraded_errors_trust", trust_secondary)
set_config("rollout_upgraded_errors_execute", exec_both)
previous_time = settings.ERRORS_UPGRADE_BEGINING_OF_TIME
settings.ERRORS_UPGRADE_BEGINING_OF_TIME = beginning_of_time
assert v2_selector_function(query, "test") == expected_value
settings.ERRORS_UPGRADE_BEGINING_OF_TIME = previous_time
OrganizationRateLimiterProcessor,
ReferrerRateLimiterProcessor,
)
from snuba.query.processors.timeseries_processor import TimeSeriesProcessor
from snuba.query.validation.validators import (
ColumnValidationMode,
EntityRequiredColumnValidator,
)
from snuba.utils.schemas import Column
outcomes_data_model = EntityColumnSet([
Column("org_id", UInt(64)),
Column("project_id", UInt(64)),
Column("key_id", UInt(64)),
Column("timestamp", DateTime()),
Column("outcome", UInt(8)),
Column("reason", String()),
Column("quantity", UInt(64)),
Column("category", UInt(8)),
Column("times_seen", UInt(64)),
Column("time", DateTime()),
])
class OutcomesEntity(Entity):
"""
Tracks event ingestion outcomes in Sentry.
"""
def __init__(self) -> None:
# The raw table we write onto, and that potentially we could
# query.
)
from snuba.query.processors.quota_processor import ResourceQuotaProcessor
from snuba.query.validation.validators import EntityRequiredColumnValidator
profile_columns = EntityColumnSet([
Column("organization_id", UInt(64)),
Column("project_id", UInt(64)),
Column("transaction_id", UUID()),
Column("profile_id", UUID()),
Column("received", DateTime()),
Column("profile", String()),
Column("android_api_level", UInt(32, Modifiers(nullable=True))),
Column("device_classification", String()),
Column("device_locale", String()),
Column("device_manufacturer", String()),
Column("device_model", String()),
Column("device_os_build_number", String(Modifiers(nullable=True))),
Column("device_os_name", String()),
Column("device_os_version", String()),
Column("duration_ns", UInt(64)),
Column("environment", String(Modifiers(nullable=True))),
Column("platform", String()),
Column("trace_id", UUID()),
Column("transaction_name", String()),
Column("version_name", String()),
Column("version_code", String()),
])
class ProfilesEntity(Entity, ABC):
def __init__(self, ) -> None:
from snuba.clickhouse.columns import UUID, UInt
from snuba.datasets.entities import EntityKey
from snuba.datasets.entities.entity_data_model import EntityColumnSet
from snuba.query import SelectedExpression
from snuba.query.composite import CompositeQuery
from snuba.query.conditions import ConditionFunctions, binary_condition
from snuba.query.data_source.simple import Entity
from snuba.query.expressions import Column, FunctionCall, Literal
from snuba.query.logical import Query
from snuba.utils.schemas import Column as EntityColumn
from snuba.web.query import ProjectsFinder
EVENTS_SCHEMA = EntityColumnSet([
EntityColumn("event_id", UUID()),
EntityColumn("project_id", UInt(32)),
EntityColumn("group_id", UInt(32)),
])
SIMPLE_QUERY = Query(
Entity(EntityKey.EVENTS, EVENTS_SCHEMA),
selected_columns=[
SelectedExpression(
"alias",
Column("_snuba_project", None, "project_id"),
)
],
array_join=None,
condition=binary_condition(
ConditionFunctions.IN,
Column("_snuba_project", None, "project_id"),
import pytest
from snuba import settings
from snuba.datasets.entities import EntityKey
from snuba.datasets.entities.entity_data_model import EntityColumnSet
from snuba.datasets.entities.events import v2_selector_function
from snuba.query import SelectedExpression
from snuba.query.conditions import ConditionFunctions, binary_condition
from snuba.query.data_source.simple import Entity as QueryEntity
from snuba.query.expressions import Column, Expression, Literal
from snuba.query.logical import Query
from snuba.state import set_config
query = Query(
QueryEntity(EntityKey.EVENTS, EntityColumnSet([])),
selected_columns=[
SelectedExpression("column2", Column(None, None, "column2"))
],
condition=binary_condition(
ConditionFunctions.EQ,
Column("_snuba_project_id", None, "project_id"),
Literal(None, 1),
),
)
TESTS = [
pytest.param(
binary_condition(
ConditionFunctions.GTE,
Column(None, None, "timestamp"),