<commit_before>import time
import copy
import sqlalchemy

from blitzdb.queryset import QuerySet as BaseQuerySet
from functools import wraps
from sqlalchemy.sql import select,func,expression,delete,distinct,and_,union,intersect
from sqlalchemy.sql.expression import join,asc,desc,outerjoin
from ..file.serializers import JsonSerializer
from .helpers import get_value
from collections import OrderedDict
from blitzdb.fields import ManyToManyField,ForeignKeyField,OneToManyField

class QuerySet(BaseQuerySet):

    def __init__(self, backend, table, cls,
                 condition = None,
                 select = None,
                 intersects = None,
                 raw = False,
                 include = None,
                 only = None,
                 joins = None,
                 extra_fields = None,
                 group_bys = None,
                 objects = None,
                 havings = None,
                 limit = None,
                 offset = None
                 ):
        super(QuerySet,self).__init__(backend = backend,cls = cls)

        self.joins = joins
        self.backend = backend
        self.condition = condition
        self.select = select
        self.havings = havings
        self.only = only
        self.include = include
        self.extra_fields = extra_fields
        self.group_bys = group_bys
        self.cls = cls
        self._limit = limit
        self._offset = offset
        self.table = table
        self.raw = raw
        self.intersects = intersects
        self.objects = objects
        if self.objects:
            self.pop_objects = self.objects[:]

        self.deserialized_objects = None
        self.deserialized_pop_objects = None
        self._it = None
        self.order_bys = None
        self.count = None
        self.result = None

    def limit(self,limit):
        self._limit = limit
        return self

    def offset(self,offset):
        self._offset = offset
        return self

    def deserialize(self, data):

        d,lazy = self.backend.deserialize_db_data(data)

        if self.raw:
            return d

        obj = self.backend.create_instance(self.cls, d,lazy = lazy)
        obj.attributes = self.backend.deserialize(obj.lazy_attributes,create_instance = False)

        return obj

    def sort(self, keys,direction = None):
        #we sort by a single argument
        if direction:
            keys = ((keys,direction),)
        order_bys = []
        for key,direction in keys:
            if direction > 0:
                direction = asc
            else:
                direction = desc
            try:
                column = self.backend.get_column_for_key(self.cls,key)
            except KeyError:
                raise AttributeError("Attempting to sort results by a non-indexed field %s" % key)
            order_bys.append((column,direction))
        self.order_bys = order_bys
        self.objects = None
        return self

    def next(self):
        if self._it is None:
            self._it = iter(self)
        return self._it.next()

    __next__ = next

    def __iter__(self):
        if self.deserialized_objects is None:
            self.get_deserialized_objects()
        for obj in self.deserialized_objects:
            yield obj
        raise StopIteration

    def __contains__(self, obj):
        pks = self.distinct_pks()
        if isinstance(obj, list) or isinstance(obj, tuple):
            obj_list = obj
        else:
            obj_list = [obj]
        for obj in obj_list:
            if obj.pk not in pks:
                return False
        return True

    def get_deserialized_objects(self):
        if self.objects is None:
            self.get_objects()

        self.deserialized_objects = [self.deserialize(obj) for obj in self.objects]
        self.deserialized_pop_objects = self.deserialized_objects[:]

    def get_select_table_and_rows(self):

        columns = []
        joins = []

        def join_table(collection,table,key,params,path = None):
            if path is None:
                path = []
            if isinstance(params['relation']['field'],ManyToManyField):
                join_many_to_many(collection,table,key,params,path)
            elif isinstance(params['relation']['field'],ForeignKeyField):
                join_foreign_key(collection,table,key,params,path)
            elif isinstance(params['relation']['field'],OneToManyField):
                join_one_to_many(collection,table,key,params,path)
            else:
                raise AttributeError

        def process_fields_and_subkeys(related_collection,related_table,params,path):

            params['table_fields'] = {}
            for field,column_name in params['fields'].items():
                column_label = '_'.join(path+[column_name])
                params['table_fields'][field] = column_label
                columns.append(related_table.c[column_name].label(column_label))

            for subkey,subparams in sorted(params['joins'].items(),key = lambda i : i[0]):
                join_table(params['collection'],related_table,subkey,subparams,path = path)

        def join_one_to_many(collection,table,key,params,path):
            related_table = params['table'].alias()
            related_collection = params['relation']['collection']
            condition = table.c['pk'] == related_table.c[params['relation']['backref']['column']]
            joins.append((related_table,condition))
            process_fields_and_subkeys(related_collection,related_table,params,path+\
                                        [params['relation']['backref']['column']])

        def join_foreign_key(collection,table,key,params,path):
            related_table = params['table'].alias()
            related_collection = params['relation']['collection']
            condition = table.c[params['relation']['column']] == related_table.c.pk
            joins.append((related_table,condition))
            process_fields_and_subkeys(related_collection,related_table,params,path+\
                                        [params['relation']['column']])

        def join_many_to_many(collection,table,key,params,path):
            relationship_table = params['relation']['relationship_table'].alias()
            related_collection = params['relation']['collection']
            related_table = self.backend.get_collection_table(related_collection).alias()
            left_condition = relationship_table.c['pk_%s' % collection] == table.c.pk
            right_condition = relationship_table.c['pk_%s' % related_collection] == related_table.c.pk
            joins.append((relationship_table,left_condition))
            joins.append((related_table,right_condition))
            process_fields_and_subkeys(related_collection,related_table,params,path+[key])

        if self.include:
            include = copy.deepcopy(self.include)
            if not isinstance(include,(list,tuple)):
                raise AttributeError("include must be a list/tuple")
        else:
            include = ()

        exclude = []
        if self.only:
            if isinstance(self.only,dict):
                only = []
                for key,value in self.only.items():
                    if value is False:
                        exclude.append(key)
                    else:
                        only.append(key)
            else:
                only = set(self.only)
            include = set(include)
            for only_key in only:
                include.add(only_key)

        self.include_joins = self.backend.get_include_joins(self.cls,includes = include,excludes = exclude)

        my_columns = self.include_joins['fields'].values()+\
                     [params['relation']['column'] for params in self.include_joins['joins'].values()
                      if isinstance(params['relation']['field'],ForeignKeyField)]
        s = self.get_select(fields = [self.table.c[column] for column in my_columns])
        s_cte = s.cte()<commit_msg>Automatically include all sort_by keys in include list.<commit_after>import time
import copy
import sqlalchemy

from blitzdb.queryset import QuerySet as BaseQuerySet
from functools import wraps
from sqlalchemy.sql import select,func,expression,delete,distinct,and_,union,intersect
from sqlalchemy.sql.expression import join,asc,desc,outerjoin
from ..file.serializers import JsonSerializer
from .helpers import get_value
from collections import OrderedDict
from blitzdb.fields import ManyToManyField,ForeignKeyField,OneToManyField

class QuerySet(BaseQuerySet):

    def __init__(self, backend, table, cls,
                 condition = None,
                 select = None,
                 intersects = None,
                 raw = False,
                 include = None,
                 only = None,
                 joins = None,
                 extra_fields = None,
                 group_bys = None,
                 objects = None,
                 havings = None,
                 limit = None,
                 offset = None
                 ):
        super(QuerySet,self).__init__(backend = backend,cls = cls)

        self.joins = joins
        self.backend = backend
        self.condition = condition
        self.select = select
        self.havings = havings
        self.only = only
        self.include = include
        self.extra_fields = extra_fields
        self.group_bys = group_bys
        self.cls = cls
        self._limit = limit
        self._offset = offset
        self.table = table
        self.raw = raw
        self.intersects = intersects
        self.objects = objects
        if self.objects:
            self.pop_objects = self.objects[:]

        self.deserialized_objects = None
        self.deserialized_pop_objects = None
        self._it = None
        self.order_bys = None
        self.order_bys_keys = None
        self.count = None
        self.result = None

    def limit(self,limit):
        self._limit = limit
        return self

    def offset(self,offset):
        self._offset = offset
        return self

    def deserialize(self, data):

        d,lazy = self.backend.deserialize_db_data(data)

        if self.raw:
            return d

        obj = self.backend.create_instance(self.cls, d,lazy = lazy)
        obj.attributes = self.backend.deserialize(obj.lazy_attributes,create_instance = False)

        return obj

    def sort(self, keys,direction = None):
        #we sort by a single argument
        if direction:
            keys = ((keys,direction),)
        order_bys = []
        for key,direction in keys:
            if direction > 0:
                direction = asc
            else:
                direction = desc
            try:
                column = self.backend.get_column_for_key(self.cls,key)
            except KeyError:
                raise AttributeError("Attempting to sort results by a non-indexed field %s" % key)
            order_bys.append((column,direction))
        self.order_bys = order_bys
        self.order_bys_keys = keys
        self.objects = None
        return self

    def next(self):
        if self._it is None:
            self._it = iter(self)
        return self._it.next()

    __next__ = next

    def __iter__(self):
        if self.deserialized_objects is None:
            self.get_deserialized_objects()
        for obj in self.deserialized_objects:
            yield obj
        raise StopIteration

    def __contains__(self, obj):
        pks = self.distinct_pks()
        if isinstance(obj, list) or isinstance(obj, tuple):
            obj_list = obj
        else:
            obj_list = [obj]
        for obj in obj_list:
            if obj.pk not in pks:
                return False
        return True

    def get_deserialized_objects(self):
        if self.objects is None:
            self.get_objects()

        self.deserialized_objects = [self.deserialize(obj) for obj in self.objects]
        self.deserialized_pop_objects = self.deserialized_objects[:]

    def get_select_table_and_rows(self):

        columns = []
        joins = []

        def join_table(collection,table,key,params,path = None):
            if path is None:
                path = []
            if isinstance(params['relation']['field'],ManyToManyField):
                join_many_to_many(collection,table,key,params,path)
            elif isinstance(params['relation']['field'],ForeignKeyField):
                join_foreign_key(collection,table,key,params,path)
            elif isinstance(params['relation']['field'],OneToManyField):
                join_one_to_many(collection,table,key,params,path)
            else:
                raise AttributeError

        def process_fields_and_subkeys(related_collection,related_table,params,path):

            params['table_fields'] = {}
            for field,column_name in params['fields'].items():
                column_label = '_'.join(path+[column_name])
                params['table_fields'][field] = column_label
                columns.append(related_table.c[column_name].label(column_label))

            for subkey,subparams in sorted(params['joins'].items(),key = lambda i : i[0]):
                join_table(params['collection'],related_table,subkey,subparams,path = path)

        def join_one_to_many(collection,table,key,params,path):
            related_table = params['table'].alias()
            related_collection = params['relation']['collection']
            condition = table.c['pk'] == related_table.c[params['relation']['backref']['column']]
            joins.append((related_table,condition))
            process_fields_and_subkeys(related_collection,related_table,params,path+\
                                        [params['relation']['backref']['column']])

        def join_foreign_key(collection,table,key,params,path):
            related_table = params['table'].alias()
            related_collection = params['relation']['collection']
            condition = table.c[params['relation']['column']] == related_table.c.pk
            joins.append((related_table,condition))
            process_fields_and_subkeys(related_collection,related_table,params,path+\
                                        [params['relation']['column']])

        def join_many_to_many(collection,table,key,params,path):
            relationship_table = params['relation']['relationship_table'].alias()
            related_collection = params['relation']['collection']
            related_table = self.backend.get_collection_table(related_collection).alias()
            left_condition = relationship_table.c['pk_%s' % collection] == table.c.pk
            right_condition = relationship_table.c['pk_%s' % related_collection] == related_table.c.pk
            joins.append((relationship_table,left_condition))
            joins.append((related_table,right_condition))
            process_fields_and_subkeys(related_collection,related_table,params,path+[key])

        if self.include:
            include = copy.deepcopy(self.include)
            if isinstance(include,tuple):
                include = list(include)
            if not isinstance(include,list):
                raise AttributeError("include must be a list/tuple")
        else:
            include = []

        exclude = []
        if self.only:
            if isinstance(self.only,dict):
                only = []
                for key,value in self.only.items():
                    if value is False:
                        exclude.append(key)
                    else:
                        only.append(key)
            else:
                only = set(self.only)
            for only_key in only:
                if not only_key in include:
                    include.append(only_key)

        if self.order_bys_keys:
            for key,direction in self.order_bys_keys:
                if not key in include:
                    include.append(key)

        self.include_joins = self.backend.get_include_joins(self.cls,includes = include,excludes = exclude)

        #we only select the columns that we actually need
        my_columns = self.include_joins['fields'].values()+\
                     [params['relation']['column'] for params in self.include_joins['joins'].values()
                      if isinstance(params['relation']['field'],ForeignKeyField)]
        s = self.get_select(fields = [self.table.c[column] for column in my_columns])
        s_cte = s.cte()
{'fraction_non_alphanumeric': 0.09562941319698076, 'fraction_numerical': 0.00024348672997321646, 'mean_word_length': 2.0508820798514393, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 5, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 4, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '1663588', 'n_tokens_mistral': 4367, 'n_tokens_neox': 4215, 'n_words': 892}