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