def _get_object(self, object_id, eagerload=False, **kwargs):
"""
Given the object_id and extra route params, get an instance of
``self.model_class``
"""
confirmed = self._get_querystring_boolean_field('confirmed')
active = self._get_querystring_boolean_field('active')
status = flask.request.args.get('status')
extra_query = None
if status and status in Vulnerability.STATUSES:
extra_query = f"status='{status}'"
self._validate_object_id(object_id)
query = db.session.query(Workspace).filter_by(name=object_id)
if active is not None:
query = query.filter_by(active=active)
query = query.options(
with_expression(
Workspace.vulnerability_web_count,
_make_vuln_count_property('vulnerability_web',
confirmed=confirmed,
extra_query=extra_query,
use_column_property=False),
),
with_expression(
Workspace.vulnerability_standard_count,
_make_vuln_count_property('vulnerability',
confirmed=confirmed,
extra_query=extra_query,
use_column_property=False)),
with_expression(
Workspace.vulnerability_total_count,
_make_vuln_count_property(type_=None,
confirmed=confirmed,
extra_query=extra_query,
use_column_property=False)),
with_expression(
Workspace.vulnerability_code_count,
_make_vuln_count_property('vulnerability_code',
extra_query=extra_query,
use_column_property=False),
),
with_expression(
Workspace.active_agents_count,
_make_active_agents_count_property(),
),
)
query = count_vulnerability_severities(query,
Workspace,
status=status,
confirmed=confirmed,
all_severities=True)
try:
obj = query.one()
except NoResultFound:
flask.abort(404, f'Object with name "{object_id}" not found')
return obj
def _get_object(self, object_id, eagerload=False, **kwargs):
"""
Given the object_id and extra route params, get an instance of
``self.model_class``
"""
confirmed = flask.request.values.get('confirmed', None)
if confirmed:
try:
confirmed = bool(int(confirmed))
except ValueError:
if confirmed.lower() == 'false':
confirmed = False
elif confirmed.lower() == 'true':
confirmed = True
self._validate_object_id(object_id)
query = db.session.query(Workspace).filter_by(name=object_id)
query = query.options(
with_expression(
Workspace.vulnerability_web_count,
_make_vuln_count_property('vulnerability_web',
confirmed=confirmed,
use_column_property=False),
),
with_expression(
Workspace.vulnerability_standard_count,
_make_vuln_count_property('vulnerability',
confirmed=confirmed,
use_column_property=False)
),
with_expression(
Workspace.vulnerability_total_count,
_make_vuln_count_property(type_=None,
confirmed=confirmed,
use_column_property=False)
),
with_expression(
Workspace.vulnerability_code_count,
_make_vuln_count_property('vulnerability_code',
use_column_property=False)
),
)
try:
obj = query.one()
except NoResultFound:
flask.abort(404, 'Object with id "%s" not found' % object_id)
return obj
def test_query_expr(self):
(User, ) = self.classes("User")
self._run_cache_key_fixture(
lambda: (
with_expression(User.name, true()),
with_expression(User.name, null()),
with_expression(User.name, func.foobar()),
with_expression(User.name, User.name == "test"),
Load(User).with_expression(User.name, true()),
Load(User).with_expression(User.name, null()),
Load(User).with_expression(User.name, func.foobar()),
Load(User).with_expression(User.name, User.name == "test"),
),
compare_values=True,
)
def most_copies(cls, num=5):
count_expr = func.count(cls.copies)
query_string = cls.query.options(with_expression(cls.expr, count_expr))\
.group_by(Book.id).order_by(count_expr.desc()).limit(num)
return query_string.all()
def test_simple_expr(self):
A = self.classes.A
s = Session()
a1 = s.query(A).options(
with_expression(A.my_expr, A.x + A.y)).filter(A.x > 1).\
order_by(A.id)
eq_(a1.all(), [A(my_expr=5), A(my_expr=15), A(my_expr=12)])
def _inject_omics_data_summary(self, db: Session, query: Query) -> Query:
aggs = []
for omics_class in workflow_search_classes:
pipeline_model = omics_class().table.model
table_name = pipeline_model.__tablename__ # type: ignore
filter_conditions = [
c for c in self.conditions if c.table.value in
{"omics_processing", table_name, "biosample"}
]
query_schema = omics_class(conditions=filter_conditions)
omics_subquery = self._count_omics_data_query(
db, query_schema).subquery()
study_id = getattr(omics_subquery.c, f"{table_name}_study_id")
query = query.join(
omics_subquery,
self.table.model.id == study_id, # type: ignore
isouter=True,
)
aggs.append(
func.json_build_object(
"type", table_name, "count",
getattr(omics_subquery.c, f"{table_name}_count")))
op_filter_conditions = [
c for c in self.conditions
if c.table.value in {"omics_processing", "biosample"}
]
op_summary_subquery = self._count_omics_processing_summary(
db, op_filter_conditions).subquery()
query = query.join(
op_summary_subquery,
op_summary_subquery.c.omics_processing_study_id == models.Study.id,
isouter=True,
)
aggregation = func.json_build_array(*aggs)
return query.populate_existing().options(
with_expression(models.Study.omics_counts, aggregation),
with_expression(
models.Study.omics_processing_counts,
op_summary_subquery.c.omics_processing_summary,
),
)
def _get_object(self, object_id, eagerload=False, **kwargs):
"""
Given the object_id and extra route params, get an instance of
``self.model_class``
"""
confirmed = self._get_querystring_boolean_field('confirmed')
active = self._get_querystring_boolean_field('active')
self._validate_object_id(object_id)
query = db.session.query(Workspace).filter_by(name=object_id)
if active is not None:
query = query.filter_by(active=active)
query = query.options(
with_expression(
Workspace.vulnerability_web_count,
_make_vuln_count_property('vulnerability_web',
confirmed=confirmed,
use_column_property=False),
),
with_expression(
Workspace.vulnerability_standard_count,
_make_vuln_count_property('vulnerability',
confirmed=confirmed,
use_column_property=False)
),
with_expression(
Workspace.vulnerability_total_count,
_make_vuln_count_property(type_=None,
confirmed=confirmed,
use_column_property=False)
),
with_expression(
Workspace.vulnerability_code_count,
_make_vuln_count_property('vulnerability_code',
use_column_property=False)
),
)
try:
obj = query.one()
except NoResultFound:
flask.abort(404, 'Object with id "%s" not found' % object_id)
return obj
def test_dont_explode_on_expire_whole(self):
A = self.classes.A
s = Session()
q = s.query(A).options(
with_expression(A.my_expr, A.x + A.y)).filter(A.x > 1).\
order_by(A.id)
a1 = q.first()
eq_(a1.my_expr, 5)
s.expire(a1)
eq_(a1.my_expr, None)
# comes back
q = s.query(A).options(
with_expression(A.my_expr, A.x + A.y)).filter(A.x > 1).\
order_by(A.id)
q.first()
eq_(a1.my_expr, 5)
def test_dont_explode_on_expire_whole(self):
A = self.classes.A
s = Session()
q = s.query(A).options(
with_expression(A.my_expr, A.x + A.y)).filter(A.x > 1).\
order_by(A.id)
a1 = q.first()
eq_(a1.my_expr, 5)
s.expire(a1)
eq_(a1.my_expr, None)
# comes back
q = s.query(A).options(
with_expression(A.my_expr, A.x + A.y)).filter(A.x > 1).\
order_by(A.id)
q.first()
eq_(a1.my_expr, 5)
def _get_object(self, object_id, eagerload=False, **kwargs):
"""
Given the object_id and extra route params, get an instance of
``self.model_class``
"""
confirmed = self._get_querystring_boolean_field('confirmed')
active = self._get_querystring_boolean_field('active')
self._validate_object_id(object_id)
query = db.session.query(Workspace).filter_by(name=object_id)
if active is not None:
query = query.filter_by(active=active)
query = query.options(
with_expression(
Workspace.vulnerability_web_count,
_make_vuln_count_property('vulnerability_web',
confirmed=confirmed,
use_column_property=False),
),
with_expression(
Workspace.vulnerability_standard_count,
_make_vuln_count_property('vulnerability',
confirmed=confirmed,
use_column_property=False)),
with_expression(
Workspace.vulnerability_total_count,
_make_vuln_count_property(type_=None,
confirmed=confirmed,
use_column_property=False)),
with_expression(
Workspace.vulnerability_code_count,
_make_vuln_count_property('vulnerability_code',
use_column_property=False)),
)
try:
obj = query.one()
except NoResultFound:
flask.abort(404, 'Object with name "%s" not found' % object_id)
return obj
def filter_within_radius(self, lat, lng, radius):
"""
Filter user within radius from a center (lat, lng) coordinate
"""
# Define center point
point = 'POINT(%f %f)' % (lng, lat)
wkb_element = WKTElement(point, srid=4326)
# Define expression to calculate distance
# from center point to users location
if db.engine.name == 'sqlite':
distance = func \
.distance(User.location, wkb_element, 1) \
.label('distance')
else:
distance = User.location \
.distance_centroid(wkb_element) \
.cast(db.Float) \
.label('distance')
# Define lat, lng query set
lat = func.ST_Y(User.location).label('lat')
lng = func.ST_X(User.location).label('lng')
# Filter user within radius from center point
if db.engine.name == 'sqlite':
qs = User.query.filter(
func.PtDistWithin(User.location, wkb_element, radius))
else:
qs = User.query.filter(
func.ST_DWithin(User.location, wkb_element, radius))
# Append Query-time SQL expressions distance as mapped attributes
# https://docs.sqlalchemy.org/en/latest/orm/mapped_sql_expr.html
qs = qs.options(with_expression(User.distance, distance),
with_expression(User.lat, lat),
with_expression(User.lng, lng))
return qs
def test_reuse_expr(self):
A = self.classes.A
s = Session()
# so people will obv. want to say, "filter(A.my_expr > 10)".
# but that means Query or Core has to post-modify the statement
# after construction.
expr = A.x + A.y
a1 = s.query(A).options(
with_expression(A.my_expr, expr)).filter(expr > 10).\
order_by(expr)
eq_(a1.all(), [A(my_expr=12), A(my_expr=15)])
def test_simple_expr(self):
A = self.classes.A
s = Session()
a1 = s.query(A).options(
with_expression(A.my_expr, A.x + A.y)).filter(A.x > 1).\
order_by(A.id)
eq_(
a1.all(),
[
A(my_expr=5), A(my_expr=15), A(my_expr=12)
]
)
def query_next_requests(self, max_n_requests, partitions):
partitions_count = len(self.partitioner.partitions)
score_window = partitions_count * self.score_window
partition = DynamicQueueModel.partition_seed % partitions_count
netloc = text("(regexp_matches(url, '^(?:://)?([^/]+)'))[1]")
score_query = self.session.query(
self.queue_model,
func.row_number().over(
order_by=[
self.queue_model.score.desc(),
self.queue_model.crawl_at
]).label('score_rank')
).\
filter(and_(self.queue_model.crawl_at <= utcnow_timestamp(),
partition.in_(partitions))).\
limit(score_window).\
subquery()
netloc_query = self.session.query(
self.queue_model,
score_query.c.score_rank,
func.row_number().over(
partition_by=[
partition,
netloc
],
order_by=score_query.c.score_rank
).label('netloc_rank')
).\
select_entity_from(score_query).\
subquery()
partition_query = self.session.query(
self.queue_model,
func.row_number().over(
partition_by=partition,
order_by=netloc_query.c.score_rank
).label('partition_rank')
).\
select_entity_from(netloc_query).\
filter(netloc_query.c.netloc_rank <= self.max_request_per_host).\
subquery()
return self.session.query(self.queue_model).\
select_entity_from(partition_query).\
options(with_expression(self.queue_model.partition_id, partition)).\
filter(partition_query.c.partition_rank <= max_n_requests)
def test_reuse_expr(self):
A = self.classes.A
s = Session()
# so people will obv. want to say, "filter(A.my_expr > 10)".
# but that means Query or Core has to post-modify the statement
# after construction.
expr = A.x + A.y
a1 = s.query(A).options(
with_expression(A.my_expr, expr)).filter(expr > 10).\
order_by(expr)
eq_(
a1.all(),
[A(my_expr=12), A(my_expr=15)]
)
def assignable(self, task_set_id=None, has_skippable=False):
assign = so.aliased(Assignment)
uid = sa.bindparam("user_id")
stmt = (session.query(
assign.task_id,
sa.func.count(assign.task_id).label("assignments_count"),
sa.func.bool_or(assign.user_id != uid).label("started_by_anyone"),
sa.func.bool_or(
assign.user_id == uid).label("previously_assigned"),
).filter(assign.is_relevant()).group_by(assign.task_id).subquery())
query = (self.join(
TaskSet, Task.task_set).join(Group, TaskSet.groups).join(
UserGroup, Group.id == UserGroup.group_id).outerjoin(
Exclusion,
(Exclusion.task_id == Task.id) &
(Exclusion.user_id == uid),
).outerjoin(stmt).filter(UserGroup.user_id == uid).filter(
Task.status == Task.Status.PROC).filter(
Exclusion.user_id.is_(None)).filter(
Task.assigned_group_id.is_(None)
| (Task.assigned_group_id == Group.id)).
filter(~sa.func.coalesce(stmt.c.previously_assigned, False)
).filter(
sa.func.coalesce(stmt.c.assignments_count, 0) <
Task.assignments_required).order_by(
sa.desc(Task.priority)))
if task_set_id:
query = query.filter(Task.task_set_id == task_set_id)
if has_skippable:
is_skippable_expr = (
TaskSet.allow_skip
& (Task.type != Task.Type.CC)
& (~Task.is_pushback)
& (Task.assigned_group_id.is_(None)
| (Task.assigned_group_id == TaskSet.default_group_id)
| ~sa.func.coalesce(stmt.c.started_by_anyone, False)))
query = query.options(
so.with_expression(Task.is_skippable, is_skippable_expr))
return query
def execute(self, db: Session) -> Query:
sample_subquery = BiosampleQuerySchema(
conditions=self.conditions).query(db).subquery()
sample_count = (db.query(
models.Biosample.study_id.label("study_id"),
func.count(models.Biosample.id).label("sample_count"),
).join(sample_subquery,
models.Biosample.id == sample_subquery.c.id).group_by(
models.Biosample.study_id)).subquery()
model = self.table.model
subquery = self.query(db).subquery()
return self._inject_omics_data_summary(
db,
db.query(model).join(subquery,
model.id == subquery.c.id) # type: ignore
.join(sample_count,
model.id == sample_count.c.study_id,
isouter=True) # type: ignore
.options(
with_expression(models.Study.sample_count,
sample_count.c.sample_count)),
)
from sqlalchemy import Column
from sqlalchemy import Integer
from sqlalchemy.orm import declarative_base
from sqlalchemy.orm import Mapped
from sqlalchemy.orm import query_expression
from sqlalchemy.orm import Session
from sqlalchemy.orm import with_expression
Base = declarative_base()
class User(Base):
__tablename__ = "users"
id = Column(Integer, primary_key=True)
foo = Column(Integer)
question_count: Mapped[int] = query_expression()
answer_count: int = query_expression()
s = Session()
q = s.query(User).options(with_expression(User.question_count, User.foo + 5))
u1: User = cast(User, q.first())
qc: int = u1.question_count
print(qc)
