Cloned DMOJ

judge/utils/raw_sql.py

@@ -0,0 +1,96 @@
+from copy import copy
+
+from django.db import connections
+from django.db.models import Field
+from django.db.models.expressions import RawSQL
+from django.db.models.sql.constants import INNER, LOUTER
+from django.db.models.sql.datastructures import Join
+from django.utils import six
+
+from judge.utils.cachedict import CacheDict
+
+
+def unique_together_left_join(queryset, model, link_field_name, filter_field_name, filter_value, parent_model=None):
+    link_field = copy(model._meta.get_field(link_field_name).remote_field)
+    filter_field = model._meta.get_field(filter_field_name)
+
+    def restrictions(where_class, alias, related_alias):
+        cond = where_class()
+        cond.add(filter_field.get_lookup('exact')(filter_field.get_col(alias), filter_value), 'AND')
+        return cond
+
+    link_field.get_extra_restriction = restrictions
+
+    if parent_model is not None:
+        parent_alias = parent_model._meta.db_table
+    else:
+        parent_alias = queryset.query.get_initial_alias()
+    return queryset.query.join(Join(model._meta.db_table, parent_alias, None, LOUTER, link_field, True))
+
+
+class RawSQLJoin(Join):
+    def __init__(self, subquery, subquery_params, parent_alias, table_alias, join_type, join_field, nullable,
+                 filtered_relation=None):
+        self.subquery_params = subquery_params
+        super().__init__(subquery, parent_alias, table_alias, join_type, join_field, nullable, filtered_relation)
+
+    def as_sql(self, compiler, connection):
+        compiler.quote_cache[self.table_name] = '(%s)' % self.table_name
+        sql, params = super().as_sql(compiler, connection)
+        return sql, self.subquery_params + params
+
+
+class FakeJoinField:
+    def __init__(self, joining_columns):
+        self.joining_columns = joining_columns
+
+    def get_joining_columns(self):
+        return self.joining_columns
+
+    def get_extra_restriction(self, where_class, alias, remote_alias):
+        pass
+
+
+def join_sql_subquery(queryset, subquery, params, join_fields, alias, join_type=INNER, parent_model=None):
+    if parent_model is not None:
+        parent_alias = parent_model._meta.db_table
+    else:
+        parent_alias = queryset.query.get_initial_alias()
+    queryset.query.external_aliases.add(alias)
+    join = RawSQLJoin(subquery, params, parent_alias, alias, join_type, FakeJoinField(join_fields), join_type == LOUTER)
+    queryset.query.join(join)
+    join.table_alias = alias
+
+
+def RawSQLColumn(model, field=None):
+    if isinstance(model, Field):
+        field = model
+        model = field.model
+    if isinstance(field, six.string_types):
+        field = model._meta.get_field(field)
+    return RawSQL('%s.%s' % (model._meta.db_table, field.get_attname_column()[1]), ())
+
+
+def make_straight_join_query(QueryType):
+    class Query(QueryType):
+        def join(self, join, *args, **kwargs):
+            alias = super().join(join, *args, **kwargs)
+            join = self.alias_map[alias]
+            if join.join_type == INNER:
+                join.join_type = 'STRAIGHT_JOIN'
+            return alias
+
+    return Query
+
+
+straight_join_cache = CacheDict(make_straight_join_query)
+
+
+def use_straight_join(queryset):
+    if connections[queryset.db].vendor != 'mysql':
+        return
+    try:
+        cloner = queryset.query.chain
+    except AttributeError:
+        cloner = queryset.query.clone
+    queryset.query = cloner(straight_join_cache[type(queryset.query)])