<s> import copy <EOL> from sqlalchemy import inspect <EOL> from sqlalchemy import orm <EOL> from sqlalchemy import sql <EOL> from sqlalchemy import types as sqltypes <EOL> from oslo_db . sqlalchemy import utils <EOL> def update_on_match ( <EOL> query , <EOL> specimen , <EOL> surrogate_key , <EOL> values = None , <EOL> attempts = <NUM_LIT:3> , <EOL> include_only = None , <EOL> process_query = None , <EOL> handle_failure = None <EOL> ) : <EOL> """<STR_LIT>""" <EOL> if values is None : <EOL> values = { } <EOL> entity = inspect ( specimen ) <EOL> mapper = entity . mapper <EOL> assert [ desc [ '<STR_LIT:type>' ] for desc in query . column_descriptions ] == [ mapper . class_ ] , "<STR_LIT>" <EOL> criteria = manufacture_entity_criteria ( <EOL> specimen , include_only = include_only , exclude = [ surrogate_key ] ) <EOL> query = query . filter ( criteria ) <EOL> if process_query : <EOL> query = process_query ( query ) <EOL> surrogate_key_arg = ( <EOL> surrogate_key , entity . attrs [ surrogate_key ] . loaded_value ) <EOL> pk_value = None <EOL> for attempt in range ( attempts ) : <EOL> try : <EOL> pk_value = query . update_returning_pk ( values , surrogate_key_arg ) <EOL> except MultiRowsMatched : <EOL> raise <EOL> except NoRowsMatched : <EOL> if handle_failure and handle_failure ( query ) : <EOL> break <EOL> else : <EOL> break <EOL> else : <EOL> raise NoRowsMatched ( "<STR_LIT>" % attempts ) <EOL> if pk_value is None : <EOL> pk_value = entity . mapper . primary_key_from_instance ( specimen ) <EOL> values = copy . copy ( values ) <EOL> values [ surrogate_key ] = surrogate_key_arg [ <NUM_LIT:1> ] <EOL> persistent_obj = manufacture_persistent_object ( <EOL> query . session , specimen . __class__ ( ) , values , pk_value ) <EOL> return persistent_obj <EOL> def manufacture_persistent_object ( <EOL> session , specimen , values = None , primary_key = None ) : <EOL> """<STR_LIT>""" <EOL> state = inspect ( specimen ) <EOL> mapper = state . mapper <EOL> for k , v in values . items ( ) : <EOL> orm . attributes . set_committed_value ( specimen , k , v ) <EOL> pk_attrs = [ <EOL> mapper . get_property_by_column ( col ) . key <EOL> for col in mapper . primary_key <EOL> ] <EOL> if primary_key is not None : <EOL> for key , value in zip ( pk_attrs , primary_key ) : <EOL> orm . attributes . set_committed_value ( <EOL> specimen , <EOL> key , <EOL> value <EOL> ) <EOL> for key in pk_attrs : <EOL> if state . attrs [ key ] . loaded_value is orm . attributes . NO_VALUE : <EOL> raise ValueError ( "<STR_LIT>" ) <EOL> orm . make_transient_to_detached ( specimen ) <EOL> if state . key not in session . identity_map : <EOL> session . add ( specimen ) <EOL> return specimen <EOL> else : <EOL> return session . merge ( specimen , load = False ) <EOL> def manufacture_entity_criteria ( entity , include_only = None , exclude = None ) : <EOL> """<STR_LIT>""" <EOL> state = inspect ( entity ) <EOL> exclude = set ( exclude ) if exclude is not None else set ( ) <EOL> existing = dict ( <EOL> ( attr . key , attr . loaded_value ) <EOL> for attr in state . attrs <EOL> if attr . loaded_value is not orm . attributes . NO_VALUE <EOL> and attr . key not in exclude <EOL> ) <EOL> if include_only : <EOL> existing = dict ( <EOL> ( k , existing [ k ] ) <EOL> for k in set ( existing ) . intersection ( include_only ) <EOL> ) <EOL> return manufacture_criteria ( state . mapper , existing ) <EOL> def manufacture_criteria ( mapped , values ) : <EOL> """<STR_LIT>""" <EOL> mapper = inspect ( mapped ) <EOL> value_keys = set ( values ) <EOL> keys = [ k for k in mapper . column_attrs . keys ( ) if k in value_keys ] <EOL> return sql . and_ ( * [ <EOL> _sql_crit ( mapper . column_attrs [ key ] . expression , values [ key ] ) <EOL> for key in keys <EOL> ] ) <EOL> def _sql_crit ( expression , value ) : <EOL> """<STR_LIT>""" <EOL> values = utils . to_list ( value , default = ( None , ) ) <EOL> if len ( values ) == <NUM_LIT:1> : <EOL> if values [ <NUM_LIT:0> ] is None : <EOL> return expression == sql . null ( ) <EOL> else : <EOL> return expression == values [ <NUM_LIT:0> ] <EOL> elif _none_set . intersection ( values ) : <EOL> return sql . or_ ( <EOL> expression == sql . null ( ) , <EOL> _sql_crit ( expression , set ( values ) . difference ( _none_set ) ) <EOL> ) <EOL> else : <EOL> return expression . in_ ( values ) <EOL> def update_returning_pk ( query , values , surrogate_key ) : <EOL> """<STR_LIT>""" <EOL> entity = query . column_descriptions [ <NUM_LIT:0> ] [ '<STR_LIT:type>' ] <EOL> mapper = inspect ( entity ) . mapper <EOL> session = query . session <EOL> bind = session . connection ( mapper = mapper ) <EOL> if bind . dialect . implicit_returning : <EOL> pk_strategy = _pk_strategy_returning <EOL> elif bind . dialect . name == '<STR_LIT>' and len ( mapper . primary_key ) == <NUM_LIT:1> and isinstance ( <EOL> mapper . primary_key [ <NUM_LIT:0> ] . type , sqltypes . Integer ) : <EOL> pk_strategy = _pk_strategy_mysql_last_insert_id <EOL> else : <EOL> pk_strategy = _pk_strategy_refetch <EOL> return pk_strategy ( query , mapper , values , surrogate_key ) <EOL> def _assert_single_row ( rows_updated ) : <EOL> if rows_updated == <NUM_LIT:1> : <EOL> return rows_updated <EOL> elif rows_updated > <NUM_LIT:1> : <EOL> raise MultiRowsMatched ( "<STR_LIT>" % rows_updated ) <EOL> else : <EOL> raise NoRowsMatched ( "<STR_LIT>" ) <EOL> def _pk_strategy_refetch ( query , mapper , values , surrogate_key ) : <EOL> surrogate_key_name , surrogate_key_value = surrogate_key <EOL> surrogate_key_col = mapper . attrs [ surrogate_key_name ] . expression <EOL> rowcount = query . filter ( surrogate_key_col == surrogate_key_value ) . update ( values , synchronize_session = False ) <EOL> _assert_single_row ( rowcount ) <EOL> fetch_query = query . session . query ( <EOL> * mapper . primary_key ) . filter ( <EOL> surrogate_key_col == surrogate_key_value ) <EOL> primary_key = fetch_query . one ( ) <EOL> return primary_key <EOL> def _pk_strategy_returning ( query , mapper , values , surrogate_key ) : <EOL> surrogate_key_name , surrogate_key_value = surrogate_key <EOL> surrogate_key_col = mapper . attrs [ surrogate_key_name ] . expression <EOL> update_stmt = _update_stmt_from_query ( mapper , query , values ) <EOL> update_stmt = update_stmt . where ( surrogate_key_col == surrogate_key_value ) <EOL> update_stmt = update_stmt . returning ( * mapper . primary_key ) <EOL> result = query . session . execute ( update_stmt ) <EOL> rowcount = result . rowcount <EOL> _assert_single_row ( rowcount ) <EOL> primary_key = tuple ( result . first ( ) ) <EOL> return primary_key <EOL> def _pk_strategy_mysql_last_insert_id ( query , mapper , values , surrogate_key ) : <EOL> surrogate_key_name , surrogate_key_value = surrogate_key <EOL> surrogate_key_col = mapper . attrs [ surrogate_key_name ] . expression <EOL> surrogate_pk_col = mapper . primary_key [ <NUM_LIT:0> ] <EOL> update_stmt = _update_stmt_from_query ( mapper , query , values ) <EOL> update_stmt = update_stmt . where ( surrogate_key_col == surrogate_key_value ) <EOL> update_stmt = update_stmt . values ( <EOL> { surrogate_pk_col : sql . func . last_insert_id ( surrogate_pk_col ) } ) <EOL> result = query . session . execute ( update_stmt ) <EOL> rowcount = result . rowcount <EOL> _assert_single_row ( rowcount ) <EOL> primary_key = query . session . scalar ( sql . func . last_insert_id ( ) ) , <EOL> return primary_key <EOL> def _update_stmt_from_query ( mapper , query , values ) : <EOL> upd_values = dict ( <EOL> ( <EOL> mapper . column_attrs [ key ] , value <EOL> ) for key , value in values . items ( ) <EOL> ) <EOL> query = query . enable_eagerloads ( False ) <EOL> context = query . _compile_context ( ) <EOL> primary_table = context . statement . froms [ <NUM_LIT:0> ] <EOL> update_stmt = sql . update ( primary_table , <EOL> context . whereclause , <EOL> upd_values ) <EOL> return update_stmt <EOL> _none_set = frozenset ( [ None ] ) <EOL> class CantUpdateException ( Exception ) : <EOL> pass <EOL> class NoRowsMatched ( CantUpdateException ) : <EOL> pass <EOL> class MultiRowsMatched ( CantUpdateException ) : <EOL> pass </s>
