indexer.py 3.48 KB
Newer Older
Sebastian Schüpbach's avatar
Sebastian Schüpbach committed
1
2
3
4
import logging
import os
import time

Sebastian Schüpbach's avatar
Sebastian Schüpbach committed
5
6
import mysql.connector as mariadb

Sebastian Schüpbach's avatar
Sebastian Schüpbach committed
7
8
9
10
11
12
13
14
15
16
17
18
19
20

# noinspection SqlResolve,SqlNoDataSourceInspection
class Indexer:
    def __init__(self):
        """
        Start MariaDB client
        """
        self.mariadb_connection, self.mariadb_cursor = self._connect_to_mariadb()

    def _connect_to_mariadb(self, retries=0):
        """
        Connect to MariaDB. Abort after configured retries.
        """
        try:
Sebastian Schüpbach's avatar
Sebastian Schüpbach committed
21
22
            logging.debug(f"Connecting to DB {os.environ['MARIADB_DATABASE']} on "
                          f"{os.environ['MARIADB_HOST']}:{os.environ['MARIADB_PORT']}")
Sebastian Schüpbach's avatar
Sebastian Schüpbach committed
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
            mariadb_connection = mariadb.connect(user=os.environ['MARIADB_USER'],
                                                 password=os.environ['MARIADB_PASSWORD'],
                                                 host=os.environ['MARIADB_HOST'],
                                                 port=int(os.environ['MARIADB_PORT']),
                                                 database=os.environ['MARIADB_DATABASE'])
            mariadb_connection.autocommit = False
            mariadb_cursor = mariadb_connection.cursor()
            return mariadb_connection, mariadb_cursor
        except Exception as ex:
            status = 'Exception: ' + str(ex)
            logging.error(status)
            if retries < int(os.environ['MARIADB_CONNECTION_RETRIES']):
                time.sleep(30 * (retries + 1))
                self._connect_to_mariadb(retries + 1)
            exit(1)

    @staticmethod
40
    def _create_sql_stmt(table_name, record, fields) -> (str, tuple):
Sebastian Schüpbach's avatar
Sebastian Schüpbach committed
41
42
43
        """
        Create SQL statement
        """
44
45
46
47
48
49
        db_values = [record[f] for f in fields if f in record and record[f]]
        db_values.extend([record[f] for f in fields if f != 'sig' and f in record and record[f]])
        db_fields = ','.join([f for f in fields if f in record and record[f]])
        db_value_placeholders = ', '.join(['?' for f in fields if f in record and record[f]])
        key_value = ", ".join([f"{f}=?" for f in fields
                               if f != 'sig' and f in record and record[f]])
Sebastian Schüpbach's avatar
Sebastian Schüpbach committed
50
51
        # noinspection SqlNoDataSourceInspection
        return 'INSERT INTO {} ({}) VALUES ({}) ON DUPLICATE KEY UPDATE {}'.format(
52
            table_name, db_fields, db_value_placeholders, key_value), tuple(db_values)
Sebastian Schüpbach's avatar
Sebastian Schüpbach committed
53

54
    def insert_in_db(self, record) -> (bool, str):
Sebastian Schüpbach's avatar
Sebastian Schüpbach committed
55
56
57
        """
        Insert record in DB
        """
58
59
        stmt, values = Indexer._create_sql_stmt('entities', record,
                                                ['sig', 'uri', 'access', 'proto'])
Sebastian Schüpbach's avatar
Sebastian Schüpbach committed
60
        try:
61
62
63
64
65
            self.mariadb_cursor.execute(stmt, values)
            stmt, values = Indexer._create_sql_stmt('metadata', record,
                                                    ['sig', 'mimetype', 'height',
                                                     'width', 'duration', 'type'])
            self.mariadb_cursor.execute(stmt, values)
66
67
            return True, ""
        except mariadb.Error as ex:
68
69
            logging.error(f'Problems in sql statement (statement: "{stmt}", '
                          f'parameters: {values}): {ex}')
70
            return False, str(ex)
Sebastian Schüpbach's avatar
Sebastian Schüpbach committed
71
72
73
74
75

    def commit(self):
        """
        Commit changes to DB
        """
Sebastian Schüpbach's avatar
Sebastian Schüpbach committed
76
        logging.debug("Commiting changes to DB")
Sebastian Schüpbach's avatar
Sebastian Schüpbach committed
77
78
79
80
81
82
        self.mariadb_connection.commit()

    def rollback(self):
        """
        Rollback changes
        """
Sebastian Schüpbach's avatar
Sebastian Schüpbach committed
83
        logging.info("Rollback changes")
Sebastian Schüpbach's avatar
Sebastian Schüpbach committed
84
85
        self.mariadb_cursor.reset()
        self.mariadb_connection.rollback()