Get 3 sample rows from a table in JSON format to help AI understand the data
(ds: CoreDatasource, table_name: str, fields: list)
| 433 | |
| 434 | |
| 435 | def get_table_sample_data(ds: CoreDatasource, table_name: str, fields: list) -> str: |
| 436 | """Get 3 sample rows from a table in JSON format to help AI understand the data""" |
| 437 | if not fields: |
| 438 | return "" |
| 439 | |
| 440 | db = DB.get_db(ds.type) |
| 441 | # Get prefix/suffix for identifier quoting |
| 442 | prefix = db.prefix if hasattr(db, 'prefix') else '"' |
| 443 | suffix = db.suffix if hasattr(db, 'suffix') else '"' |
| 444 | |
| 445 | # Build field list with proper quoting |
| 446 | field_names = [] |
| 447 | for field in fields[:10]: # Limit to first 10 fields to avoid too wide results |
| 448 | field_name = f"{prefix}{field.field_name}{suffix}" |
| 449 | field_names.append(field_name) |
| 450 | |
| 451 | # Build LIMIT query based on database type |
| 452 | if equals_ignore_case(ds.type, "sqlServer"): |
| 453 | query = f"SELECT TOP 3 {','.join(field_names)} FROM {prefix}{table_name}{suffix}" |
| 454 | elif equals_ignore_case(ds.type, "ck"): |
| 455 | query = f"SELECT {','.join(field_names)} FROM {table_name} LIMIT 3" |
| 456 | elif equals_ignore_case(ds.type, "hive"): |
| 457 | query = f"SELECT {','.join(field_names)} FROM {table_name} LIMIT 3" |
| 458 | elif equals_ignore_case(ds.type, "oracle"): |
| 459 | query = f"SELECT {','.join(field_names)} FROM \"{table_name}\" WHERE ROWNUM <= 3" |
| 460 | elif equals_ignore_case(ds.type, "dm"): |
| 461 | query = f"SELECT {','.join(field_names)} FROM \"{table_name}\" WHERE ROWNUM <= 3" |
| 462 | else: |
| 463 | query = f"SELECT {','.join(field_names)} FROM {prefix}{table_name}{suffix} LIMIT 3" |
| 464 | |
| 465 | try: |
| 466 | result = exec_sql(ds=ds, sql=query, origin_column=True) |
| 467 | if result and result.get('data') and len(result['data']) > 0: |
| 468 | import json |
| 469 | # Truncate long string values for readability |
| 470 | json_rows = [] |
| 471 | for row in result['data'][:3]: |
| 472 | truncated_row = {} |
| 473 | for key, value in row.items(): |
| 474 | if value is None: |
| 475 | truncated_row[key] = None |
| 476 | elif isinstance(value, str): |
| 477 | # Truncate long strings |
| 478 | if len(value) > 100: |
| 479 | value = value[:100] + '...' |
| 480 | truncated_row[key] = value.replace('\n', ' ').replace('\r', ' ') |
| 481 | else: |
| 482 | truncated_row[key] = value |
| 483 | json_rows.append(truncated_row) |
| 484 | return json.dumps(json_rows, ensure_ascii=False, indent=2) |
| 485 | except Exception: |
| 486 | pass |
| 487 | return "" |
| 488 | |
| 489 | |
| 490 | def get_tables_sample_data(session: SessionDep, current_user: CurrentUser, ds: CoreDatasource, |
no test coverage detected