# SPDX-License-Identifier: GPL-2.0-or-later
#
#   database.py
#   Library for processing results from XMLSQLparser and
#   query a PostgreSQL database based on the input data
#
#   Copyright 2009 - 2013   David Sommerseth <davids@redhat.com>
#

import psycopg2
import types

class Database(object):
    def __init__(self, host=None, port=None, user=None, password=None, database=None,
                 noaction=False, debug=False):
        self.noaction = noaction
        self.debug = debug

        dsnd = {}
        if host is not None:
            dsnd['host'] = host
            dsnd['sslmode'] = 'require'
        if port is not None:
            dsnd['port'] = str(port)
            dsnd['sslmode'] = 'require'
        if user is not None:
            dsnd['user'] = user
        if password is not None:
            dsnd['password'] = password
        if database is not None:
            dsnd['dbname'] = database

        dsn = " ".join(["%s='%s'" %(k,v) for (k,v) in list(dsnd.items())])
        self.conn = not self.noaction and psycopg2.connect(dsn) or None


    def INSERT(self, sqlvars):
        #
        # Validate input data
        #
        if type(sqlvars) is not dict:
            raise AttributeError('Input parameter is not a Python dict')

        try:
            sqlvars['table']
            sqlvars['fields']
            sqlvars['records']
        except KeyError as err:
            raise KeyError("Input dictionary do not contain a required element: %s").with_traceback(str(err))

        if type(sqlvars['fields']) is not list:
            raise AttributeError("The 'fields' element is not a list of fields")

        if type(sqlvars['records']) is not list:
            raise AttributeError("The 'records' element is not a list of fields")

        if len(sqlvars['records']) == 0:
            return True

        try:
            sqlvars['returning']
        except:
            sqlvars['returning'] = None

        #
        # Build SQL template
        #
        sqlstub = "INSERT INTO %s (%s) VALUES (%s)" % (
            sqlvars['table'],
            ",".join(sqlvars['fields']),
            ",".join(["%%(%s)s" % f for f in sqlvars['fields']])
            )

        # Get a database cursor
        curs = not self.noaction and self.conn.cursor() or None

        #
        # Loop through all records and insert them into the database
        #
        results = []
        for rec in sqlvars['records']:
            if type(rec) is not list:
                raise AttributeError("The field values inside the 'records' list must be in a list")

            # Create a dictionary, which will be used for the SQL operation
            values = {}
            for i in range(0, len(sqlvars['fields'])):
                values[sqlvars['fields'][i]] = rec[i]

            if self.debug:
                print("SQL QUERY: ==> %s" % (sqlstub % values))

            # Do the INSERT query
            if not self.noaction:
                curs.execute(sqlstub, values)

            # If a return value for the INSERT is defined, catch that one
            if not self.noaction and sqlvars['returning']:
                # The psycopg2 do not handle INSERT INTO ... RETURNING column queries, so we can only use
                # this on tables with oid and do the look up that way
                vls = {"table": sqlvars['table'], 'colname': sqlvars['returning'], 'oid': str(curs.lastrowid)}
                curs.execute("SELECT %(colname)s FROM %(table)s WHERE oid='%(oid)s'" % vls)
                results.append(curs.fetchone()[0])
            else:
                results.append(True)

        if not self.noaction:
            curs.close()
        return results


    def DELETE(self, table, where):
        try:
            sql = "DELETE FROM %s WHERE %s" % (
                table,
                " AND ".join(["%s = %%(%s)s" % (k,k) for (k,v) in list(where.items())])
                )

            if self.debug:
                print("SQL QUERY ==> %s" % (sql % where))

            if not self.noaction:
                curs = self.conn.cursor()
                curs.execute(sql, where)
                delrows = curs.rowcount
                curs.close()
                return delrows
            else:
                return 0
        except Exception as err:
            raise Exception("** SQL ERROR ** %s\n** SQL ERROR ** Message: %s" % ((sql % where), str(err)))

    def SELECT(self, table, fields, joins=None, where=None):
        curs = not self.noaction and self.conn.cursor() or None

        # Query
        try:
            sql = "SELECT %s FROM %s %s %s" % (
                ",".join(fields),
                table,
                joins and "%s" % joins or "",
                where and "WHERE %s" % " AND ".join(["%s = %%(%s)s" % (k,k) for (k,v) in list(where.items())] or "")
                )
            if self.debug:
                print("SQL QUERY: ==> %s" % (sql % where))
            if not self.noaction:
                curs.execute(sql, where)
            else:
                # If no action is setup (mainly for debugging), return empty result set
                return {"table": table, "fields": [], "records": []}
        except Exception as err:
            raise Exception("** SQL ERROR *** %s\n** SQL ERROR ** Message: %s" % (where and (sql % where) or sql, str(err)))

        # Extract field names
        fields = []
        for fn in curs.description:
            fields.append(fn[0])

        # Extract records
        records = []
        for dbrec in curs.fetchall():
            values = []
            for val in dbrec:
                values.append(val)
            records.append(values)

        curs.close()
        if self.debug:
            print("database::SELECT() result ** Fields: %s\nRecords: %s" % (fields, records))
        return {"table": table, "fields": fields, "records": records}

    def COMMIT(self):
        # Commit the work
        if not self.noaction:
            self.conn.commit()

    def ROLLBACK(self):
        # Abort / rollback the current work
        if not self.noaction:
            self.conn.rollback()


    def GetValue(self, dbres, recidx, field):
        "Helper function to easy extract a field from a record set"

        # Check that input data good
        if type(dbres) is not dict:
            raise AttributeError('Database result parameter is not a Python dict')

        try:
            dbres['table']
            dbres['fields']
            dbres['records']
        except KeyError as err:
            raise KeyError("Database result parameter do not contain a required element: %s").with_traceback(str(err))

        if type(dbres['fields']) is not list:
            raise AttributeError("The 'fields' element is not a list of fields")

        if type(dbres['records']) is not list:
            raise AttributeError("The 'records' element is not a list of fields")

        # Return None when we're going out of boundaries
        if recidx >= len(dbres['records']):
            return None

        if type(field) == bytes:
            # Find the field index of the field name in the records set
            try:
                fidx = dbres['fields'].index(field)
            except ValueError:
                raise Exception("Field '%s' is not found in the database result" % field)
        elif type(field) == int:
            # If the field value is integer, assume it is the numeric field id
            if field >= len(dbres['fields']):
                raise Exception("Field id '%i' is too high.  No field available" % field)
            fidx = field

        # Return the value
        return dbres['records'][recidx][fidx]


    def NumTuples(self, dbres):
        # Check that input data good
        if type(dbres) is not dict:
            raise AttributeError('Database result parameter is not a Python dict')

        try:
            dbres['table']
            dbres['fields']
            dbres['records']
        except KeyError as err:
            raise KeyError("Database result parameter do not contain a required element: %s").with_traceback(str(err))

        if type(dbres['records']) is not list:
            raise AttributeError("The 'records' element is not a list of fields")

        return len(dbres['records'])
