<commit_before># -*- coding: utf-8 -*-
"""Database related tools."""
import logging

from datetime import datetime

import dateutil.parser

from sqlalchemy import (
    Column,
    MetaData,
    Table,
    create_engine,
    inspect,
    select,
    type_coerce,
)
from sqlalchemy.exc import DatabaseError
from sqlalchemy.types import (
    BIGINT,
    BLOB,
    BOOLEAN,
    DATE,
    DATETIME,
    INTEGER,
    NUMERIC,
    SMALLINT,
    TEXT,
    TIMESTAMP,
    TypeDecorator,
)

from esis.util import datetime_to_timestamp

logger = logging.getLogger(__name__)


class Database(object):

    """Generic database object.

    This class is subclassed to provide additional functionality specific to
    artifacts and/or documents.

    :param db_filename: Path to the sqlite database file
    :type db_filename: str

    """

    def __init__(self, db_filename):
        """Connect to database and create session object."""
        self.db_filename = db_filename
        self.engine = create_engine(
            'sqlite:///{}'.format(db_filename),
        )
        self.connection = None
        self.metadata = MetaData(bind=self.engine)

    def connect(self):
        """Create connection."""
        logger.debug('Connecting to SQLite database: %r', self.db_filename)
        self.connection = self.engine.connect()

    def disconnect(self):
        """Close connection."""
        assert not self.connection.closed
        logger.debug(
            'Disconnecting from SQLite database: %r', self.db_filename)
        self.connection.close()

    def __enter__(self):
        """Connect on entering context."""
        self.connect()
        return self

    def __exit__(self, exc_type, exc_value, traceback):
        """Disconnect on exiting context."""
        self.disconnect()

    def __getitem__(self, table_name):
        """Get table object in database.

        :param table_name: Name of the table
        :type table_name: str
        :return: Table object that can be used in queries
        :rtype: sqlalchemy.schema.Table

        """
        table = self.metadata.tables.get(table_name)
        if table is None:
            table = Table(table_name, self.metadata, autoload=True)
        return table

    def reflect(self, table_names):
        """Get table metadata through reflection.

        sqlalchemy already provides a reflect method, but it will stop at the
        first failure, while this method will try to get as much as possible.

        :param table_names: Table names to inspect
        :type table_names: list(str)

        """
        inspector = inspect(self.engine)
        for table_name in table_names:
            columns = []
            for column_data in inspector.get_columns(table_name):
                # Rename 'type' to 'type_' to create column object
                column_type = column_data.pop('type', None)
                column_data['type_'] = column_type
                columns.append(Column(**column_data))
            Table(table_name, self.metadata, *columns)

    def run_quick_check(self):
        """Check database integrity.

        Some files, especially those files created after carving, might not
        contain completely valid data.

        """
        try:
            result = self.connection.execute('PRAGMA quick_check;')
        except DatabaseError:
            return False

        passed = result.fetchone()[0] == 'ok'
        if not passed:
            logger.warning('Integrity check failure: %r', self.db_filename)
        return passed


class DBReader(object):

    """Iterate through all db tables and rows easily.

    :param database: Database to traverse
    :type database: esis.db.Database

    """

    # Name suffixes for the shadow tables that support full text search
    FTS_SUFFIXES = (
        'content',
        'segdir',
        'segments',
        'stat',
        'docsize',
    )

    def __init__(self, database):
        """Connect to database and get table metadata."""
        self.database = database

        master_table = Table('sqlite_master', database.metadata, autoload=True)
        query = (
            select([master_table.c.name])
            .where(master_table.c.type == 'table')
        )
        result = database.connection.execute(query)
        all_table_names = set(row[0] for row in result.fetchall())

        ignored_table_names = ['sqlite_master']
        sequence_table_name = 'sqlite_sequence'
        if sequence_table_name in all_table_names:
            ignored_table_names.append(sequence_table_name)
        fts_table_names = self._get_fts_table_names(all_table_names)
        ignored_table_names.extend(fts_table_names)
        logger.debug(
            '%d tables ignored: %r',
            len(ignored_table_names),
            ', '.join(sorted(ignored_table_names)))

        ignored_table_names = set(ignored_table_names)
        table_names = all_table_names - ignored_table_names
        database.reflect(table_names)

        self.db_tables = [
            database.metadata.tables[table_name]
            for table_name in table_names
        ]
        logger.info('%d tables found', len(self.db_tables))

    def _get_fts_table_names(self, all_table_names):
        """Get a list of FTS-related table names.

        :param all_table_names: The names for all tables in the database
        :type all_table_names: set(str)

        """
        master_table = Table('sqlite_master', self.database.metadata, autoload=True)
        query = (
            select([master_table.c.name])
            .where(master_table.c.sql.like('%USING fts%'))
        )
        result = self.database.connection.execute(query)
        fts_table_names = [row[0] for row in result.fetchall()]
<commit_msg>Access table objects by name<commit_after># -*- coding: utf-8 -*-
"""Database related tools."""
import logging

from datetime import datetime

import dateutil.parser

from sqlalchemy import (
    Column,
    MetaData,
    Table,
    create_engine,
    inspect,
    select,
    type_coerce,
)
from sqlalchemy.exc import DatabaseError
from sqlalchemy.types import (
    BIGINT,
    BLOB,
    BOOLEAN,
    DATE,
    DATETIME,
    INTEGER,
    NUMERIC,
    SMALLINT,
    TEXT,
    TIMESTAMP,
    TypeDecorator,
)

from esis.util import datetime_to_timestamp

logger = logging.getLogger(__name__)


class Database(object):

    """Generic database object.

    This class is subclassed to provide additional functionality specific to
    artifacts and/or documents.

    :param db_filename: Path to the sqlite database file
    :type db_filename: str

    """

    def __init__(self, db_filename):
        """Connect to database and create session object."""
        self.db_filename = db_filename
        self.engine = create_engine(
            'sqlite:///{}'.format(db_filename),
        )
        self.connection = None
        self.metadata = MetaData(bind=self.engine)

    def connect(self):
        """Create connection."""
        logger.debug('Connecting to SQLite database: %r', self.db_filename)
        self.connection = self.engine.connect()

    def disconnect(self):
        """Close connection."""
        assert not self.connection.closed
        logger.debug(
            'Disconnecting from SQLite database: %r', self.db_filename)
        self.connection.close()

    def __enter__(self):
        """Connect on entering context."""
        self.connect()
        return self

    def __exit__(self, exc_type, exc_value, traceback):
        """Disconnect on exiting context."""
        self.disconnect()

    def __getitem__(self, table_name):
        """Get table object in database.

        :param table_name: Name of the table
        :type table_name: str
        :return: Table object that can be used in queries
        :rtype: sqlalchemy.schema.Table

        """
        table = self.metadata.tables.get(table_name)
        if table is None:
            self.reflect([table_name])
            table = Table(table_name, self.metadata, autoload=True)
        return table

    def reflect(self, table_names):
        """Get table metadata through reflection.

        sqlalchemy already provides a reflect method, but it will stop at the
        first failure, while this method will try to get as much as possible.

        :param table_names: Table names to inspect
        :type table_names: list(str)

        """
        inspector = inspect(self.engine)
        for table_name in table_names:
            columns = []
            for column_data in inspector.get_columns(table_name):
                # Rename 'type' to 'type_' to create column object
                column_type = column_data.pop('type', None)
                column_data['type_'] = column_type
                columns.append(Column(**column_data))
            Table(table_name, self.metadata, *columns)

    def run_quick_check(self):
        """Check database integrity.

        Some files, especially those files created after carving, might not
        contain completely valid data.

        """
        try:
            result = self.connection.execute('PRAGMA quick_check;')
        except DatabaseError:
            return False

        passed = result.fetchone()[0] == 'ok'
        if not passed:
            logger.warning('Integrity check failure: %r', self.db_filename)
        return passed


class DBReader(object):

    """Iterate through all db tables and rows easily.

    :param database: Database to traverse
    :type database: esis.db.Database

    """

    # Name suffixes for the shadow tables that support full text search
    FTS_SUFFIXES = (
        'content',
        'segdir',
        'segments',
        'stat',
        'docsize',
    )

    def __init__(self, database):
        """Connect to database and get table metadata."""
        self.database = database

        master_table = database['sqlite_master']
        query = (
            select([master_table.c.name])
            .where(master_table.c.type == 'table')
        )
        result = database.connection.execute(query)
        all_table_names = set(row[0] for row in result.fetchall())

        ignored_table_names = ['sqlite_master']
        sequence_table_name = 'sqlite_sequence'
        if sequence_table_name in all_table_names:
            ignored_table_names.append(sequence_table_name)
        fts_table_names = self._get_fts_table_names(all_table_names)
        ignored_table_names.extend(fts_table_names)
        logger.debug(
            '%d tables ignored: %r',
            len(ignored_table_names),
            ', '.join(sorted(ignored_table_names)))

        ignored_table_names = set(ignored_table_names)
        table_names = all_table_names - ignored_table_names

        self.db_tables = [
            database[table_name]
            for table_name in table_names
        ]
        logger.info('%d tables found', len(self.db_tables))

    def _get_fts_table_names(self, all_table_names):
        """Get a list of FTS-related table names.

        :param all_table_names: The names for all tables in the database
        :type all_table_names: set(str)

        """
        master_table = Table('sqlite_master', self.database.metadata, autoload=True)
        query = (
            select([master_table.c.name])
            .where(master_table.c.sql.like('%USING fts%'))
        )
        result = self.database.connection.execute(query)
        fts_table_names = [row[0] for row in result.fetchall()]

{'fraction_non_alphanumeric': 0.09344162081914244, 'fraction_numerical': 0.0006986289407038686, 'mean_word_length': 2.3603286384976525, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 3, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '6554568', 'n_tokens_mistral': 3249, 'n_tokens_neox': 3146, 'n_words': 860}