from supertenant import consts
from supertenant.supermeter.data.base import BaseData


class DBData(BaseData):
    def __init__(self, resource_type, integration_module):
        # type: (str, str) -> None
        super(DBData, self).__init__(resource_type, integration_module)

    def set_db(self, db):
        # type: (str) -> None
        self.set_tag(consts.LABEL_DB_DATABASE, db)

    def set_statement(self, statement):
        # type: (str) -> None
        self.set_tag(consts.LABEL_DB_STATEMENT, statement)

    def set_user(self, user):
        # type: (str) -> None
        self.set_tag(consts.LABEL_DB_USER, user)

    def set_host(self, host):
        # type: (str) -> None
        self.set_tag(consts.LABEL_DB_HOST, host)

    def set_port(self, port):
        # type: (str) -> None
        self.set_tag(consts.LABEL_DB_PORT, port)


class PGData(DBData):
    def __init__(self, integration_module):
        # type: (str) -> None
        super(PGData, self).__init__(consts.RESOURCE_TYPE_POSTGRESQL, integration_module)
        self.set_span_type(consts.SPAN_TYPE_CLIENT_REQUEST)


class MYSQLData(DBData):
    def __init__(self, integration_module):
        # type: (str) -> None
        super(MYSQLData, self).__init__(consts.RESOURCE_TYPE_MYSQL, integration_module)
        self.set_span_type(consts.SPAN_TYPE_CLIENT_REQUEST)


class MongoData(DBData):
    def __init__(self, integration_module):
        # type: (str) -> None
        super(MongoData, self).__init__(consts.RESOURCE_TYPE_MONGODB, integration_module)
        self.set_span_type(consts.SPAN_TYPE_CLIENT_REQUEST)

    def set_mongodb_collection(self, collection):
        # type: (str) -> None
        self.set_tag(consts.LABEL_DB_MONGODB_COLLECTION, collection)

    def set_command(self, command):
        # type: (str) -> None
        self.set_tag(consts.LABEL_DB_COMMAND, command)

    def set_mongodb_filter(self, filter):
        # type: (str) -> None
        self.set_tag(consts.LABEL_DB_MONGODB_FILTER, filter)

    def set_mongodb_command_json(self, json):
        # type: (str) -> None
        self.set_tag(consts.LABEL_DB_MONGODB_COMMAND_JSON, json)