src/api/services/postgresql_service.py

import logging
from typing import List
from src.api.database import Database

logger = logging.getLogger(__name__)

class PostgresqlService:
    def __init__(self, db: Database):
        self.db = db


    async def get_db_rows_from_dataset_name(self, dataset) -> List:
        query_mapping = {
            "re-mind/product_type_embedding": {
                "column": "type",
                "target_column": "product_type",
                "table": "product_producttype",
                "optional": {
                    "column": "name",
                    "target_column": "product_category",
                    "table": "product_category",
                    "foreign_key": "category_id",
                }
            },
            "re-mind/marketplace_name_embedding": {
                "column": "name",
                "target_column": "marketplace_name",
                "table": "invoice_marketplace",
            },
            "re-mind/manufacturer_name_embedding": {
                "column": "name",
                "target_column": "manufacturer_name",
                "table": "product_manufacturer",
            },
            "re-mind/seller_name_embedding": {
                "column": "name",
                "target_column": "seller_name",
                "table": "product_seller",
            }
        }

        db_table = query_mapping[dataset]["table"]
        column = query_mapping[dataset]["column"]
        target_column = query_mapping[dataset]["target_column"]
        optional = query_mapping[dataset].get("optional")
        optional_query_1 = ""
        optional_query_2 = ""
        if optional:
            optional_table = optional["table"]
            optional_column = optional["column"]
            optional_target_column = optional["target_column"]
            optional_foreign_key = optional["foreign_key"]

            optional_query_1 = f", {optional_table}.{optional_column} AS {optional_target_column}"
            optional_query_2 = f"INNER JOIN {optional_table} ON {db_table}.{optional_foreign_key} = {optional_table}.id"
        query = f"SELECT {db_table}.{column} AS {target_column} {optional_query_1} FROM {db_table} {optional_query_2}"

        logger.info(query)
        results = await self.db.fetch(query)

        return results