<s> """<STR_LIT>""" <EOL> import itertools <EOL> import os <EOL> import re <EOL> import stat <EOL> import tempfile <EOL> import thread <EOL> import threading <EOL> import time <EOL> import sqlite3 <EOL> import logging <EOL> from grr . lib import aff4 <EOL> from grr . lib import config_lib <EOL> from grr . lib import data_store <EOL> from grr . lib import rdfvalue <EOL> from grr . lib import utils <EOL> from grr . lib . data_stores import common <EOL> SQLITE_EXTENSION = "<STR_LIT>" <EOL> SQLITE_TIMEOUT = <NUM_LIT> <EOL> SQLITE_ISOLATION = "<STR_LIT>" <EOL> SQLITE_SUBJECT_SPEC = "<STR_LIT>" <EOL> SQLITE_DETECT_TYPES = <NUM_LIT:0> <EOL> SQLITE_FACTORY = sqlite3 . Connection <EOL> SQLITE_CACHED_STATEMENTS = <NUM_LIT:20> <EOL> SQLITE_PAGE_SIZE = <NUM_LIT> <EOL> class SqliteConnectionCache ( utils . FastStore ) : <EOL> """<STR_LIT>""" <EOL> template = None <EOL> def _CreateModelDatabase ( self ) : <EOL> try : <EOL> if not os . path . isdir ( self . root_path ) : <EOL> os . makedirs ( self . root_path ) <EOL> except OSError : <EOL> pass <EOL> fd , model = tempfile . mkstemp ( dir = self . root_path ) <EOL> os . close ( fd ) <EOL> conn = sqlite3 . connect ( model , SQLITE_TIMEOUT , SQLITE_DETECT_TYPES , <EOL> SQLITE_ISOLATION , False , SQLITE_FACTORY , <EOL> SQLITE_CACHED_STATEMENTS ) <EOL> cursor = conn . cursor ( ) <EOL> cursor . execute ( "<STR_LIT>" ) <EOL> cursor . execute ( "<STR_LIT>" ) <EOL> cursor . execute ( "<STR_LIT>" ) <EOL> cursor . execute ( "<STR_LIT>" ) <EOL> cursor . execute ( "<STR_LIT>" % SQLITE_PAGE_SIZE ) <EOL> query = """<STR_LIT>""" % { "<STR_LIT>" : SQLITE_SUBJECT_SPEC } <EOL> cursor . execute ( query ) <EOL> query = """<STR_LIT>""" % { "<STR_LIT>" : SQLITE_SUBJECT_SPEC } <EOL> cursor . execute ( query ) <EOL> query = """<STR_LIT>""" <EOL> cursor . execute ( query ) <EOL> query = """<STR_LIT>""" <EOL> cursor . execute ( query ) <EOL> conn . commit ( ) <EOL> cursor . close ( ) <EOL> conn . close ( ) <EOL> with open ( model , "<STR_LIT:rb>" ) as model_file : <EOL> self . template = model_file . read ( ) <EOL> os . unlink ( model ) <EOL> def _WaitUntilReadable ( self , target_path ) : <EOL> start_time = time . time ( ) <EOL> while True : <EOL> if os . access ( target_path , os . R_OK ) : <EOL> return <EOL> time . sleep ( <NUM_LIT> ) <EOL> if time . time ( ) - start_time >= <NUM_LIT> : <EOL> raise IOError ( "<STR_LIT>" % target_path ) <EOL> def _EnsureDatabaseExists ( self , target_path ) : <EOL> if os . path . exists ( target_path ) : <EOL> self . _WaitUntilReadable ( target_path ) <EOL> return <EOL> umask_original = os . umask ( <NUM_LIT:0> ) <EOL> write_permissions = stat . S_IWUSR | stat . S_IWGRP <EOL> read_permissions = stat . S_IRUSR | stat . S_IRGRP <EOL> try : <EOL> fd = os . open ( target_path , os . O_CREAT | os . O_EXCL | os . O_WRONLY , <EOL> write_permissions ) <EOL> os . close ( fd ) <EOL> with open ( target_path , "<STR_LIT:wb>" ) as target_file : <EOL> target_file . write ( self . template ) <EOL> os . chmod ( target_path , write_permissions | read_permissions ) <EOL> except OSError : <EOL> if os . path . exists ( target_path ) : <EOL> self . _WaitUntilReadable ( target_path ) <EOL> else : <EOL> logging . error ( "<STR_LIT>" <EOL> "<STR_LIT>" <EOL> "<STR_LIT>" , target_path ) <EOL> finally : <EOL> os . umask ( umask_original ) <EOL> def __init__ ( self , max_size , path ) : <EOL> super ( SqliteConnectionCache , self ) . __init__ ( max_size = max_size ) <EOL> self . root_path = path or config_lib . CONFIG . Get ( "<STR_LIT>" ) <EOL> self . _CreateModelDatabase ( ) <EOL> self . RecreatePathing ( ) <EOL> def RecreatePathing ( self , pathing = None ) : <EOL> if not pathing : <EOL> pathing = config_lib . CONFIG . Get ( "<STR_LIT>" ) <EOL> try : <EOL> self . path_regexes = [ re . compile ( path ) for path in pathing ] <EOL> self . pathing = pathing <EOL> except re . error : <EOL> raise data_store . Error ( "<STR_LIT>" ) <EOL> def RootPath ( self ) : <EOL> return self . root_path <EOL> def ChangePath ( self , new_path ) : <EOL> self . root_path = new_path <EOL> def KillObject ( self , conn ) : <EOL> conn . Close ( ) <EOL> @ utils . Synchronized <EOL> def Get ( self , subject ) : <EOL> """<STR_LIT>""" <EOL> filename , directory = common . ResolveSubjectDestination ( subject , <EOL> self . path_regexes ) <EOL> key = common . MakeDestinationKey ( directory , filename ) <EOL> try : <EOL> return super ( SqliteConnectionCache , self ) . Get ( key ) <EOL> except KeyError : <EOL> dirname = utils . JoinPath ( self . root_path , directory ) <EOL> path = utils . JoinPath ( dirname , filename ) + SQLITE_EXTENSION <EOL> dirname = utils . SmartStr ( dirname ) <EOL> path = utils . SmartStr ( path ) <EOL> if not os . path . isdir ( dirname ) : <EOL> try : <EOL> os . makedirs ( dirname ) <EOL> except OSError : <EOL> pass <EOL> self . _EnsureDatabaseExists ( path ) <EOL> connection = SqliteConnection ( path ) <EOL> super ( SqliteConnectionCache , self ) . Put ( key , connection ) <EOL> return connection <EOL> def GetPrefix ( self , subject_prefix ) : <EOL> """<STR_LIT>""" <EOL> components = common . Components ( subject_prefix ) <EOL> components = [ common . ConvertStringToFilename ( x ) for x in components ] <EOL> path_prefix = utils . JoinPath ( * components ) <EOL> if path_prefix == "<STR_LIT:/>" : <EOL> path_prefix = "<STR_LIT>" <EOL> for regex in self . path_regexes : <EOL> result = common . EvaluatePrefix ( path_prefix , regex ) <EOL> if result == "<STR_LIT>" : <EOL> yield self . Get ( subject_prefix ) <EOL> return <EOL> if result == "<STR_LIT>" : <EOL> for conn in self . DatabasesByPath ( path_prefix ) : <EOL> yield conn <EOL> return <EOL> yield self . Get ( subject_prefix ) <EOL> def DatabasesInDir ( self , directory ) : <EOL> """<STR_LIT>""" <EOL> for ( path , dirs , files ) in os . walk ( directory , topdown = True ) : <EOL> dirs . sort ( ) <EOL> files . sort ( ) <EOL> for f in files : <EOL> if f . endswith ( SQLITE_EXTENSION ) : <EOL> f = f [ : - len ( SQLITE_EXTENSION ) ] <EOL> yield utils . JoinPath ( path , f ) <EOL> @ utils . Synchronized <EOL> def DatabasesByPath ( self , path_prefix ) : <EOL> """<STR_LIT>""" <EOL> dir_prefix = utils . JoinPath ( self . root_path , path_prefix ) <EOL> shortened_path_prefix = path_prefix <EOL> databases_found = set ( ) <EOL> while True : <EOL> shortened_path = utils . JoinPath ( self . root_path , shortened_path_prefix ) <EOL> if os . path . isdir ( shortened_path ) : <EOL> for db in self . DatabasesInDir ( shortened_path ) : <EOL> if db in databases_found : <EOL> continue <EOL> mod_db = db <EOL> if mod_db == utils . JoinPath ( self . root_path , "<STR_LIT>" ) : <EOL> mod_db = self . root_path <EOL> if mod_db . startswith ( dir_prefix ) or dir_prefix . startswith ( mod_db ) : <EOL> databases_found . add ( db ) <EOL> yield SqliteConnection ( db + SQLITE_EXTENSION ) <EOL> if not shortened_path_prefix : <EOL> break <EOL> components = shortened_path_prefix . split ( os . path . sep ) <EOL> shortened_path_prefix = utils . JoinPath ( * ( components [ : - <NUM_LIT:1> ] ) ) <EOL> if shortened_path_prefix == "<STR_LIT:/>" : <EOL> shortened_path_prefix = "<STR_LIT>" <EOL> class SqliteConnection ( object ) : <EOL> """<STR_LIT>""" <EOL> def __init__ ( self , filename ) : <EOL> self . filename = filename <EOL> self . conn = sqlite3 . connect ( filename , SQLITE_TIMEOUT , SQLITE_DETECT_TYPES , <EOL> SQLITE_ISOLATION , False , SQLITE_FACTORY , <EOL> SQLITE_CACHED_STATEMENTS ) <EOL> self . conn . text_factory = str <EOL> self . cursor = self . conn . cursor ( ) <EOL> self . Execute ( "<STR_LIT>" ) <EOL> self . Execute ( "<STR_LIT>" ) <EOL> self . Execute ( "<STR_LIT>" ) <EOL> self . Execute ( "<STR_LIT>" ) <EOL> self . lock = threading . RLock ( ) <EOL> self . dirty = False <EOL> self . deleted = <NUM_LIT:0> <EOL> self . next_vacuum_check = config_lib . CONFIG [ "<STR_LIT>" ] <EOL> def Filename ( self ) : <EOL> return self . filename <EOL> def Execute ( self , * args ) : <EOL> try : <EOL> return self . cursor . execute ( * args ) <EOL> except sqlite3 . DatabaseError : <EOL> logging . exception ( "<STR_LIT>" , self . filename , <EOL> args ) <EOL> raise <EOL> @ utils . Synchronized <EOL> def GetLock ( self , subject ) : <EOL> """<STR_LIT>""" <EOL> subject = utils . SmartStr ( subject ) <EOL> query = "<STR_LIT>" <EOL> args = ( subject , ) <EOL> data = self . Execute ( query , args ) . fetchone ( ) <EOL> if data : <EOL> return data [ <NUM_LIT:0> ] , data [ <NUM_LIT:1> ] <EOL> else : <EOL> return None , None <EOL> @ utils . Synchronized <EOL> def SetLock ( self , subject , expires , token ) : <EOL> """<STR_LIT>""" <EOL> subject = utils . SmartStr ( subject ) <EOL> query = "<STR_LIT>" <EOL> args = ( subject , expires , token ) <EOL> self . Execute ( query , args ) <EOL> self . dirty = True <EOL> @ utils . Synchronized <EOL> def RemoveLock ( self , subject ) : <EOL> """<STR_LIT>""" <EOL> subject = utils . SmartStr ( subject ) <EOL> query = "<STR_LIT>" <EOL> args = ( subject , ) <EOL> self . Execute ( query , args ) <EOL> self . dirty = True <EOL> @ utils . Synchronized <EOL> def GetNewestValue ( self , subject , attribute ) : <EOL> """<STR_LIT>""" <EOL> subject = utils . SmartStr ( subject ) <EOL> attribute = utils . SmartStr ( attribute ) <EOL> query = """<STR_LIT>""" <EOL> args = ( subject , attribute ) <EOL> data = self . Execute ( query , args ) . fetchone ( ) <EOL> if data : <EOL> return ( data [ <NUM_LIT:0> ] , data [ <NUM_LIT:1> ] ) <EOL> else : <EOL> return None <EOL> @ utils . Synchronized <EOL> def GetNewestFromPrefix ( self , subject , prefix , limit = None ) : <EOL> """<STR_LIT>""" <EOL> pattern = prefix + "<STR_LIT:%>" <EOL> subject = utils . SmartStr ( subject ) <EOL> query = """<STR_LIT>""" <EOL> if limit : <EOL> query += "<STR_LIT>" <EOL> args = ( subject , pattern , limit ) <EOL> else : <EOL> args = ( subject , pattern ) <EOL> data = self . Execute ( query , args ) . fetchall ( ) <EOL> return [ ( pred , val , ts ) for pred , ts , val in data ] <EOL> @ utils . Synchronized <EOL> def GetValuesFromPrefix ( self , subject , prefix , start , end , limit = None ) : <EOL> """<STR_LIT>""" <EOL> pattern = prefix + "<STR_LIT:%>" <EOL> subject = utils . SmartStr ( subject ) <EOL> query = """<STR_LIT>""" <EOL> if limit : <EOL> query += "<STR_LIT>" <EOL> args = ( subject , pattern , start , end , limit ) <EOL> else : <EOL> args = ( subject , pattern , start , end ) <EOL> data = self . Execute ( query , args ) . fetchall ( ) <EOL> return data <EOL> @ utils . Synchronized <EOL> def GetValues ( self , subject , attribute , start , end , limit = None ) : <EOL> """<STR_LIT>""" <EOL> subject = utils . SmartStr ( subject ) <EOL> attribute = utils . SmartStr ( attribute ) <EOL> query = """<STR_LIT>""" <EOL> if limit : <EOL> query += "<STR_LIT>" <EOL> args = ( subject , attribute , start , end , limit ) <EOL> else : <EOL> args = ( subject , attribute , start , end ) <EOL> data = self . Execute ( query , args ) . fetchall ( ) <EOL> return data <EOL> def ScanAttributes ( self , <EOL> subject_prefix , <EOL> attributes , <EOL> after_urn = None , <EOL> max_records = None ) : <EOL> """<STR_LIT>""" <EOL> cursor = self . conn . cursor ( ) <EOL> cursor . execute ( "<STR_LIT>" ) <EOL> cursor . execute ( "<STR_LIT>" ) <EOL> query = """<STR_LIT>""" % "<STR_LIT:U+002C>" . join ( "<STR_LIT:?>" * len ( attributes ) ) <EOL> subject_prefix = utils . SmartStr ( subject_prefix ) <EOL> if after_urn : <EOL> after_urn = utils . SmartStr ( after_urn ) <EOL> else : <EOL> after_urn = "<STR_LIT>" <EOL> args = [ subject_prefix + "<STR_LIT:%>" , after_urn ] + attributes <EOL> if max_records : <EOL> query += "<STR_LIT>" <EOL> args . append ( max_records * len ( attributes ) ) <EOL> cursor . execute ( query , args ) <EOL> for r in cursor : <EOL> yield r <EOL> @ utils . Synchronized <EOL> def DeleteAttribute ( self , subject , attribute ) : <EOL> """<STR_LIT>""" <EOL> subject = utils . SmartStr ( subject ) <EOL> attribute = utils . SmartStr ( attribute ) <EOL> query = "<STR_LIT>" <EOL> args = ( subject , attribute ) <EOL> self . Execute ( query , args ) <EOL> self . dirty = True <EOL> self . deleted += self . cursor . rowcount <EOL> @ utils . Synchronized <EOL> def SetAttribute ( self , subject , attribute , value , timestamp ) : <EOL> """<STR_LIT>""" <EOL> subject = utils . SmartStr ( subject ) <EOL> attribute = utils . SmartStr ( attribute ) <EOL> query = "<STR_LIT>" <EOL> args = ( subject , attribute , timestamp , value ) <EOL> self . Execute ( query , args ) <EOL> self . dirty = True <EOL> self . deleted = max ( <NUM_LIT:0> , self . deleted - self . cursor . rowcount ) <EOL> @ utils . Synchronized <EOL> def DeleteAttributeRange ( self , subject , attribute , start , end ) : <EOL> """<STR_LIT>""" <EOL> subject = utils . SmartStr ( subject ) <EOL> attribute = utils . SmartStr ( attribute ) <EOL> query = """<STR_LIT>""" <EOL> args = ( subject , attribute , int ( start ) , int ( end ) ) <EOL> self . Execute ( query , args ) <EOL> self . dirty = True <EOL> self . deleted += self . cursor . rowcount <EOL> @ utils . Synchronized <EOL> def DeleteSubject ( self , subject ) : <EOL> """<STR_LIT>""" <EOL> subject = utils . SmartStr ( subject ) <EOL> query = "<STR_LIT>" <EOL> args = ( subject , ) <EOL> self . Execute ( query , args ) <EOL> self . dirty = True <EOL> self . deleted += self . cursor . rowcount <EOL> def PrettyPrint ( self ) : <EOL> """<STR_LIT>""" <EOL> query = "<STR_LIT>" <EOL> for sub , pred , ts , val in self . Execute ( query ) : <EOL> print "<STR_LIT>" % ( sub , pred , ts , val ) <EOL> print "<STR_LIT>" <EOL> def __enter__ ( self ) : <EOL> self . lock . acquire ( ) <EOL> return self <EOL> def __exit__ ( self , exc_type , exc_value , traceback ) : <EOL> if self . dirty : <EOL> self . Flush ( ) <EOL> self . dirty = False <EOL> self . lock . release ( ) <EOL> @ utils . Synchronized <EOL> def Flush ( self ) : <EOL> """<STR_LIT>""" <EOL> if self . conn : <EOL> try : <EOL> self . conn . commit ( ) <EOL> except sqlite3 . OperationalError : <EOL> pass <EOL> if self . deleted >= self . next_vacuum_check : <EOL> if self . _NeedsVacuum ( ) and not self . _HasRecentVacuum ( ) : <EOL> self . Vacuum ( ) <EOL> self . deleted = <NUM_LIT:0> <EOL> self . next_vacuum_check = max ( <EOL> config_lib . CONFIG [ "<STR_LIT>" ] , <EOL> self . next_vacuum_check / <NUM_LIT:2> ) <EOL> else : <EOL> self . next_vacuum_check *= <NUM_LIT:2> <EOL> def _NeedsVacuum ( self ) : <EOL> """<STR_LIT>""" <EOL> pages_result = self . Execute ( "<STR_LIT>" ) . fetchone ( ) <EOL> if not pages_result : <EOL> return False <EOL> pages = int ( pages_result [ <NUM_LIT:0> ] ) <EOL> vacuum_minsize = config_lib . CONFIG [ "<STR_LIT>" ] <EOL> if pages * SQLITE_PAGE_SIZE < vacuum_minsize : <EOL> return False <EOL> free_pages_result = self . Execute ( "<STR_LIT>" ) . fetchone ( ) <EOL> if not free_pages_result : <EOL> return False <EOL> free_pages = int ( free_pages_result [ <NUM_LIT:0> ] ) <EOL> vacuum_ratio = config_lib . CONFIG [ "<STR_LIT>" ] <EOL> return <NUM_LIT> * float ( free_pages ) / float ( pages ) >= vacuum_ratio <EOL> def _HasRecentVacuum ( self ) : <EOL> """<STR_LIT>""" <EOL> query = "<STR_LIT>" <EOL> data = self . Execute ( query ) . fetchone ( ) <EOL> if not data : <EOL> return False <EOL> try : <EOL> last_vacuum = int ( str ( data [ <NUM_LIT:0> ] ) ) <EOL> vacuum_frequency = config_lib . CONFIG [ "<STR_LIT>" ] <EOL> return time . time ( ) - last_vacuum < vacuum_frequency <EOL> except ValueError : <EOL> return False <EOL> def Vacuum ( self ) : <EOL> """<STR_LIT>""" <EOL> now = time . time ( ) <EOL> logging . debug ( "<STR_LIT>" ) <EOL> self . Execute ( "<STR_LIT>" ) <EOL> query = "<STR_LIT>" <EOL> args = ( str ( int ( now ) ) , ) <EOL> self . Execute ( query , args ) <EOL> try : <EOL> self . conn . commit ( ) <EOL> except sqlite3 . OperationalError : <EOL> pass <EOL> @ utils . Synchronized <EOL> def Close ( self ) : <EOL> """<STR_LIT>""" <EOL> if self . dirty : <EOL> self . Flush ( ) <EOL> self . cursor . close ( ) <EOL> self . conn . close ( ) <EOL> self . conn = None <EOL> self . cursor = None <EOL> class SqliteDataStore ( data_store . DataStore ) : <EOL> """<STR_LIT>""" <EOL> cache = None <EOL> def __init__ ( self , path = None ) : <EOL> self . _CalculateAttributeStorageTypes ( ) <EOL> super ( SqliteDataStore , self ) . __init__ ( ) <EOL> self . cache = SqliteConnectionCache ( <EOL> config_lib . CONFIG [ "<STR_LIT>" ] , path ) <EOL> def RecreatePathing ( self , pathing ) : <EOL> self . cache . RecreatePathing ( pathing ) <EOL> def _CalculateAttributeStorageTypes ( self ) : <EOL> """<STR_LIT>""" <EOL> self . _attribute_types = { } <EOL> for attribute in aff4 . Attribute . PREDICATES . values ( ) : <EOL> self . _attribute_types [ attribute . predicate ] = ( <EOL> attribute . attribute_type . data_store_type ) <EOL> def _Encode ( self , value ) : <EOL> """<STR_LIT>""" <EOL> try : <EOL> return buffer ( value . SerializeToString ( ) ) <EOL> except AttributeError : <EOL> if isinstance ( value , ( int , long ) ) : <EOL> return value <EOL> else : <EOL> return buffer ( utils . SmartStr ( value ) ) <EOL> def _Decode ( self , attribute , value ) : <EOL> required_type = self . _attribute_types . get ( attribute , "<STR_LIT>" ) <EOL> if isinstance ( value , buffer ) : <EOL> value = str ( value ) <EOL> if required_type in ( "<STR_LIT>" , "<STR_LIT>" ) : <EOL> return int ( value ) <EOL> elif required_type == "<STR_LIT:string>" : <EOL> return utils . SmartUnicode ( value ) <EOL> else : <EOL> return value <EOL> def MultiSet ( self , subject , values , timestamp = None , replace = True , <EOL> sync = True , to_delete = None , token = None ) : <EOL> """<STR_LIT>""" <EOL> self . security_manager . CheckDataStoreAccess ( token , [ subject ] , "<STR_LIT:w>" ) <EOL> _ = sync <EOL> if timestamp is None or timestamp == self . NEWEST_TIMESTAMP : <EOL> timestamp = time . time ( ) * <NUM_LIT> <EOL> if to_delete is None : <EOL> to_delete = [ ] <EOL> with self . cache . Get ( subject ) as sqlite_connection : <EOL> if replace : <EOL> to_delete . extend ( values . keys ( ) ) <EOL> if to_delete : <EOL> for attribute in to_delete : <EOL> sqlite_connection . DeleteAttribute ( subject , attribute ) <EOL> for attribute , seq in values . items ( ) : <EOL> for v in seq : <EOL> element_timestamp = None <EOL> if isinstance ( v , ( list , tuple ) ) : <EOL> v , element_timestamp = v <EOL> if element_timestamp is None : <EOL> element_timestamp = timestamp <EOL> element_timestamp = long ( element_timestamp ) <EOL> value = self . _Encode ( v ) <EOL> sqlite_connection . SetAttribute ( subject , attribute , value , <EOL> element_timestamp ) <EOL> def DeleteAttributes ( self , subject , attributes , start = None , end = None , <EOL> sync = True , token = None ) : <EOL> """<STR_LIT>""" <EOL> self . security_manager . CheckDataStoreAccess ( token , [ subject ] , "<STR_LIT:w>" ) <EOL> _ = sync <EOL> if isinstance ( attributes , basestring ) : <EOL> raise ValueError ( <EOL> "<STR_LIT>" ) <EOL> with self . cache . Get ( subject ) as sqlite_connection : <EOL> if start is None and end is None : <EOL> for attribute in list ( attributes ) : <EOL> sqlite_connection . DeleteAttribute ( subject , attribute ) <EOL> else : <EOL> start = start or <NUM_LIT:0> <EOL> if end is None : <EOL> end = ( <NUM_LIT:2> ** <NUM_LIT> ) - <NUM_LIT:1> <EOL> for attribute in list ( attributes ) : <EOL> sqlite_connection . DeleteAttributeRange ( subject , attribute , start , <EOL> end ) <EOL> def DeleteSubject ( self , subject , sync = False , token = None ) : <EOL> _ = sync <EOL> self . security_manager . CheckDataStoreAccess ( token , [ subject ] , "<STR_LIT:w>" ) <EOL> with self . cache . Get ( subject ) as sqlite_connection : <EOL> sqlite_connection . DeleteSubject ( subject ) <EOL> def MultiResolvePrefix ( self , subjects , attribute_prefix , timestamp = None , <EOL> limit = None , token = None ) : <EOL> """<STR_LIT>""" <EOL> result = { } <EOL> remaining_limit = limit <EOL> for subject in subjects : <EOL> values = self . ResolvePrefix ( subject , attribute_prefix , token = token , <EOL> timestamp = timestamp , limit = remaining_limit ) <EOL> if values : <EOL> if limit : <EOL> if len ( values ) >= remaining_limit : <EOL> result [ subject ] = values [ : remaining_limit ] <EOL> return result . iteritems ( ) <EOL> remaining_limit -= len ( values ) <EOL> result [ subject ] = values <EOL> return result . iteritems ( ) <EOL> def _GetStartEndTimestamp ( self , timestamp ) : <EOL> if timestamp == self . ALL_TIMESTAMPS or timestamp is None : <EOL> return <NUM_LIT:0> , ( <NUM_LIT:2> ** <NUM_LIT> ) - <NUM_LIT:1> <EOL> elif timestamp == self . NEWEST_TIMESTAMP : <EOL> return - <NUM_LIT:1> , - <NUM_LIT:1> <EOL> elif isinstance ( timestamp , int ) : <EOL> return timestamp , timestamp <EOL> else : <EOL> try : <EOL> start , end = timestamp <EOL> return int ( start ) , int ( end ) <EOL> except ValueError : <EOL> return timestamp , timestamp <EOL> def ResolvePrefix ( self , subject , attribute_prefix , timestamp = None , <EOL> limit = None , token = None ) : <EOL> """<STR_LIT>""" <EOL> self . security_manager . CheckDataStoreAccess ( <EOL> token , [ subject ] , self . GetRequiredResolveAccess ( attribute_prefix ) ) <EOL> if isinstance ( attribute_prefix , str ) : <EOL> attribute_prefix = [ attribute_prefix ] <EOL> start , end = self . _GetStartEndTimestamp ( timestamp ) <EOL> results = [ ] <EOL> with self . cache . Get ( subject ) as sqlite_connection : <EOL> for prefix in attribute_prefix : <EOL> nr_results = len ( results ) <EOL> if limit and nr_results >= limit : <EOL> break <EOL> new_limit = limit <EOL> if new_limit : <EOL> new_limit -= nr_results <EOL> if timestamp == self . NEWEST_TIMESTAMP : <EOL> data = sqlite_connection . GetNewestFromPrefix ( subject , prefix , <EOL> new_limit ) <EOL> for attribute , value , ts in data : <EOL> value = self . _Decode ( attribute , value ) <EOL> results . append ( ( attribute , value , ts ) ) <EOL> else : <EOL> data = sqlite_connection . GetValuesFromPrefix ( subject , prefix , start , <EOL> end , new_limit ) <EOL> for attribute , value , ts in data : <EOL> value = self . _Decode ( attribute , value ) <EOL> results . append ( ( attribute , value , ts ) ) <EOL> return results <EOL> def _GroupSubjects ( self , collection , max_records ) : <EOL> """<STR_LIT>""" <EOL> record_count = <NUM_LIT:0> <EOL> current_subject = None <EOL> current_results = { } <EOL> for subject , attribute , timestamp , value in collection : <EOL> if not current_subject : <EOL> current_subject = subject <EOL> if current_subject != subject : <EOL> yield ( current_subject , current_results ) <EOL> record_count += <NUM_LIT:1> <EOL> if max_records and record_count >= max_records : <EOL> return <EOL> current_results = { } <EOL> current_subject = subject <EOL> current_results [ attribute ] = ( timestamp , self . _Decode ( attribute , value ) ) <EOL> if current_results : <EOL> yield ( current_subject , current_results ) <EOL> def ScanAttributes ( self , <EOL> subject_prefix , <EOL> attributes , <EOL> after_urn = None , <EOL> max_records = None , <EOL> token = None , <EOL> relaxed_order = False ) : <EOL> self . security_manager . CheckDataStoreAccess ( token , [ subject_prefix ] , "<STR_LIT>" ) <EOL> subject_prefix = utils . SmartStr ( rdfvalue . RDFURN ( subject_prefix ) ) <EOL> if subject_prefix [ - <NUM_LIT:1> ] != "<STR_LIT:/>" : <EOL> subject_prefix += "<STR_LIT:/>" <EOL> if after_urn : <EOL> after_urn = str ( after_urn ) <EOL> connection_iter = self . cache . GetPrefix ( subject_prefix ) <EOL> if relaxed_order : <EOL> for sqlite_connection in connection_iter : <EOL> with sqlite_connection : <EOL> for r in self . _GroupSubjects ( list ( sqlite_connection . ScanAttributes ( <EOL> subject_prefix , <EOL> attributes , <EOL> after_urn = after_urn , <EOL> max_records = max_records ) ) , max_records ) : <EOL> yield r <EOL> return <EOL> first_connections = [ ] <EOL> try : <EOL> first_connections . append ( connection_iter . next ( ) ) <EOL> first_connections . append ( connection_iter . next ( ) ) <EOL> except StopIteration : <EOL> pass <EOL> if not first_connections : <EOL> return <EOL> if len ( first_connections ) == <NUM_LIT:1> : <EOL> with first_connections [ <NUM_LIT:0> ] as sqlite_connection : <EOL> for r in self . _GroupSubjects ( list ( sqlite_connection . ScanAttributes ( <EOL> subject_prefix , <EOL> attributes , <EOL> after_urn = after_urn , <EOL> max_records = max_records ) ) , max_records ) : <EOL> yield r <EOL> return <EOL> raw_results = [ ] <EOL> for sqlite_connection in itertools . chain ( first_connections , <EOL> connection_iter ) : <EOL> for record in sqlite_connection . ScanAttributes ( subject_prefix , <EOL> attributes , <EOL> after_urn = after_urn , <EOL> max_records = max_records ) : <EOL> raw_results . append ( record ) <EOL> for r in self . _GroupSubjects ( sorted ( raw_results , key = lambda x : x [ <NUM_LIT:0> ] ) , <EOL> max_records ) : <EOL> yield r <EOL> def ResolveMulti ( self , subject , attributes , timestamp = None , <EOL> limit = None , token = None ) : <EOL> """<STR_LIT>""" <EOL> self . security_manager . CheckDataStoreAccess ( <EOL> token , [ subject ] , self . GetRequiredResolveAccess ( attributes ) ) <EOL> results = [ ] <EOL> start , end = self . _GetStartEndTimestamp ( timestamp ) <EOL> with self . cache . Get ( subject ) as sqlite_connection : <EOL> for attribute in attributes : <EOL> if timestamp == self . NEWEST_TIMESTAMP : <EOL> ret = sqlite_connection . GetNewestValue ( subject , attribute ) <EOL> if ret : <EOL> value , ts = ret <EOL> value = self . _Decode ( attribute , value ) <EOL> results . append ( ( attribute , value , ts ) ) <EOL> if limit and len ( results ) >= limit : <EOL> break <EOL> else : <EOL> new_limit = limit <EOL> if new_limit : <EOL> new_limit = limit - len ( results ) <EOL> values = sqlite_connection . GetValues ( subject , attribute , start , end , <EOL> new_limit ) <EOL> for value , ts in values : <EOL> value = self . _Decode ( attribute , value ) <EOL> results . append ( ( attribute , value , ts ) ) <EOL> if limit and len ( results ) >= limit : <EOL> break <EOL> return results <EOL> def DumpDatabase ( self , token = None ) : <EOL> self . security_manager . CheckDataStoreAccess ( token , [ ] , "<STR_LIT:r>" ) <EOL> for _ , sql_connection in self . cache : <EOL> sql_connection . PrettyPrint ( ) <EOL> def Size ( self ) : <EOL> root_path = self . Location ( ) <EOL> if not os . path . exists ( root_path ) : <EOL> return <NUM_LIT:0> <EOL> if not os . path . isdir ( root_path ) : <EOL> raise IOError ( "<STR_LIT>" % <EOL> root_path ) <EOL> size , _ = common . DatabaseDirectorySize ( root_path , self . FileExtension ( ) ) <EOL> return size <EOL> @ staticmethod <EOL> def FileExtension ( ) : <EOL> return SQLITE_EXTENSION <EOL> def Location ( self ) : <EOL> """<STR_LIT>""" <EOL> return self . cache . RootPath ( ) <EOL> def ChangeLocation ( self , location ) : <EOL> self . cache . ChangePath ( location ) <EOL> def Transaction ( self , subject , lease_time = None , token = None ) : <EOL> return SqliteTransaction ( self , subject , lease_time = lease_time , token = token ) <EOL> class SqliteTransaction ( data_store . CommonTransaction ) : <EOL> """<STR_LIT>""" <EOL> lock_creation_lock = threading . Lock ( ) <EOL> locked = False <EOL> def __init__ ( self , store , subject , lease_time = None , token = None ) : <EOL> """<STR_LIT>""" <EOL> super ( SqliteTransaction , self ) . __init__ ( store , utils . SmartUnicode ( subject ) , <EOL> lease_time = lease_time , token = token ) <EOL> if lease_time is None : <EOL> lease_time = config_lib . CONFIG [ "<STR_LIT>" ] <EOL> self . lock_token = thread . get_ident ( ) <EOL> sqlite_connection = store . cache . Get ( self . subject ) <EOL> with sqlite_connection : <EOL> locked_until , stored_token = sqlite_connection . GetLock ( subject ) <EOL> if locked_until and time . time ( ) < float ( locked_until ) : <EOL> raise data_store . TransactionError ( "<STR_LIT>" % subject ) <EOL> self . expires = time . time ( ) + lease_time <EOL> sqlite_connection . SetLock ( subject , self . expires , self . lock_token ) <EOL> locked_until , stored_token = sqlite_connection . GetLock ( subject ) <EOL> if stored_token != self . lock_token : <EOL> raise data_store . TransactionError ( "<STR_LIT>" % subject ) <EOL> self . locked = True <EOL> def UpdateLease ( self , duration ) : <EOL> self . expires = time . time ( ) + duration <EOL> with self . store . cache . Get ( self . subject ) as sqlite_connection : <EOL> self . expires = time . time ( ) + duration <EOL> sqlite_connection . SetLock ( self . subject , self . expires , self . lock_token ) <EOL> def Abort ( self ) : <EOL> if self . locked : <EOL> self . _RemoveLock ( ) <EOL> def Commit ( self ) : <EOL> if self . locked : <EOL> super ( SqliteTransaction , self ) . Commit ( ) <EOL> self . _RemoveLock ( ) <EOL> def _RemoveLock ( self ) : <EOL> with self . store . cache . Get ( self . subject ) as sqlite_connection : <EOL> sqlite_connection . RemoveLock ( self . subject ) <EOL> self . locked = False </s>
