<s> """<STR_LIT>""" <EOL> import uuid <EOL> from peewee import * <EOL> from peewee import Expression <EOL> from peewee import logger <EOL> from peewee import Node <EOL> from peewee import OP <EOL> from peewee import Param <EOL> from peewee import Passthrough <EOL> from peewee import returns_clone <EOL> from peewee import QueryCompiler <EOL> from peewee import SelectQuery <EOL> from peewee import UUIDField <EOL> try : <EOL> from psycopg2cffi import compat <EOL> compat . register ( ) <EOL> except ImportError : <EOL> pass <EOL> from psycopg2 . extensions import adapt <EOL> from psycopg2 . extensions import AsIs <EOL> from psycopg2 . extensions import register_adapter <EOL> from psycopg2 . extras import register_hstore <EOL> try : <EOL> from psycopg2 . extras import Json <EOL> except : <EOL> Json = None <EOL> @ Node . extend ( clone = False ) <EOL> def cast ( self , as_type ) : <EOL> return Expression ( self , OP . CAST , SQL ( as_type ) ) <EOL> class _LookupNode ( Node ) : <EOL> def __init__ ( self , node , parts ) : <EOL> self . node = node <EOL> self . parts = parts <EOL> super ( _LookupNode , self ) . __init__ ( ) <EOL> def clone_base ( self ) : <EOL> return type ( self ) ( self . node , list ( self . parts ) ) <EOL> def cast ( self , as_type ) : <EOL> return Expression ( Clause ( self , parens = True ) , OP . CAST , SQL ( as_type ) ) <EOL> class _JsonLookupBase ( _LookupNode ) : <EOL> def __init__ ( self , node , parts , as_json = False ) : <EOL> super ( _JsonLookupBase , self ) . __init__ ( node , parts ) <EOL> self . _as_json = as_json <EOL> def clone_base ( self ) : <EOL> return type ( self ) ( self . node , list ( self . parts ) , self . _as_json ) <EOL> @ returns_clone <EOL> def as_json ( self , as_json = True ) : <EOL> self . _as_json = as_json <EOL> def contains ( self , other ) : <EOL> clone = self . as_json ( True ) <EOL> if isinstance ( other , ( list , dict ) ) : <EOL> return Expression ( clone , OP . JSONB_CONTAINS , Json ( other ) ) <EOL> return Expression ( clone , OP . JSONB_EXISTS , other ) <EOL> def contains_any ( self , * keys ) : <EOL> return Expression ( <EOL> self . as_json ( True ) , <EOL> OP . JSONB_CONTAINS_ANY_KEY , <EOL> Passthrough ( list ( keys ) ) ) <EOL> def contains_all ( self , * keys ) : <EOL> return Expression ( <EOL> self . as_json ( True ) , <EOL> OP . JSONB_CONTAINS_ALL_KEYS , <EOL> Passthrough ( list ( keys ) ) ) <EOL> class JsonLookup ( _JsonLookupBase ) : <EOL> _node_type = '<STR_LIT>' <EOL> def __getitem__ ( self , value ) : <EOL> return JsonLookup ( self . node , self . parts + [ value ] , self . _as_json ) <EOL> class JsonPath ( _JsonLookupBase ) : <EOL> _node_type = '<STR_LIT>' <EOL> class ObjectSlice ( _LookupNode ) : <EOL> _node_type = '<STR_LIT>' <EOL> @ classmethod <EOL> def create ( cls , node , value ) : <EOL> if isinstance ( value , slice ) : <EOL> parts = [ value . start or <NUM_LIT:0> , value . stop or <NUM_LIT:0> ] <EOL> elif isinstance ( value , int ) : <EOL> parts = [ value ] <EOL> else : <EOL> parts = map ( int , value . split ( '<STR_LIT::>' ) ) <EOL> return cls ( node , parts ) <EOL> def __getitem__ ( self , value ) : <EOL> return ObjectSlice . create ( self , value ) <EOL> class _Array ( Node ) : <EOL> def __init__ ( self , field , items ) : <EOL> self . field = field <EOL> self . items = items <EOL> super ( _Array , self ) . __init__ ( ) <EOL> def adapt_array ( arr ) : <EOL> conn = arr . field . model_class . _meta . database . get_conn ( ) <EOL> items = adapt ( arr . items ) <EOL> items . prepare ( conn ) <EOL> return AsIs ( '<STR_LIT>' % ( <EOL> items , <EOL> arr . field . get_column_type ( ) , <EOL> '<STR_LIT>' * arr . field . dimensions ) ) <EOL> register_adapter ( _Array , adapt_array ) <EOL> class IndexedFieldMixin ( object ) : <EOL> default_index_type = '<STR_LIT>' <EOL> def __init__ ( self , index_type = None , * args , ** kwargs ) : <EOL> kwargs . setdefault ( '<STR_LIT:index>' , True ) <EOL> super ( IndexedFieldMixin , self ) . __init__ ( * args , ** kwargs ) <EOL> if self . index : <EOL> self . index_type = index_type or self . default_index_type <EOL> else : <EOL> self . index_type = None <EOL> class ArrayField ( IndexedFieldMixin , Field ) : <EOL> default_index_type = '<STR_LIT>' <EOL> def __init__ ( self , field_class = IntegerField , dimensions = <NUM_LIT:1> , * args , <EOL> ** kwargs ) : <EOL> self . __field = field_class ( * args , ** kwargs ) <EOL> self . dimensions = dimensions <EOL> self . db_field = self . __field . get_db_field ( ) <EOL> super ( ArrayField , self ) . __init__ ( * args , ** kwargs ) <EOL> def __ddl_column__ ( self , column_type ) : <EOL> sql = self . __field . __ddl_column__ ( column_type ) <EOL> sql . value += '<STR_LIT>' * self . dimensions <EOL> return sql <EOL> def db_value ( self , value ) : <EOL> if value is not None and not isinstance ( value , ( list , _Array ) ) : <EOL> return list ( value ) <EOL> return value <EOL> def __getitem__ ( self , value ) : <EOL> return ObjectSlice . create ( self , value ) <EOL> def contains ( self , * items ) : <EOL> return Expression ( self , OP . ACONTAINS , _Array ( self , list ( items ) ) ) <EOL> def contains_any ( self , * items ) : <EOL> return Expression ( self , OP . ACONTAINS_ANY , _Array ( self , list ( items ) ) ) <EOL> class DateTimeTZField ( DateTimeField ) : <EOL> db_field = '<STR_LIT>' <EOL> class HStoreField ( IndexedFieldMixin , Field ) : <EOL> db_field = '<STR_LIT>' <EOL> def __getitem__ ( self , key ) : <EOL> return Expression ( self , OP . HKEY , Param ( key ) ) <EOL> def keys ( self ) : <EOL> return fn . akeys ( self ) <EOL> def values ( self ) : <EOL> return fn . avals ( self ) <EOL> def items ( self ) : <EOL> return fn . hstore_to_matrix ( self ) <EOL> def slice ( self , * args ) : <EOL> return fn . slice ( self , Param ( list ( args ) ) ) <EOL> def exists ( self , key ) : <EOL> return fn . exist ( self , key ) <EOL> def defined ( self , key ) : <EOL> return fn . defined ( self , key ) <EOL> def update ( self , ** data ) : <EOL> return Expression ( self , OP . HUPDATE , data ) <EOL> def delete ( self , * keys ) : <EOL> return fn . delete ( self , Param ( list ( keys ) ) ) <EOL> def contains ( self , value ) : <EOL> if isinstance ( value , dict ) : <EOL> return Expression ( self , OP . HCONTAINS_DICT , Param ( value ) ) <EOL> elif isinstance ( value , ( list , tuple ) ) : <EOL> return Expression ( self , OP . HCONTAINS_KEYS , Param ( value ) ) <EOL> return Expression ( self , OP . HCONTAINS_KEY , value ) <EOL> def contains_any ( self , * keys ) : <EOL> return Expression ( self , OP . HCONTAINS_ANY_KEY , Param ( list ( keys ) ) ) <EOL> class JSONField ( Field ) : <EOL> db_field = '<STR_LIT>' <EOL> def __init__ ( self , dumps = None , * args , ** kwargs ) : <EOL> if Json is None : <EOL> raise Exception ( '<STR_LIT>' ) <EOL> self . dumps = dumps <EOL> super ( JSONField , self ) . __init__ ( * args , ** kwargs ) <EOL> def db_value ( self , value ) : <EOL> if value is None : <EOL> return value <EOL> if not isinstance ( value , Json ) : <EOL> return Json ( value , dumps = self . dumps ) <EOL> return value <EOL> def __getitem__ ( self , value ) : <EOL> return JsonLookup ( self , [ value ] ) <EOL> def path ( self , * keys ) : <EOL> return JsonPath ( self , keys ) <EOL> class BinaryJSONField ( IndexedFieldMixin , JSONField ) : <EOL> db_field = '<STR_LIT>' <EOL> default_index_type = '<STR_LIT>' <EOL> def contains ( self , other ) : <EOL> if isinstance ( other , ( list , dict ) ) : <EOL> return Expression ( self , OP . JSONB_CONTAINS , Json ( other ) ) <EOL> return Expression ( self , OP . JSONB_EXISTS , Passthrough ( other ) ) <EOL> def contained_by ( self , other ) : <EOL> return Expression ( self , OP . JSONB_CONTAINED_BY , Json ( other ) ) <EOL> def contains_any ( self , * items ) : <EOL> return Expression ( <EOL> self , <EOL> OP . JSONB_CONTAINS_ANY_KEY , <EOL> Passthrough ( list ( items ) ) ) <EOL> def contains_all ( self , * items ) : <EOL> return Expression ( <EOL> self , <EOL> OP . JSONB_CONTAINS_ALL_KEYS , <EOL> Passthrough ( list ( items ) ) ) <EOL> class TSVectorField ( IndexedFieldMixin , TextField ) : <EOL> db_field = '<STR_LIT>' <EOL> default_index_type = '<STR_LIT>' <EOL> def match ( self , query ) : <EOL> return Expression ( self , OP . TS_MATCH , fn . to_tsquery ( query ) ) <EOL> def Match ( field , query ) : <EOL> return Expression ( fn . to_tsvector ( field ) , OP . TS_MATCH , fn . to_tsquery ( query ) ) <EOL> OP . update ( <EOL> HKEY = '<STR_LIT:key>' , <EOL> HUPDATE = '<STR_LIT>' , <EOL> HCONTAINS_DICT = '<STR_LIT>' , <EOL> HCONTAINS_KEYS = '<STR_LIT>' , <EOL> HCONTAINS_KEY = '<STR_LIT>' , <EOL> HCONTAINS_ANY_KEY = '<STR_LIT>' , <EOL> ACONTAINS = '<STR_LIT>' , <EOL> ACONTAINS_ANY = '<STR_LIT>' , <EOL> TS_MATCH = '<STR_LIT>' , <EOL> JSONB_CONTAINS = '<STR_LIT>' , <EOL> JSONB_CONTAINED_BY = '<STR_LIT>' , <EOL> JSONB_CONTAINS_ANY_KEY = '<STR_LIT>' , <EOL> JSONB_CONTAINS_ALL_KEYS = '<STR_LIT>' , <EOL> JSONB_EXISTS = '<STR_LIT>' , <EOL> CAST = '<STR_LIT>' , <EOL> ) <EOL> class PostgresqlExtCompiler ( QueryCompiler ) : <EOL> def _create_index ( self , model_class , fields , unique = False ) : <EOL> clause = super ( PostgresqlExtCompiler , self ) . _create_index ( <EOL> model_class , fields , unique ) <EOL> index_type = None <EOL> for field in fields : <EOL> if isinstance ( field , IndexedFieldMixin ) : <EOL> index_type = field . index_type <EOL> if index_type : <EOL> clause . nodes . insert ( - <NUM_LIT:1> , SQL ( '<STR_LIT>' % index_type ) ) <EOL> return clause <EOL> def _parse_object_slice ( self , node , alias_map , conv ) : <EOL> sql , params = self . parse_node ( node . node , alias_map , conv ) <EOL> parts = [ str ( part + <NUM_LIT:1> ) for part in node . parts ] <EOL> sql = '<STR_LIT>' % ( sql , '<STR_LIT::>' . join ( parts ) ) <EOL> return sql , params <EOL> def _parse_json_lookup ( self , node , alias_map , conv ) : <EOL> sql , params = self . parse_node ( node . node , alias_map , conv ) <EOL> lookups = [ sql ] <EOL> for part in node . parts : <EOL> part_sql , part_params = self . parse_node ( <EOL> part , alias_map , conv ) <EOL> lookups . append ( part_sql ) <EOL> params . extend ( part_params ) <EOL> if node . _as_json : <EOL> sql = '<STR_LIT>' . join ( lookups ) <EOL> else : <EOL> head , tail = lookups [ : - <NUM_LIT:1> ] , lookups [ - <NUM_LIT:1> ] <EOL> sql = '<STR_LIT>' . join ( ( '<STR_LIT>' . join ( head ) , tail ) ) <EOL> return sql , params <EOL> def _parse_json_path ( self , node , alias_map , conv ) : <EOL> sql , params = self . parse_node ( node . node , alias_map , conv ) <EOL> if node . _as_json : <EOL> operand = '<STR_LIT>' <EOL> else : <EOL> operand = '<STR_LIT>' <EOL> params . append ( '<STR_LIT>' % '<STR_LIT:U+002C>' . join ( map ( str , node . parts ) ) ) <EOL> return operand . join ( ( sql , self . interpolation ) ) , params <EOL> def get_parse_map ( self ) : <EOL> parse_map = super ( PostgresqlExtCompiler , self ) . get_parse_map ( ) <EOL> parse_map . update ( <EOL> object_slice = self . _parse_object_slice , <EOL> json_lookup = self . _parse_json_lookup , <EOL> json_path = self . _parse_json_path ) <EOL> return parse_map <EOL> class PostgresqlExtDatabase ( PostgresqlDatabase ) : <EOL> compiler_class = PostgresqlExtCompiler <EOL> def __init__ ( self , * args , ** kwargs ) : <EOL> self . server_side_cursors = kwargs . pop ( '<STR_LIT>' , False ) <EOL> self . register_hstore = kwargs . pop ( '<STR_LIT>' , True ) <EOL> super ( PostgresqlExtDatabase , self ) . __init__ ( * args , ** kwargs ) <EOL> def get_cursor ( self , name = None ) : <EOL> if name : <EOL> return self . get_conn ( ) . cursor ( name = name ) <EOL> return self . get_conn ( ) . cursor ( ) <EOL> def execute_sql ( self , sql , params = None , require_commit = True , <EOL> named_cursor = False ) : <EOL> logger . debug ( ( sql , params ) ) <EOL> use_named_cursor = ( named_cursor or ( <EOL> self . server_side_cursors and <EOL> sql . lower ( ) . startswith ( '<STR_LIT>' ) ) ) <EOL> with self . exception_wrapper ( ) : <EOL> if use_named_cursor : <EOL> cursor = self . get_cursor ( name = str ( uuid . uuid1 ( ) ) ) <EOL> require_commit = False <EOL> else : <EOL> cursor = self . get_cursor ( ) <EOL> try : <EOL> cursor . execute ( sql , params or ( ) ) <EOL> except Exception as exc : <EOL> if self . get_autocommit ( ) and self . autorollback : <EOL> self . rollback ( ) <EOL> raise <EOL> else : <EOL> if require_commit and self . get_autocommit ( ) : <EOL> self . commit ( ) <EOL> return cursor <EOL> def _connect ( self , database , ** kwargs ) : <EOL> conn = super ( PostgresqlExtDatabase , self ) . _connect ( database , ** kwargs ) <EOL> if self . register_hstore : <EOL> register_hstore ( conn , globally = True ) <EOL> return conn <EOL> class ServerSideSelectQuery ( SelectQuery ) : <EOL> @ classmethod <EOL> def clone_from_query ( cls , query ) : <EOL> clone = ServerSideSelectQuery ( query . model_class ) <EOL> return query . _clone_attributes ( clone ) <EOL> def _execute ( self ) : <EOL> sql , params = self . sql ( ) <EOL> return self . database . execute_sql ( <EOL> sql , params , require_commit = False , named_cursor = True ) <EOL> PostgresqlExtDatabase . register_fields ( { <EOL> '<STR_LIT>' : '<STR_LIT>' , <EOL> '<STR_LIT>' : '<STR_LIT>' , <EOL> '<STR_LIT>' : '<STR_LIT>' , <EOL> '<STR_LIT>' : '<STR_LIT>' , <EOL> '<STR_LIT>' : '<STR_LIT>' , <EOL> } ) <EOL> PostgresqlExtDatabase . register_ops ( { <EOL> OP . HCONTAINS_DICT : '<STR_LIT>' , <EOL> OP . HCONTAINS_KEYS : '<STR_LIT>' , <EOL> OP . HCONTAINS_KEY : '<STR_LIT:?>' , <EOL> OP . HCONTAINS_ANY_KEY : '<STR_LIT>' , <EOL> OP . HKEY : '<STR_LIT>' , <EOL> OP . HUPDATE : '<STR_LIT>' , <EOL> OP . ACONTAINS : '<STR_LIT>' , <EOL> OP . ACONTAINS_ANY : '<STR_LIT>' , <EOL> OP . TS_MATCH : '<STR_LIT>' , <EOL> OP . JSONB_CONTAINS : '<STR_LIT>' , <EOL> OP . JSONB_CONTAINED_BY : '<STR_LIT>' , <EOL> OP . JSONB_CONTAINS_ANY_KEY : '<STR_LIT>' , <EOL> OP . JSONB_CONTAINS_ALL_KEYS : '<STR_LIT>' , <EOL> OP . JSONB_EXISTS : '<STR_LIT:?>' , <EOL> OP . CAST : '<STR_LIT>' , <EOL> } ) <EOL> def ServerSide ( select_query ) : <EOL> clone = ServerSideSelectQuery . clone_from_query ( select_query ) <EOL> with clone . database . transaction ( ) : <EOL> query_result = clone . execute ( ) <EOL> select_query . _qr = query_result <EOL> for obj in query_result . iterator ( ) : <EOL> yield obj <EOL> def LateralJoin ( lhs , rhs , join_type = '<STR_LIT>' , condition = True ) : <EOL> return Clause ( <EOL> lhs , <EOL> SQL ( '<STR_LIT>' % join_type ) , <EOL> rhs , <EOL> SQL ( '<STR_LIT>' , condition ) ) </s>
