"""Vendor model."""
from sqlalchemy import Column, Integer

from artist.connectors import mysql
from artist.models.country import Country


class Vendor(mysql.BaseModel):
    """Vendor Model.

    Represents Vendor metadata.
    """

    __tablename__ = 'vendor'

    vendor_id = Column(
        Integer, primary_key=True, autoincrement=True, nullable=False)
    country = Column(Integer)


def get_country_code_by_vendor_id(vendor_id):
    """Get country code by vendor ID."""
    with mysql.db_session() as session:
        country_id = session.query(Vendor).filter_by(vendor_id=vendor_id).first().country
        if not country_id:
            return country_id
        country_code = session.query(
            Country).filter_by(id=country_id).first().country_code
    return country_code