Spaces:
Running on Zero
Running on Zero
| # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. | |
| # SPDX-License-Identifier: OpenMDW-1.1 | |
| from functools import lru_cache | |
| from pathlib import Path | |
| import psycopg2 | |
| import yaml | |
| from psycopg2.extras import NamedTupleCursor | |
| def get_recipe_cursor(is_production_database: bool = False) -> NamedTupleCursor: | |
| config_path = Path("credentials/config.yaml") | |
| if not config_path.exists(): | |
| raise FileNotFoundError(f"Credential file for Recipe System is not found at {config_path}") | |
| dbname = None | |
| user = None | |
| password = None | |
| endpoint = None | |
| with open(config_path, "r") as file: | |
| config = yaml.safe_load(file) | |
| postgres_profiles = config.get("postgres", {}).get("profiles", {}) | |
| if is_production_database: | |
| dbname = "dirdb_prd" | |
| profile_name = "dirdb_viewer_prd" | |
| else: | |
| dbname = "dirdb_dev" | |
| profile_name = "dirdb_viewer_dev" | |
| profile = postgres_profiles.get(profile_name, {}) | |
| user = profile.get("user") | |
| password = profile.get("password") | |
| endpoint = profile.get("endpoint") | |
| if not all([user, password, endpoint]): | |
| raise ValueError(f"Missing database credentials for profile: {profile_name}") | |
| try: | |
| conn = psycopg2.connect( | |
| database=dbname, | |
| user=user, | |
| password=password, | |
| host=endpoint, | |
| port=5432, | |
| ) | |
| postgres_cursor = conn.cursor(cursor_factory=NamedTupleCursor) | |
| return postgres_cursor | |
| except psycopg2.Error as e: | |
| raise ValueError(f"Failed to connect to the database for Recipe System: {e}") from e | |
| def get_datasource_sql(datacollection_name: str, storage_type: str, data_type: str) -> str: | |
| return f""" | |
| SELECT collection.datasource_name,collection.ratio,source.name,source.sensitivity,source.text_only | |
| FROM datacollection_to_datasource as collection | |
| INNER JOIN datasource as source | |
| ON collection.datasource_name=source.name | |
| WHERE datacollection_name='{datacollection_name}' AND collection.data_type='{data_type}' AND collection.storage_type='{storage_type}' AND source.storage_type='{storage_type}' | |
| """ | |
| def get_wdinfo_sql(datasource: str, storage_type: str, data_type: str, split_type: str = "train") -> str: | |
| """Get wdinfo SQL query for a datasource. | |
| Args: | |
| datasource: Name of the datasource | |
| storage_type: Storage type (s3, pdx, neb_eu) | |
| data_type: Data type (vlm, vfm, etc) | |
| split_type: Dataset split (train, val, test etc.), defaults to train | |
| Returns: | |
| SQL query string | |
| """ | |
| return f""" | |
| SELECT datasource_name, wdinfo | |
| FROM datasource_to_wdinfo | |
| WHERE datasource_name='{datasource}' AND data_type='{data_type}' AND storage_type='{storage_type}' AND split_type='{split_type}' | |
| """ | |