from airflow.models import BaseOperator
from airflow.utils.decorators import apply_defaults
from airflow.providers.postgres.hooks.postgres import PostgresHook
from airflow.providers.slack.hooks.slack import SlackHook
from config import Config
import pandas as pd
import json
import uuid
from datetime import datetime
from typing import Dict, List, Optional
import io
# import git --Airflow gives ModuleNotFoundError: No module named 'git'
import os
from tempfile import TemporaryDirectory
from util import Util
[docs]
class BedrockDataComparisonOperator(BaseOperator):
"""
Enhanced Bedrock operator for intelligent data comparison
Extends existing Bedrock operators with specific comparison capabilities
"""
ui_color = "#7FB3D5"
template_fields = [
"source_path",
"target_path",
"key_columns",
"column_mappings",
"semantic_equivalents",
]
@apply_defaults
def __init__(
self,
source_path: str, # Can be S3 path or table name
target_path: str, # Can be S3 path or table name
source_sql: Optional[str], # Optional custom SQL for source
target_sql: Optional[str], # Optional custom SQL for target
key_columns: List[str] = [],
column_mappings: Optional[Dict[str, str]] = None,
semantic_equivalents: Optional[
Dict[str, str]
] = None, # New parameter for user-defined semantic mappings
model_id: str = Config.CLAUDE_BEDROCK_SONNET,
user_prompt: str = "",
gdrive_folder_id: Optional[str] = None,
slack_webhook_conn_id: Optional[str] = None,
slack_channel: Optional[str] = "#amazon-q", # temporarily use this one
postgres_conn_id: str = "postgres_default",
comparison_threshold: float = 0.85,
sample_size: int = 1000,
inference_config: Dict = None,
git_repo_path: Optional[str] = Config.DAG_ROOT
+ "/bedrock/data_comparison_agent/SQL/",
mode: str = "overwrite",
*args,
**kwargs,
):
super().__init__(*args, **kwargs)
self.source_path = source_path
self.target_path = target_path
self.source_sql = source_sql
self.target_sql = target_sql
self.comparison_name = (
source_path.split("/")[-1].split(".")[0]
+ "_vs_"
+ target_path.split("/")[-1].split(".")[0]
)
self.key_columns = key_columns
self.column_mappings = column_mappings or {}
self.semantic_equivalents = (
semantic_equivalents or {}
) # Store user-defined semantic equivalents
self.model_id = model_id
self.gdrive_folder_id = gdrive_folder_id
self.slack_webhook_conn_id = slack_webhook_conn_id
self.slack_channel = slack_channel
self.postgres_conn_id = postgres_conn_id
self.comparison_threshold = comparison_threshold
self.sample_size = sample_size
self.inference_config = inference_config or {
"maxTokens": 4096,
"temperature": 0.1, # Low temperature for more consistent comparison
"topP": 0.9,
}
self.git_repo_path = git_repo_path
self.mode = mode
self.session_id = str(uuid.uuid4())
self.source_data_origin = None
self.target_data_origing = None
self.user_prompt = user_prompt
[docs]
def get_llm_client(self):
"""Get ChatBedrockConverse client for structured output"""
from langchain_aws import ChatBedrockConverse
from botocore.config import Config as BotocoreConfig
config = BotocoreConfig(read_timeout=600)
llm = ChatBedrockConverse(
model=self.model_id,
temperature=self.inference_config.get("temperature", 0.1),
max_tokens=self.inference_config.get("maxTokens", 4096),
top_p=self.inference_config.get("topP", 0.9),
config=config
)
return llm
[docs]
def run_query(self, data_origin, sql):
import awswrangler as wr
if data_origin == "pandas":
raise Exception("Not yet imlemented")
elif data_origin == "athena":
df = wr.athena.read_sql_query(
sql,
database=Config.DATA_LAKE_DB,
)
data_origin = "athena"
else:
# Assume it's a table name in PostgreSQL
postgres_hook = PostgresHook(self.postgres_conn_id)
df = postgres_hook.get_pandas_df(sql)
return df
[docs]
def analyze_df(self, sample_df: pd.DataFrame, path, sql, data_origin) -> dict:
import pandas as pd
import awswrangler as wr
result = {}
if not sql:
origin_sql = f"SELECT * FROM {path}"
else:
origin_sql = sql
for col, dtype in sample_df.dtypes.items():
# if data_origin=='pandas':
# # Compute stats
# mean = sample_df[col].mean()
# min_val = sample_df[col].min()
# max_val = sample_df[col].max()
# stdev = sample_df[col].std()
# distinct = sample_df[col].unique().tolist()
print("Analyzing column:", col, "of type", dtype)
if dtype in ["int64", "float64", "object", "bool", "datetime64[ns]"]:
# SQL for stats
stats_sql = f"""
SELECT
AVG(t.{col}) as mean,
MIN(t.{col}) as min,
MAX(t.{col}) as max,
STDDEV({col}) as stdev
FROM ({origin_sql}) t;
"""
print(stats_sql)
stats = self.run_query(data_origin, stats_sql)
else:
stats = pd.DataFrame(
[{"mean": None, "min": None, "max": None, "stdev": None}]
)
# SQL for distinct
distinct_sql = f"""
SELECT DISTINCT t.{col} distincts
FROM ({origin_sql}) t;
"""
print(distinct_sql)
distincts = self.run_query(data_origin, distinct_sql)
result[col] = {
"min": stats.iloc[0]["min"],
"max": stats.iloc[0]["max"],
"stdev": stats.iloc[0]["stdev"],
"mean": stats.iloc[0]["mean"],
"distinct_values": distincts["distincts"].tolist(),
"dtype": str(dtype),
}
return result
[docs]
def load_sample_data(self, path: str, sql: str):
"""Load data from S3 or PostgreSQL"""
import pandas as pd
import awswrangler as wr
if not sql:
sample_sql = f"SELECT * FROM {path} LIMIT {self.sample_size}"
else:
sample_sql = sql
data_origin = None
if path.startswith("s3://"):
# Load from S3 using awswrangler
print("Loading from S3:", path)
print("SQLs are not used when loading from S3 path directly.")
df = (
wr.s3.read_parquet(path)
if path.endswith(".parquet")
else wr.s3.read_csv(path)
)
data_origin = "pandas"
elif path == "":
df = wr.athena.read_sql_query(
sample_sql,
database=Config.DATA_LAKE_DB,
)
data_origin = "athena"
else:
# Assume it's a table name in PostgreSQL
postgres_hook = PostgresHook(self.postgres_conn_id)
df = postgres_hook.get_pandas_df(sample_sql)
data_origin = "pg"
return df, data_origin
[docs]
def invoke_bedrock_for_schema_analysis(self, source_schema, target_schema) -> Dict:
"""Use Bedrock to analyze schema differences with structured output"""
import json
from langchain_core.messages import HumanMessage, SystemMessage
print("Invoking Bedrock for schema analysis...")
print("Source Schema:", source_schema)
print("Target Schema:", target_schema)
# Get LLM client
llm = self.get_llm_client()
llm_output_schema = {
"$schema": "http://json-schema.org/draft-07/schema#",
"type": "object",
"title": "SchemaComparison",
"description": "Structured output for schema comparison",
"properties": {
"matched_columns": {
"type": "array",
"description": "List of columns that match exactly",
"items": {"type": "string"},
},
"semantic_matches": {
"type": "object",
"description": "Mapping of source columns to target columns based on semantics",
"additionalProperties": {"type": "string"},
"patternProperties": {"^.*$": {"type": "string"}},
},
"source_only": {
"type": "array",
"description": "Columns only present in the source",
"items": {"type": "string"},
},
"target_only": {
"type": "array",
"description": "Columns only present in the target",
"items": {"type": "string"},
},
"type_mismatches": {
"type": "array",
"description": "List of columns with type mismatches",
"items": {
"type": "object",
"properties": {
"column": {
"type": "string",
"description": "Name of the column with mismatch",
},
"source_type": {
"type": "string",
"description": "Data type in the source",
},
"target_type": {
"type": "string",
"description": "Data type in the target",
},
},
"required": ["column", "source_type", "target_type"],
"additionalProperties": False,
},
},
"recommendations": {
"type": "array",
"description": "List of recommendations for handling differences",
"items": {"type": "string"},
},
"confidence_score": {
"type": "number",
"description": "Confidence score for the matching results",
"minimum": 0,
"maximum": 1,
},
},
"required": [
"matched_columns",
"semantic_matches",
"source_only",
"target_only",
"type_mismatches",
"recommendations",
"confidence_score",
],
"additionalProperties": False,
}
llm_with_structured = llm.with_structured_output(llm_output_schema)
# Combine default and user-defined semantic equivalents
all_semantic_mappings = {
"brand_name": "provider_name",
"submission_date": "submitted_at",
"location_id": "loc_id",
"application_id": "app_id",
**self.semantic_equivalents, # Add user-defined mappings
}
system_prompt = """You are a data analyst expert. Compare these two data schemas and provide detailed analysis.
Return your response as a valid JSON object with the exact structure specified."""
user_prompt = f"""Compare these two data schemas:
Source Schema: {source_schema}
Target Schema: {target_schema}
Known Column Mappings: {json.dumps(self.column_mappings, indent=2)}
User-Defined Semantic Equivalents: {json.dumps(self.semantic_equivalents, indent=2)}
Consider these semantic equivalents when analyzing:
{json.dumps(all_semantic_mappings, indent=2)}
{self.user_prompt}
"""
messages = [
SystemMessage(content=system_prompt),
HumanMessage(content=user_prompt),
]
try:
response = llm_with_structured.invoke(messages)
print("Bedrock schema analysis response:", response)
# Parse JSON response
return response
except Exception as e:
# Fallback if response isn't valid JSON
self.log.warning(f"Failed to parse JSON response: {response}, error: {e}")
[docs]
def generate_comparison_sql(self, schema_analysis: Dict) -> str:
"""Generate SQL for data comparison based on schema analysis"""
matched_cols = schema_analysis.get("matched_columns", [])
semantic_matches = schema_analysis.get("semantic_matches", {})
# Build SELECT columns for comparison
select_columns = []
comparison_columns = []
# Add key columns
for key_col in self.key_columns:
select_columns.append(f"COALESCE(s.{key_col}, t.{key_col}) AS {key_col}")
# Add matched columns comparison
for col in matched_cols:
if col not in self.key_columns:
select_columns.append(f"s.{col} AS s_{col}")
select_columns.append(f"t.{col} AS t_{col}")
comparison_columns.append(
f"CASE WHEN s.{col} = t.{col} OR (s.{col} IS NULL AND t.{col} IS NULL) "
f"THEN 1 ELSE 0 END AS {col}_match"
)
# Add semantic matches comparison
for source_col, target_col in semantic_matches.items():
select_columns.append(f"s.{source_col} AS s_{source_col}")
select_columns.append(f"t.{target_col} AS t_{target_col}")
comparison_columns.append(
f"CASE WHEN s.{source_col} = t.{target_col} OR "
f"(s.{source_col} IS NULL AND t.{target_col} IS NULL) "
f"THEN 1 ELSE 0 END AS {source_col}_to_{target_col}_match"
)
# Build join conditions
join_conditions = " AND ".join([f"s.{k} = t.{k}" for k in self.key_columns])
# Source table reference
source_table = (
self.source_path # this is a table name
if not self.source_path.startswith("s3://")
else f"source_temp_{self.comparison_name}"
)
target_table = (
self.target_path
if not self.target_path.startswith("s3://")
else f"target_temp_{self.comparison_name}"
)
sql = f"""
-- Data Comparison SQL Generated by Bedrock AI
-- Comparison: {self.comparison_name}
-- Generated: {datetime.now().isoformat()}
-- Key Columns: {', '.join(self.key_columns)}
WITH comparison_details AS (
SELECT
{', '.join(select_columns)},
{', '.join(comparison_columns) if comparison_columns else '1 as dummy'},
CASE
WHEN s.{self.key_columns[0]} IS NULL THEN 'target_only'
WHEN t.{self.key_columns[0]} IS NULL THEN 'source_only'
ELSE 'both'
END AS record_status
FROM {source_table} s
FULL OUTER JOIN {target_table} t
ON {join_conditions}
),
summary_stats AS (
SELECT
record_status,
COUNT(*) as record_count,
{', '.join([f'SUM({col}_match) AS {col}_matches' for col in matched_cols if col not in self.key_columns]) if matched_cols else 'COUNT(*) as total_count'}
FROM comparison_details
GROUP BY record_status
),
overall_metrics AS (
SELECT
SUM(CASE WHEN record_status = 'both' THEN record_count ELSE 0 END) as matching_records,
SUM(CASE WHEN record_status = 'source_only' THEN record_count ELSE 0 END) as source_only_records,
SUM(CASE WHEN record_status = 'target_only' THEN record_count ELSE 0 END) as target_only_records,
SUM(record_count) as total_records
FROM summary_stats
)
SELECT * FROM summary_stats
UNION ALL
SELECT
'OVERALL' as record_status,
total_records as record_count,
{', '.join([f'NULL AS {col}_matches' for col in matched_cols if col not in self.key_columns]) if matched_cols else 'total_records as total_count'}
FROM overall_metrics;
"""
return sql
[docs]
def analyze_content_differences(
self,
source_df: pd.DataFrame,
target_df: pd.DataFrame,
schema_analysis: Dict,
) -> Dict:
"""Use Bedrock to analyze actual content differences including distinct values"""
from langchain_core.messages import HumanMessage, SystemMessage
# Get LLM client
llm = self.get_llm_client()
# Get sample of differences
merged_df = pd.merge(
source_df.head(100),
target_df.head(100),
left_on=self.key_columns,
right_on=self.key_columns,
how="outer",
suffixes=("_source", "_target"),
indicator=True,
)
# Find rows with differences
diff_samples = []
for _, row in merged_df[merged_df["_merge"] == "both"].head(20).iterrows():
differences = {}
for col in schema_analysis.get("matched_columns", []):
if col not in self.key_columns:
source_val = row.get(f"{col}_source")
target_val = row.get(f"{col}_target")
if source_val != target_val:
differences[col] = {
"source": str(source_val),
"target": str(target_val),
}
if differences:
diff_samples.append(
{
"key": {k: row[k] for k in self.key_columns if k in row},
"differences": differences,
}
)
# Analyze distinct values in each column
distinct_values_analysis = {}
for col in schema_analysis.get("matched_columns", []):
if col in source_df.columns and col in target_df.columns:
source_distinct = source_df[col].nunique()
target_distinct = target_df[col].nunique()
source_sample_values = source_df[col].value_counts().head(5).to_dict()
target_sample_values = target_df[col].value_counts().head(5).to_dict()
distinct_values_analysis[col] = {
"source_distinct_count": int(source_distinct),
"target_distinct_count": int(target_distinct),
"source_top_values": {
str(k): int(v) for k, v in source_sample_values.items()
},
"target_top_values": {
str(k): int(v) for k, v in target_sample_values.items()
},
}
system_prompt = """You are a data quality analyst.
Analyze data differences and provide insights.
Return your response as a valid JSON object with the exact structure specified."""
user_prompt = f"""Analyze these data differences and provide insights:
Sample Differences:
{json.dumps(diff_samples, indent=2)}
Distinct Values Analysis:
{json.dumps(distinct_values_analysis, indent=2)}
Schema Context:
- Matched Columns: {schema_analysis.get('matched_columns', [])}
- Semantic Matches: {schema_analysis.get('semantic_matches', {})}
Analyze and return JSON with:
{{
"difference_patterns": [
{{"pattern": "description", "frequency": "high/medium/low", "severity": "critical/major/minor"}}
],
"data_quality_issues": [
{{"issue": "description", "affected_columns": ["col1", "col2"], "recommendation": "action"}}
],
"distinct_values_summary": {{
"columns_with_cardinality_diff": ["columns with significant difference in distinct values"],
"potential_issues": ["issues identified from distinct value analysis"]
}},
"formatting_differences": ["list of formatting-only differences"],
"true_differences": ["list of actual data differences"],
"quality_score": 85,
"severity": "minor/major/critical",
"summary": "brief summary of findings",
"recommendations": ["actionable recommendations"]
}}"""
messages = [
SystemMessage(content=system_prompt),
HumanMessage(content=user_prompt),
]
try:
response = llm.invoke(messages)
response_text = response.content
result = json.loads(response_text)
# Add the raw distinct values analysis to the result
result["distinct_values_raw"] = distinct_values_analysis
return result
except json.JSONDecodeError:
return {
"difference_patterns": [],
"data_quality_issues": [],
"distinct_values_summary": {
"columns_with_cardinality_diff": [],
"potential_issues": [],
},
"distinct_values_raw": distinct_values_analysis,
"quality_score": 50,
"severity": "unknown",
"summary": "Analysis failed to parse",
"recommendations": ["Manual review required"],
}
[docs]
def generate_markdown_report(self, results: Dict) -> str:
"""Generate comprehensive markdown report"""
distinct_values = results.get("content_analysis", {}).get(
"distinct_values_raw", {}
)
distinct_summary = results.get("content_analysis", {}).get(
"distinct_values_summary", {}
)
report = f"""# Data Comparison Report - {self.comparison_name}
**Generated**: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}
**Source**: `{self.source_path}`
**Target**: `{self.target_path}`
**Session ID**: `{self.session_id}`
## Executive Summary
| Metric | Value |
|--------|-------|
| **Overall Match Rate** | {results.get('match_percentage', 0):.1f}% |
| **Data Quality Score** | {results.get('content_analysis', {}).get('quality_score', 0)}/100 |
| **Severity Level** | {results.get('content_analysis', {}).get('severity', 'N/A')} |
| **Confidence Score** | {results.get('schema_analysis', {}).get('confidence_score', 0):.2f} |
## Schema Analysis
### ✅ Matched Columns ({len(results.get('schema_analysis', {}).get('matched_columns', []))})
{chr(10).join(['- `' + col + '`' for col in results.get('schema_analysis', {}).get('matched_columns', [])])}
### 🔄 Semantic Mappings Identified
{chr(10).join(['- `' + k + '` → `' + v + '`' for k, v in results.get('schema_analysis', {}).get('semantic_matches', {}).items()])}
### ⚠️ Schema Differences
**Source Only Columns** ({len(results.get('schema_analysis', {}).get('source_only', []))}):
{', '.join(['`' + col + '`' for col in results.get('schema_analysis', {}).get('source_only', [])])}
**Target Only Columns** ({len(results.get('schema_analysis', {}).get('target_only', []))}):
{', '.join(['`' + col + '`' for col in results.get('schema_analysis', {}).get('target_only', [])])}
### 🔧 Type Mismatches
{chr(10).join(['- `' + tm['column'] + '`: ' + tm['source_type'] + ' → ' + tm['target_type'] for tm in results.get('schema_analysis', {}).get('type_mismatches', [])])}
## Content Analysis
### Summary
{results.get('content_analysis', {}).get('summary', 'No summary available')}
### Difference Patterns
{chr(10).join(['- **' + p['pattern'] + '** (Frequency: ' + p['frequency'] + ', Severity: ' + p['severity'] + ')' for p in results.get('content_analysis', {}).get('difference_patterns', [])])}
### Data Quality Issues
{chr(10).join(['- **' + issue['issue'] + '**' + chr(10) + ' - Affected: ' + ', '.join(issue['affected_columns']) + chr(10) + ' - Action: ' + issue['recommendation'] for issue in results.get('content_analysis', {}).get('data_quality_issues', [])])}
### Distinct Values Analysis
**Columns with Cardinality Differences:**
{', '.join(['`' + col + '`' for col in distinct_summary.get('columns_with_cardinality_diff', [])])}
**Potential Issues from Distinct Values:**
{chr(10).join(['- ' + issue for issue in distinct_summary.get('potential_issues', [])])}
## Generated SQL
```sql
{results.get('generated_sql', '-- No SQL generated')}
```
## Recommendations
{chr(10).join(['1. ' + rec for rec in results.get('content_analysis', {}).get('recommendations', [])])}
## Metadata
- **Comparison Threshold**: {self.comparison_threshold}
- **Sample Size**: {self.sample_size}
- **Model Used**: `{self.model_id}`
- **User-Defined Semantic Equivalents**: {len(self.semantic_equivalents)} mappings
---
*Report generated by BedrockDataComparisonOperator*
"""
return report
[docs]
def upload_results_to_gdrive(
self, results_df: pd.DataFrame, report_content: str
) -> Optional[str]:
"""Upload results DataFrame to Google Drive using GDriveUtil"""
if not self.gdrive_folder_id:
return None
gdrive_service = Util.construct_gdrive_service()
timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
try:
timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
# Upload results DataFrame as Google Sheet
sheet_name = f"{self.comparison_name}_results_{timestamp}"
sheet_id = Util.write_df_to_gsheet(
df=results_df,
file_name=sheet_name,
sheet_name="Results",
parent_folder_id=self.gdrive_folder_id,
overwrite=True,
)
# Also upload the markdown report
drive_service = Util.construct_gdrive_service()
# Save report to temporary file and upload
with TemporaryDirectory() as temp_dir:
temp_file_path = os.path.join(
temp_dir, f"{self.comparison_name}_report_{timestamp}.md"
)
with open(temp_file_path, "w") as f:
f.write(report_content)
report_file_id = Util.gdrive_upload_file(
drive_service=drive_service,
file_path=temp_file_path,
file_name=f"{self.comparison_name}_report_{timestamp}.md",
parent_folder_id=self.gdrive_folder_id,
overwrite=False,
)
# Return the sheet URL
return f"https://docs.google.com/spreadsheets/d/{sheet_id}/edit"
except Exception as e:
self.log.error(f"Failed to upload to Google Drive: {e}")
return None
[docs]
def send_slack_notification(self, results: Dict, report_url: str) -> None:
"""Send Slack notification with feedback request"""
try:
environment = "#prod" if Config.PRODUCTION else "#dev"
slack_hook = SlackHook(
slack_conn_id=Config.SLACK_DEFAULT_API,
)
# Create Slack message
message = {
"text": f"Data Comparison Complete: {self.comparison_name}",
"blocks": [
{
"type": "section",
"text": {
"type": "mrkdwn",
"text": f"*Data Comparison Complete*\n"
f"environment: {environment}\n"
f"📊 *Name*: {self.comparison_name}\n"
f"✅ *Match Rate*: {results.get('match_percentage', 0):.1f}%\n"
f"📈 *Quality Score*: {results.get('content_analysis', {}).get('quality_score', 0)}/100\n"
f"⚠️ *Severity*: {results.get('content_analysis', {}).get('severity', 'N/A')}",
},
},
{
"type": "section",
"text": {
"type": "mrkdwn",
"text": f"<{report_url}|📄 View Full Report>",
},
},
{"type": "divider"},
{
"type": "section",
"text": {
"type": "mrkdwn",
"text": "*Please rate this comparison:*",
},
},
{
"type": "actions",
"elements": [
{
"type": "button",
"text": {"type": "plain_text", "text": "👍 Accurate"},
"style": "primary",
"value": f"accurate_{self.session_id}",
},
{
"type": "button",
"text": {
"type": "plain_text",
"text": "🤔 Needs Review",
},
"style": "danger",
"value": f"review_{self.session_id}",
},
{
"type": "button",
"text": {
"type": "plain_text",
"text": "👎 Issues Found",
},
"style": "danger",
"value": f"issues_{self.session_id}",
},
],
},
],
}
if self.slack_channel:
message["channel"] = self.slack_channel
slack_hook.send(message=message)
except Exception as e:
self.log.error(f"Failed to send Slack notification: {e}")
[docs]
def commit_sql_to_git(self, sql: str) -> None:
"""Commit generated SQL to git repository in a new branch"""
if not self.git_repo_path:
return
try:
# repo = git.Repo(self.git_repo_path) Airflow gives ModuleNotFoundError: No module named 'git'
# Create a new branch for this comparison
timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
branch_name = f"comparison_{self.comparison_name}_{timestamp}"
# Get current branch
original_branch = repo.active_branch
print(f"Current branch: {original_branch}")
# Create and checkout new branch
new_branch = repo.create_head(branch_name)
new_branch.checkout()
print(f"Switched to new branch: {branch_name}")
print("*" * 100)
# Create directory if it doesn't exist ## THIS SHOULD BE A CONFIG PARAMETER
sql_dir = os.path.join(self.git_repo_path, "generated")
os.makedirs(sql_dir, exist_ok=True)
print(f"SQL directory: {sql_dir}")
print("*" * 100)
sql_file_path = os.path.join(
sql_dir, f"{self.comparison_name}_{timestamp}.sql"
)
# Write SQL file
with open(sql_file_path, "w") as f:
f.write(sql)
# Commit to git
repo.index.add([sql_file_path])
repo.index.commit(f"Generated comparison SQL for {self.comparison_name}")
# Push the new branch to remote
origin = repo.remote(name="origin")
origin.push(refspec=f"{branch_name}:{branch_name}")
self.log.info(
f"SQL committed to git in branch '{branch_name}': {sql_file_path}"
)
# Switch back to original branch
original_branch.checkout()
except Exception as e:
self.log.error(f"Failed to commit SQL to git: {e}")
[docs]
def execute(self, context):
"""Main execution method"""
self.log.info(f"Starting data comparison: {self.comparison_name}")
# Load data
self.log.info("Loading source and target data...")
source_df, data_origin = self.load_sample_data(
self.source_path, self.source_sql
)
self.source_data_origin = data_origin
target_df, data_origin = self.load_sample_data(
self.target_path, self.target_sql
)
self.target_data_origing = data_origin
self.log.info(
f"Source shape: {source_df.shape}, Target shape: {target_df.shape}"
)
# Step 1: Schema Analysis
source_schema = self.analyze_df(
source_df, self.source_path, self.source_sql, self.source_data_origin
)
target_schema = self.analyze_df(
target_df, self.target_path, self.target_sql, self.target_data_origing
)
self.log.info("Performing schema analysis with Bedrock...")
schema_analysis = self.invoke_bedrock_for_schema_analysis(
source_schema, target_schema
)
# Step 2: Generate SQL
self.log.info("Generating comparison SQL...")
comparison_sql = self.generate_comparison_sql(schema_analysis)
# Commit SQL to git in a new branch
self.commit_sql_to_git(comparison_sql)
# Step 3: Content Analysis (now includes distinct values)
self.log.info("Analyzing content differences and distinct values...")
content_analysis = self.analyze_content_differences(
source_df, target_df, schema_analysis
)
# Step 4: Calculate metrics
total_columns = len(schema_analysis.get("matched_columns", [])) + len(
schema_analysis.get("semantic_matches", {})
)
match_percentage = (
len(schema_analysis.get("matched_columns", [])) / max(total_columns, 1)
) * 100
# Compile results
results = {
"comparison_name": self.comparison_name,
"source_path": self.source_path,
"target_path": self.target_path,
"schema_analysis": schema_analysis,
"content_analysis": content_analysis,
"generated_sql": comparison_sql,
"match_percentage": match_percentage,
"timestamp": datetime.now().isoformat(),
"session_id": self.session_id,
"user_defined_semantic_equivalents": self.semantic_equivalents,
}
# Step 5: Generate report
self.log.info("Generating markdown report...")
report = self.generate_markdown_report(results)
# Step 6: Save to S3
self.log.info("Saving results to S3...")
# Step 7: Create results DataFrame and upload to Google Drive
results_df = pd.DataFrame(
[
{
"comparison_name": self.comparison_name,
"timestamp": datetime.now().isoformat(),
"source_path": self.source_path,
"target_path": self.target_path,
"match_percentage": match_percentage,
"quality_score": content_analysis.get("quality_score", 0),
"severity": content_analysis.get("severity", "unknown"),
"session_id": self.session_id,
"user_semantic_mappings": json.dumps(self.semantic_equivalents),
"results_json": json.dumps(results),
}
]
)
gdrive_url = None
if self.gdrive_folder_id:
self.log.info("Uploading results to Google Drive...")
gdrive_url = self.upload_results_to_gdrive(results_df, report)
# Step 8: Send Slack notification
if self.slack_webhook_conn_id:
self.log.info("Sending Slack notification...")
self.send_slack_notification(results, gdrive_url)
# Push results to XCom for downstream tasks
context["task_instance"].xcom_push(
key="comparison_results",
value={
"report_gdrive_url": gdrive_url,
"match_percentage": match_percentage,
"quality_score": content_analysis.get("quality_score", 0),
"severity": content_analysis.get("severity", "unknown"),
"session_id": self.session_id,
"distinct_values_summary": content_analysis.get(
"distinct_values_summary", {}
),
},
)
if gdrive_url:
self.log.info(f"Google Drive report: {gdrive_url}")
return results