hydrus/hydrus/core/HydrusDBBase.py

import collections
import psutil
import sqlite3

from hydrus.core import HydrusData
from hydrus.core import HydrusPaths
from hydrus.core import HydrusGlobals as HG
from hydrus.core import HydrusTemp

def CheckHasSpaceForDBTransaction( db_dir, num_bytes ):
    
    if HG.no_db_temp_files:
        
        space_needed = int( num_bytes * 1.1 )
        
        approx_available_memory = psutil.virtual_memory().available * 4 / 5
        
        if approx_available_memory < num_bytes:
            
            raise Exception( 'I believe you need about ' + HydrusData.ToHumanBytes( space_needed ) + ' available memory, since you are running in no_db_temp_files mode, but you only seem to have ' + HydrusData.ToHumanBytes( approx_available_memory ) + '.' )
            
        
        db_disk_free_space = HydrusPaths.GetFreeSpace( db_dir )
        
        if db_disk_free_space < space_needed:
            
            raise Exception( 'I believe you need about ' + HydrusData.ToHumanBytes( space_needed ) + ' on your db\'s disk partition, but you only seem to have ' + HydrusData.ToHumanBytes( db_disk_free_space ) + '.' )
            
        
    else:
        
        temp_dir = HydrusTemp.GetCurrentTempDir()
        
        temp_disk_free_space = HydrusPaths.GetFreeSpace( temp_dir )
        
        temp_and_db_on_same_device = HydrusPaths.GetDevice( temp_dir ) == HydrusPaths.GetDevice( db_dir )
        
        if temp_and_db_on_same_device:
            
            space_needed = int( num_bytes * 2.2 )
            
            if temp_disk_free_space < space_needed:
                
                raise Exception( 'I believe you need about ' + HydrusData.ToHumanBytes( space_needed ) + ' on your db\'s disk partition, which I think also holds your temporary path, but you only seem to have ' + HydrusData.ToHumanBytes( temp_disk_free_space ) + '.' )
                
            
        else:
            
            space_needed = int( num_bytes * 1.1 )
            
            if temp_disk_free_space < space_needed:
                
                raise Exception( 'I believe you need about ' + HydrusData.ToHumanBytes( space_needed ) + ' on your temporary path\'s disk partition, which I think is ' + temp_dir + ', but you only seem to have ' + HydrusData.ToHumanBytes( temp_disk_free_space ) + '.' )
                
            
            db_disk_free_space = HydrusPaths.GetFreeSpace( db_dir )
            
            if db_disk_free_space < space_needed:
                
                raise Exception( 'I believe you need about ' + HydrusData.ToHumanBytes( space_needed ) + ' on your db\'s disk partition, but you only seem to have ' + HydrusData.ToHumanBytes( db_disk_free_space ) + '.' )
                
            
        
    
class TemporaryIntegerTableNameCache( object ):
    
    my_instance = None
    
    def __init__( self ):
        
        TemporaryIntegerTableNameCache.my_instance = self
        
        self._column_names_to_table_names = collections.defaultdict( collections.deque )
        self._column_names_counter = collections.Counter()
        
    
    @staticmethod
    def instance() -> 'TemporaryIntegerTableNameCache':
        
        if TemporaryIntegerTableNameCache.my_instance is None:
            
            raise Exception( 'TemporaryIntegerTableNameCache is not yet initialised!' )
            
        else:
            
            return TemporaryIntegerTableNameCache.my_instance
            
        
    
    def Clear( self ):
        
        self._column_names_to_table_names = collections.defaultdict( collections.deque )
        self._column_names_counter = collections.Counter()
        
    
    def GetName( self, column_name ):
        
        table_names = self._column_names_to_table_names[ column_name ]
        
        initialised = True
        
        if len( table_names ) == 0:
            
            initialised = False
            
            i = self._column_names_counter[ column_name ]
            
            table_name = 'mem.temp_int_{}_{}'.format( column_name, i )
            
            table_names.append( table_name )
            
            self._column_names_counter[ column_name ] += 1
            
        
        table_name = table_names.pop()
        
        return ( initialised, table_name )
        
    
    def ReleaseName( self, column_name, table_name ):
        
        self._column_names_to_table_names[ column_name ].append( table_name )
        
    
class TemporaryIntegerTable( object ):
    
    def __init__( self, cursor: sqlite3.Cursor, integer_iterable, column_name ):
        
        if not isinstance( integer_iterable, set ):
            
            integer_iterable = set( integer_iterable )
            
        
        self._cursor = cursor
        self._integer_iterable = integer_iterable
        self._column_name = column_name
        
        ( self._initialised, self._table_name ) = TemporaryIntegerTableNameCache.instance().GetName( self._column_name )
        
    
    def __enter__( self ):
        
        if not self._initialised:
            
            self._cursor.execute( 'CREATE TABLE IF NOT EXISTS {} ( {} INTEGER PRIMARY KEY );'.format( self._table_name, self._column_name ) )
            
        
        self._cursor.executemany( 'INSERT INTO {} ( {} ) VALUES ( ? );'.format( self._table_name, self._column_name ), ( ( i, ) for i in self._integer_iterable ) )
        
        return self._table_name
        
    
    def __exit__( self, exc_type, exc_val, exc_tb ):
        
        self._cursor.execute( 'DELETE FROM {};'.format( self._table_name ) )
        
        TemporaryIntegerTableNameCache.instance().ReleaseName( self._column_name, self._table_name )
        
        return False
        
    
class DBBase( object ):
    
    def __init__( self ):
        
        self._c = None
        
    
    def _CloseCursor( self ):
        
        if self._c is not None:
            
            self._c.close()
            
            del self._c
            
            self._c = None
            
        
    
    def _CreateIndex( self, table_name, columns, unique = False ):
        
        if unique:
            
            create_phrase = 'CREATE UNIQUE INDEX IF NOT EXISTS'
            
        else:
            
            create_phrase = 'CREATE INDEX IF NOT EXISTS'
            
        
        index_name = self._GenerateIndexName( table_name, columns )
        
        if '.' in table_name:
            
            table_name_simple = table_name.split( '.' )[1]
            
        else:
            
            table_name_simple = table_name
            
        
        statement = '{} {} ON {} ({});'.format( create_phrase, index_name, table_name_simple, ', '.join( columns ) )
        
        self._Execute( statement )
        
    
    def _Execute( self, query, *args ) -> sqlite3.Cursor:
        
        if HG.query_planner_mode and query not in HG.queries_planned:
            
            plan_lines = self._c.execute( 'EXPLAIN QUERY PLAN {}'.format( query ), *args ).fetchall()
            
            HG.query_planner_query_count += 1
            
            HG.controller.PrintQueryPlan( query, plan_lines )
            
        
        return self._c.execute( query, *args )
        
    
    def _ExecuteMany( self, query, args_iterator ):
        
        if HG.query_planner_mode and query not in HG.queries_planned:
            
            args_iterator = list( args_iterator )
            
            if len( args_iterator ) > 0:
                
                plan_lines = self._c.execute( 'EXPLAIN QUERY PLAN {}'.format( query ), args_iterator[0] ).fetchall()
                
                HG.query_planner_query_count += 1
                
                HG.controller.PrintQueryPlan( query, plan_lines )
                
            
        
        self._c.executemany( query, args_iterator )
        
    
    def _GenerateIndexName( self, table_name, columns ):
        
        return '{}_{}_index'.format( table_name, '_'.join( columns ) )
        
    
    def _GetAttachedDatabaseNames( self, include_temp = False ):
        
        if include_temp:
            
            f = lambda schema_name, path: True
            
        else:
            
            f = lambda schema_name, path: schema_name != 'temp' and path != ''
            
        
        names = [ schema_name for ( number, schema_name, path ) in self._Execute( 'PRAGMA database_list;' ) if f( schema_name, path ) ]
        
        return names
        
    
    def _GetLastRowId( self ) -> int:
        
        return self._c.lastrowid
        
    
    def _GetRowCount( self ):
        
        row_count = self._c.rowcount
        
        if row_count == -1:
            
            return 0
            
        else:
            
            return row_count
            
        
    
    def _IndexExists( self, table_name, columns ):
        
        index_name = self._GenerateIndexName( table_name, columns )
        
        return self._TableOrIndexExists( index_name, 'index' )
        
    
    def _MakeTemporaryIntegerTable( self, integer_iterable, column_name ):
        
        return TemporaryIntegerTable( self._c, integer_iterable, column_name )
        
    
    def _SetCursor( self, c: sqlite3.Cursor ):
        
        self._c = c
        
    
    def _STI( self, iterable_cursor ):
        
        # strip singleton tuples to an iterator
        
        return ( item for ( item, ) in iterable_cursor )
        
    
    def _STL( self, iterable_cursor ):
        
        # strip singleton tuples to a list
        
        return [ item for ( item, ) in iterable_cursor ]
        
    
    def _STS( self, iterable_cursor ):
        
        # strip singleton tuples to a set
        
        return { item for ( item, ) in iterable_cursor }
        
    
    def _TableExists( self, table_name ):
        
        return self._TableOrIndexExists( table_name, 'table' )
        
    
    def _TableOrIndexExists( self, name, item_type ):
        
        if '.' in name:
            
            ( schema, name ) = name.split( '.', 1 )
            
            search_schemas = [ schema ]
            
        else:
            
            search_schemas = self._GetAttachedDatabaseNames()
            
        
        for schema in search_schemas:
            
            result = self._Execute( 'SELECT 1 FROM {}.sqlite_master WHERE name = ? AND type = ?;'.format( schema ), ( name, item_type ) ).fetchone()
            
            if result is not None:
                
                return True
                
            
        
        return False
        
    
class DBCursorTransactionWrapper( DBBase ):
    
    def __init__( self, c: sqlite3.Cursor, transaction_commit_period: int ):
        
        DBBase.__init__( self )
        
        self._SetCursor( c )
        
        self._transaction_commit_period = transaction_commit_period
        
        self._transaction_start_time = 0
        self._in_transaction = False
        self._transaction_contains_writes = False
        
        self._last_mem_refresh_time = HydrusData.GetNow()
        self._last_wal_checkpoint_time = HydrusData.GetNow()
        
        self._pubsubs = []
        
    
    def BeginImmediate( self ):
        
        if not self._in_transaction:
            
            self._Execute( 'BEGIN IMMEDIATE;' )
            self._Execute( 'SAVEPOINT hydrus_savepoint;' )
            
            self._transaction_start_time = HydrusData.GetNow()
            self._in_transaction = True
            self._transaction_contains_writes = False
            
        
    
    def CleanPubSubs( self ):
        
        self._pubsubs = []
        
    
    def Commit( self ):
        
        if self._in_transaction:
            
            self.DoPubSubs()
            
            self.CleanPubSubs()
            
            self._Execute( 'COMMIT;' )
            
            self._in_transaction = False
            self._transaction_contains_writes = False
            
            if HG.db_journal_mode == 'WAL' and HydrusData.TimeHasPassed( self._last_wal_checkpoint_time + 1800 ):
                
                self._Execute( 'PRAGMA wal_checkpoint(PASSIVE);' )
                
                self._last_wal_checkpoint_time = HydrusData.GetNow()
                
            
            if HydrusData.TimeHasPassed( self._last_mem_refresh_time + 600 ):
                
                self._Execute( 'DETACH mem;' )
                self._Execute( 'ATTACH ":memory:" AS mem;' )
                
                TemporaryIntegerTableNameCache.instance().Clear()
                
                self._last_mem_refresh_time = HydrusData.GetNow()
                
            
        else:
            
            HydrusData.Print( 'Received a call to commit, but was not in a transaction!' )
            
        
    
    def CommitAndBegin( self ):
        
        if self._in_transaction:
            
            self.Commit()
            
            self.BeginImmediate()
            
        
    
    def DoPubSubs( self ):
        
        for ( topic, args, kwargs ) in self._pubsubs:
            
            HG.controller.pub( topic, *args, **kwargs )
            
        
    
    def InTransaction( self ):
        
        return self._in_transaction
        
    
    def NotifyWriteOccuring( self ):
        
        self._transaction_contains_writes = True
        
    
    def pub_after_job( self, topic, *args, **kwargs ):
        
        if len( args ) == 0 and len( kwargs ) == 0:
            
            if ( topic, args, kwargs ) in self._pubsubs:
                
                return
                
            
        
        self._pubsubs.append( ( topic, args, kwargs ) )
        
    
    def Rollback( self ):
        
        if self._in_transaction:
            
            self._Execute( 'ROLLBACK TO hydrus_savepoint;' )
            
            # any temp int tables created in this lad will be rolled back, so 'initialised' can't be trusted. just reset, no big deal
            TemporaryIntegerTableNameCache.instance().Clear()
            
            # still in transaction
            # transaction may no longer contain writes, but it isn't important to figure out that it doesn't
            
        else:
            
            HydrusData.Print( 'Received a call to rollback, but was not in a transaction!' )
            
        
    
    def Save( self ):
        
        if self._in_transaction:
            
            try:
                
                self._Execute( 'RELEASE hydrus_savepoint;' )
                
            except sqlite3.OperationalError:
                
                HydrusData.Print( 'Tried to release a database savepoint, but failed!' )
                
            
            self._Execute( 'SAVEPOINT hydrus_savepoint;' )
            
        else:
            
            HydrusData.Print( 'Received a call to save, but was not in a transaction!' )
            
        
    
    def TimeToCommit( self ):
        
        return self._in_transaction and self._transaction_contains_writes and HydrusData.TimeHasPassed( self._transaction_start_time + self._transaction_commit_period )