YDB_QA query statistic (#27570)

Co-authored-by: Copilot <[email protected]>

1 files changed, 51 insertions, 88 deletions

diff --git a/.github/scripts/analytics/github_issue_mapping.py b/.github/scripts/analytics/github_issue_mapping.py
index 12d26f58e41..1ac2f5ac27d 100755
--- a/.github/scripts/analytics/github_issue_mapping.py
+++ b/.github/scripts/analytics/github_issue_mapping.py
@@ -7,27 +7,18 @@ This table will be used by SQL queries to join muted test data with GitHub issue
 
 import os
 import ydb
-import configparser
 import time
 import sys
+from ydb_wrapper import YDBWrapper
 
 # Import shared GitHub issue utilities
 sys.path.append(os.path.join(os.path.dirname(__file__), '..'))
 from github_issue_utils import create_test_issue_mapping
 
-# Load YDB configuration
-dir = os.path.dirname(__file__)
-config = configparser.ConfigParser()
-config_file_path = f"{dir}/../../config/ydb_qa_db.ini"
-config.read(config_file_path)
 
-DATABASE_ENDPOINT = config["QA_DB"]["DATABASE_ENDPOINT"]
-DATABASE_PATH = config["QA_DB"]["DATABASE_PATH"]
 
 
-
-
-def get_github_issues_data(driver):
+def get_github_issues_data(ydb_wrapper):
     """Get GitHub issues data from the github_data/issues table"""
     query = """
     SELECT 
@@ -44,30 +35,18 @@ def get_github_issues_data(driver):
     """
     
     print("Fetching GitHub issues data...")
-    start_time = time.time()
     
-    results = []
     try:
-        scan_query = ydb.ScanQuery(query, {})
-        it = driver.table_client.scan_query(scan_query)
-        
-        while True:
-            try:
-                result = next(it)
-                results.extend(result.result_set.rows)
-            except StopIteration:
-                break
+        results = ydb_wrapper.execute_scan_query(query)
+        print(f'Fetched {len(results)} GitHub issues')
+        return results
     except Exception as e:
         print(f"Warning: Could not fetch GitHub issues data: {e}")
         print("This might be because the github_data/issues table doesn't exist yet.")
         return []
 
-    end_time = time.time()
-    print(f'Fetched {len(results)} GitHub issues, duration: {end_time - start_time}s')
-    return results
 
-
-def create_test_issue_mapping_table(session, table_path):
+def create_test_issue_mapping_table(ydb_wrapper, table_path):
     """Create the test-to-issue mapping table"""
     print(f"Creating test-to-issue mapping table: {table_path}")
     
@@ -89,7 +68,7 @@ def create_test_issue_mapping_table(session, table_path):
     """
 
     print(f"Creating table with query: {create_table_sql}")
-    session.execute_scheme(create_table_sql)
+    ydb_wrapper.create_table(table_path, create_table_sql)
 
 
 def convert_mapping_to_table_data(test_to_issue_mapping):
@@ -117,10 +96,9 @@ def convert_mapping_to_table_data(test_to_issue_mapping):
     return table_data
 
 
-def bulk_upsert_mapping_data(table_client, table_path, mapping_data):
+def bulk_upsert_mapping_data(ydb_wrapper, table_path, mapping_data):
     """Bulk upsert mapping data into the table"""
     print(f"Bulk upserting {len(mapping_data)} test-to-issue mappings to {table_path}")
-    start_time = time.time()
     
     column_types = ydb.BulkUpsertColumns()
     column_types.add_column('full_name', ydb.PrimitiveType.Utf8)
@@ -130,75 +108,60 @@ def bulk_upsert_mapping_data(table_client, table_path, mapping_data):
     column_types.add_column('github_issue_number', ydb.OptionalType(ydb.PrimitiveType.Uint64))
     column_types.add_column('github_issue_state', ydb.PrimitiveType.Utf8)
     column_types.add_column('github_issue_created_at', ydb.OptionalType(ydb.PrimitiveType.Timestamp))
-    table_client.bulk_upsert(table_path, mapping_data, column_types)
     
-    end_time = time.time()
-    print(f"Bulk upsert completed, duration: {end_time - start_time}s")
+    ydb_wrapper.bulk_upsert(table_path, mapping_data, column_types)
+    print(f"Bulk upsert completed")
 
 
 def main():
     """Main function to create the test-to-issue mapping table"""
     print("Starting GitHub issue mapping table creation")
     script_start_time = time.time()
+    script_name = os.path.basename(__file__)
     
-    # Check environment variables
-    if "CI_YDB_SERVICE_ACCOUNT_KEY_FILE_CREDENTIALS" not in os.environ:
-        print("Error: Env variable CI_YDB_SERVICE_ACCOUNT_KEY_FILE_CREDENTIALS is missing, skipping")
-        return 1
-    
-    os.environ["YDB_SERVICE_ACCOUNT_KEY_FILE_CREDENTIALS"] = os.environ[
-        "CI_YDB_SERVICE_ACCOUNT_KEY_FILE_CREDENTIALS"
-    ]
-    
-    table_path = "test_results/analytics/github_issue_mapping"
-    full_table_path = f"{DATABASE_PATH}/{table_path}"
-    
-    try:
-        with ydb.Driver(
-            endpoint=DATABASE_ENDPOINT,
-            database=DATABASE_PATH,
-            credentials=ydb.credentials_from_env_variables()
-        ) as driver:
-            driver.wait(timeout=10, fail_fast=True)
+    # Initialize YDB wrapper with context manager for automatic cleanup
+    with YDBWrapper(script_name=script_name) as ydb_wrapper:
+        # Check credentials
+        if not ydb_wrapper.check_credentials():
+            return 1
+        
+        table_path = "test_results/analytics/github_issue_mapping"
+        full_table_path = f"{ydb_wrapper.database_path}/{table_path}"
+        
+        try:
+            # Get GitHub issues data
+            issues_data = get_github_issues_data(ydb_wrapper)
+            
+            if not issues_data:
+                print("No GitHub issues data found")
+                return 0
+            
+            # Create test-to-issue mapping using shared utilities
+            print("Creating test-to-issue mapping...")
+            test_to_issue = create_test_issue_mapping(issues_data)
+            print(f"Created mapping for {len(test_to_issue)} unique test names")
+            
+            # Convert mapping to table data format
+            mapping_data = convert_mapping_to_table_data(test_to_issue)
+            print(f"Converted to {len(mapping_data)} table records")
             
-            with ydb.SessionPool(driver) as pool:
-                # Get GitHub issues data
-                issues_data = get_github_issues_data(driver)
-                
-                if not issues_data:
-                    print("No GitHub issues data found")
-                    return 0
-                
-                # Create test-to-issue mapping using shared utilities
-                print("Creating test-to-issue mapping...")
-                test_to_issue = create_test_issue_mapping(issues_data)
-                print(f"Created mapping for {len(test_to_issue)} unique test names")
-                
-                # Convert mapping to table data format
-                mapping_data = convert_mapping_to_table_data(test_to_issue)
-                print(f"Converted to {len(mapping_data)} table records")
-                
-                # Create mapping table
-                def create_table_wrapper(session):
-                    create_test_issue_mapping_table(session, table_path)
-                    return True
-                
-                pool.retry_operation_sync(create_table_wrapper)
-                
-                # Bulk upsert mapping data
-                if mapping_data:
-                    bulk_upsert_mapping_data(driver.table_client, full_table_path, mapping_data)
-                else:
-                    print("No mapping data to insert")
+            # Create mapping table
+            create_test_issue_mapping_table(ydb_wrapper, full_table_path)
+            
+            # Bulk upsert mapping data
+            if mapping_data:
+                bulk_upsert_mapping_data(ydb_wrapper, full_table_path, mapping_data)
+            else:
+                print("No mapping data to insert")
         
-        script_elapsed = time.time() - script_start_time
-        print(f"Script completed successfully, total time: {script_elapsed:.2f}s")
+            script_elapsed = time.time() - script_start_time
+            print(f"Script completed successfully, total time: {script_elapsed:.2f}s")
+            
+        except Exception as e:
+            print(f"Error during execution: {e}")
+            return 1
         
-    except Exception as e:
-        print(f"Error during execution: {e}")
-        return 1
-    
-    return 0
+        return 0
 
 
 if __name__ == "__main__":