"""
BigQuery — Lineage Collection and Push (combined)
===================================================
Imports ``collect()`` from ``collect_lineage`` and ``push()`` from
``push_lineage``, runs both in sequence.

Substitution points (search for "← SUBSTITUTE"):
  - BIGQUERY_PROJECT_ID   : GCP project ID to collect from
  - BIGQUERY_REGION       : BigQuery region for INFORMATION_SCHEMA queries (e.g. "us", "eu")
  - LOOKBACK_HOURS        : how far back to scan job history (default 24 h)
  - MCD_INGEST_ID / MCD_INGEST_TOKEN : Monte Carlo API credentials
  - MCD_RESOURCE_UUID      : UUID of the BigQuery connection in Monte Carlo

Prerequisites:
  pip install google-cloud-bigquery pycarlo
"""

from __future__ import annotations

import argparse
import os

from collect_lineage import collect, LOOKBACK_HOURS
from push_lineage import push, _BATCH_SIZE


def main() -> None:
    parser = argparse.ArgumentParser(description="Push BigQuery lineage to Monte Carlo")
    parser.add_argument("--project-id", default=os.getenv("BIGQUERY_PROJECT_ID"))  # ← SUBSTITUTE
    parser.add_argument("--region", default=os.getenv("BIGQUERY_REGION", "us"))    # ← SUBSTITUTE
    parser.add_argument("--resource-uuid", default=os.getenv("MCD_RESOURCE_UUID"))
    parser.add_argument("--key-id", default=os.getenv("MCD_INGEST_ID"))
    parser.add_argument("--key-token", default=os.getenv("MCD_INGEST_TOKEN"))
    parser.add_argument("--lookback-hours", type=int, default=LOOKBACK_HOURS)
    parser.add_argument("--output-file", default="lineage_output.json")
    parser.add_argument("--push-result-file", default="lineage_push_result.json")
    parser.add_argument(
        "--batch-size",
        type=int,
        default=_BATCH_SIZE,
        help=f"Max events per push batch (default: {_BATCH_SIZE})",
    )
    args = parser.parse_args()

    required = ["project_id", "resource_uuid", "key_id", "key_token"]
    missing = [k for k in required if getattr(args, k) is None]
    if missing:
        parser.error(f"Missing required arguments/env vars: {missing}")

    # Step 1: Collect
    collect(
        project_id=args.project_id,
        region=args.region,
        lookback_hours=args.lookback_hours,
        output_file=args.output_file,
    )

    # Step 2: Push
    push(
        input_file=args.output_file,
        resource_uuid=args.resource_uuid,
        key_id=args.key_id,
        key_token=args.key_token,
        batch_size=args.batch_size,
        output_file=args.push_result_file,
    )


if __name__ == "__main__":
    main()

"""
BigQuery — Metadata Collection and Push (combined)
===================================================
Imports ``collect()`` from ``collect_metadata`` and ``push()`` from
``push_metadata``, runs both in sequence.

Substitution points (search for "← SUBSTITUTE"):
  - BIGQUERY_PROJECT_ID   : GCP project ID to collect from
  - GOOGLE_APPLICATION_CREDENTIALS : path to service-account JSON key file
  - MCD_INGEST_ID / MCD_INGEST_TOKEN : Monte Carlo API credentials
  - MCD_RESOURCE_UUID      : UUID of the BigQuery connection in Monte Carlo
  - DATASET_EXCLUSIONS    : datasets to skip (informational / system datasets)

Prerequisites:
  pip install google-cloud-bigquery pycarlo
"""

from __future__ import annotations

import argparse
import os

from collect_metadata import collect
from push_metadata import push, _BATCH_SIZE


def main() -> None:
    parser = argparse.ArgumentParser(description="Push BigQuery metadata to Monte Carlo")
    parser.add_argument("--project-id", default=os.getenv("BIGQUERY_PROJECT_ID"))  # ← SUBSTITUTE
    parser.add_argument("--resource-uuid", default=os.getenv("MCD_RESOURCE_UUID"))
    parser.add_argument("--key-id", default=os.getenv("MCD_INGEST_ID"))
    parser.add_argument("--key-token", default=os.getenv("MCD_INGEST_TOKEN"))
    parser.add_argument("--output-file", default="metadata_output.json")
    parser.add_argument("--push-result-file", default="metadata_push_result.json")
    parser.add_argument(
        "--batch-size",
        type=int,
        default=_BATCH_SIZE,
        help=f"Max assets per push batch (default: {_BATCH_SIZE})",
    )
    args = parser.parse_args()

    missing = [k for k, v in vars(args).items() if v is None and k not in ("output_file", "push_result_file", "batch_size")]
    if missing:
        parser.error(f"Missing required arguments/env vars: {missing}")

    # Step 1: Collect
    collect(
        project_id=args.project_id,
        output_file=args.output_file,
    )

    # Step 2: Push
    push(
        input_file=args.output_file,
        resource_uuid=args.resource_uuid,
        key_id=args.key_id,
        key_token=args.key_token,
        batch_size=args.batch_size,
        output_file=args.push_result_file,
    )


if __name__ == "__main__":
    main()

"""
BigQuery — Query Log Collection and Push (combined)
=====================================================
Imports ``collect()`` from ``collect_query_logs`` and ``push()`` from
``push_query_logs``, runs both in sequence.

Substitution points (search for "← SUBSTITUTE"):
  - BIGQUERY_PROJECT_ID   : GCP project ID to collect query logs from
  - GOOGLE_APPLICATION_CREDENTIALS : path to service-account JSON key file
  - LOOKBACK_HOURS        : how many hours back to collect (default 25, skip last 1 h)
  - MCD_INGEST_ID / MCD_INGEST_TOKEN : Monte Carlo API credentials
  - MCD_RESOURCE_UUID      : UUID of the BigQuery connection in Monte Carlo

Prerequisites:
  pip install google-cloud-bigquery pycarlo
"""

from __future__ import annotations

import argparse
import os

from collect_query_logs import collect, LOOKBACK_HOURS, LOOKBACK_LAG_HOURS
from push_query_logs import push, _BATCH_SIZE


def main() -> None:
    parser = argparse.ArgumentParser(description="Push BigQuery query logs to Monte Carlo")
    parser.add_argument("--project-id", default=os.getenv("BIGQUERY_PROJECT_ID"))  # ← SUBSTITUTE
    parser.add_argument("--resource-uuid", default=os.getenv("MCD_RESOURCE_UUID"))
    parser.add_argument("--key-id", default=os.getenv("MCD_INGEST_ID"))
    parser.add_argument("--key-token", default=os.getenv("MCD_INGEST_TOKEN"))
    parser.add_argument("--lookback-hours", type=int, default=LOOKBACK_HOURS)
    parser.add_argument("--lookback-lag-hours", type=int, default=LOOKBACK_LAG_HOURS)
    parser.add_argument("--output-file", default="query_logs_output.json")
    parser.add_argument("--push-result-file", default="query_logs_push_result.json")
    parser.add_argument(
        "--batch-size",
        type=int,
        default=_BATCH_SIZE,
        help=f"Max entries per push batch (default: {_BATCH_SIZE})",
    )
    args = parser.parse_args()

    required = ["project_id", "resource_uuid", "key_id", "key_token"]
    missing = [k for k in required if getattr(args, k) is None]
    if missing:
        parser.error(f"Missing required arguments/env vars: {missing}")

    # Step 1: Collect
    collect(
        project_id=args.project_id,
        lookback_hours=args.lookback_hours,
        lookback_lag_hours=args.lookback_lag_hours,
        output_file=args.output_file,
    )

    # Step 2: Push
    push(
        input_file=args.output_file,
        resource_uuid=args.resource_uuid,
        key_id=args.key_id,
        key_token=args.key_token,
        batch_size=args.batch_size,
        output_file=args.push_result_file,
    )


if __name__ == "__main__":
    main()

"""
BigQuery — Lineage Collection (collect only)
=============================================
Collects table-level lineage from two sources:
  1. INFORMATION_SCHEMA.SCHEMATA_LINKS — cross-project dataset shares (per region)
  2. Job query history — SQL parsing for CREATE TABLE AS SELECT and INSERT INTO
     SELECT patterns to derive source->destination relationships.

Writes the collected lineage edges to a JSON manifest file.

Can be run standalone via CLI or imported (use the ``collect()`` function).

Substitution points (search for "← SUBSTITUTE"):
  - BIGQUERY_PROJECT_ID   : GCP project ID to collect from
  - BIGQUERY_REGION       : BigQuery region for INFORMATION_SCHEMA queries (e.g. "us", "eu")
  - LOOKBACK_HOURS        : how far back to scan job history (default 24 h)

Prerequisites:
  pip install google-cloud-bigquery
"""

from __future__ import annotations

import argparse
import json
import logging
import os
import re
from datetime import datetime, timedelta, timezone

from google.cloud import bigquery

logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s")
log = logging.getLogger(__name__)

RESOURCE_TYPE = "bigquery"
LOOKBACK_HOURS = int(os.getenv("LOOKBACK_HOURS", "24"))  # ← SUBSTITUTE: adjust lookback window

# Regex patterns to detect CTAS and INSERT INTO SELECT in BigQuery SQL
_CTAS_PATTERN = re.compile(
    r"CREATE\s+(?:OR\s+REPLACE\s+)?(?:TABLE|VIEW)\s+`?(?P<dest>[\w.\-]+)`?"
    r".*?(?:AS\s+)?SELECT\b",
    re.IGNORECASE | re.DOTALL,
)
_INSERT_PATTERN = re.compile(
    r"INSERT\s+(?:INTO\s+)?`?(?P<dest>[\w.\-]+)`?.*?SELECT\b",
    re.IGNORECASE | re.DOTALL,
)
_TABLE_REF_PATTERN = re.compile(r"`?([\w\-]+\.[\w\-]+\.[\w\-]+)`?", re.IGNORECASE)


def _parse_full_name(full_name: str) -> tuple[str, str, str]:
    """Split 'project.dataset.table' into (project, dataset, table)."""
    parts = full_name.replace("`", "").split(".")
    if len(parts) == 3:
        return parts[0], parts[1], parts[2]
    if len(parts) == 2:
        return "", parts[0], parts[1]
    return "", "", parts[0]


def _collect_schema_link_lineage(
    bq_client: bigquery.Client,
    project_id: str,
    region: str,
) -> list[dict]:
    """Collect cross-project lineage from INFORMATION_SCHEMA.SCHEMATA_LINKS."""
    query = f"""
        SELECT
            CATALOG_NAME            AS source_project,
            SCHEMA_NAME             AS source_dataset,
            LINKED_SCHEMA_CATALOG_NAME AS destination_project,
            LINKED_SCHEMA_NAME      AS destination_dataset
        FROM `{project_id}`.`{region}`.INFORMATION_SCHEMA.SCHEMATA_LINKS
    """  # ← SUBSTITUTE: update project_id and region as needed
    edges: list[dict] = []
    try:
        for row in bq_client.query(query).result():
            edges.append(
                {
                    "destination": {
                        "database": row.destination_project,
                        "schema": row.destination_dataset,
                        "table": "*",
                    },
                    "sources": [
                        {
                            "database": row.source_project,
                            "schema": row.source_dataset,
                            "table": "*",
                        }
                    ],
                }
            )
    except Exception:
        log.warning("SCHEMATA_LINKS query failed — skipping dataset-share lineage", exc_info=True)
    return edges


def _collect_query_lineage(
    bq_client: bigquery.Client,
    project_id: str,
    lookback_hours: int,
) -> list[dict]:
    """Derive lineage by parsing CTAS/INSERT patterns in job query history."""
    end_dt = datetime.now(timezone.utc)
    start_dt = end_dt - timedelta(hours=lookback_hours)

    edges: list[dict] = []
    for job in bq_client.list_jobs(all_users=True, min_creation_time=start_dt, max_creation_time=end_dt):
        sql: str = getattr(job, "query", None) or ""
        if not sql.strip():
            continue

        dest_match = _CTAS_PATTERN.search(sql) or _INSERT_PATTERN.search(sql)
        if not dest_match:
            continue

        dest_full = dest_match.group("dest")
        dest_project, dest_dataset, dest_table = _parse_full_name(dest_full)
        if not dest_table:
            continue

        # Collect all 3-part table references in the query as sources, excluding destination
        source_refs = [
            m.group(1)
            for m in _TABLE_REF_PATTERN.finditer(sql)
            if m.group(1) != dest_full
        ]
        if not source_refs:
            continue

        unique_sources = list(dict.fromkeys(source_refs))
        sources = []
        for ref in unique_sources:
            p, d, t = _parse_full_name(ref)
            sources.append({"database": p, "schema": d, "table": t})

        edges.append(
            {
                "destination": {
                    "database": dest_project or project_id,
                    "schema": dest_dataset,
                    "table": dest_table,
                },
                "sources": sources,
            }
        )

    return edges


def collect(
    project_id: str,
    region: str = "us",
    lookback_hours: int = LOOKBACK_HOURS,
    output_file: str = "lineage_output.json",
) -> dict:
    """
    Connect to BigQuery, collect lineage edges, and write a JSON manifest.

    Returns the manifest dict.
    """
    bq_client = bigquery.Client(project=project_id)

    log.info("Collecting lineage from project %s ...", project_id)
    schema_edges = _collect_schema_link_lineage(bq_client, project_id, region)
    query_edges = _collect_query_lineage(bq_client, project_id, lookback_hours)
    all_edges = schema_edges + query_edges

    log.info(
        "Collected %d lineage edges (%d schema-link, %d query-derived)",
        len(all_edges), len(schema_edges), len(query_edges),
    )

    manifest = {
        "resource_type": RESOURCE_TYPE,
        "collected_at": datetime.now(timezone.utc).isoformat(),
        "schema_link_edges": len(schema_edges),
        "query_derived_edges": len(query_edges),
        "edges": all_edges,
    }
    with open(output_file, "w") as fh:
        json.dump(manifest, fh, indent=2)
    log.info("Lineage manifest written to %s", output_file)

    return manifest


def main() -> None:
    parser = argparse.ArgumentParser(
        description="Collect BigQuery lineage and write to a manifest file",
    )
    parser.add_argument("--project-id", default=os.getenv("BIGQUERY_PROJECT_ID"))  # ← SUBSTITUTE
    parser.add_argument("--region", default=os.getenv("BIGQUERY_REGION", "us"))    # ← SUBSTITUTE
    parser.add_argument("--lookback-hours", type=int, default=LOOKBACK_HOURS)
    parser.add_argument("--output-file", default="lineage_output.json")
    args = parser.parse_args()

    required = ["project_id"]
    missing = [k for k in required if getattr(args, k) is None]
    if missing:
        parser.error(f"Missing required arguments/env vars: {missing}")

    collect(
        project_id=args.project_id,
        region=args.region,
        lookback_hours=args.lookback_hours,
        output_file=args.output_file,
    )


if __name__ == "__main__":
    main()

"""
BigQuery — Metadata Collection (collect only)
==============================================
Collects table schemas, row counts, byte sizes, and descriptions from all
datasets in a BigQuery project and writes them to a JSON manifest file.

Can be run standalone via CLI or imported (use the ``collect()`` function).

Substitution points (search for "← SUBSTITUTE"):
  - BIGQUERY_PROJECT_ID   : GCP project ID to collect from
  - GOOGLE_APPLICATION_CREDENTIALS : path to service-account JSON key file
  - DATASET_EXCLUSIONS    : datasets to skip (informational / system datasets)

Prerequisites:
  pip install google-cloud-bigquery
"""

from __future__ import annotations

import argparse
import json
import logging
import os
from datetime import datetime, timezone

from google.cloud import bigquery

logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s")
log = logging.getLogger(__name__)

RESOURCE_TYPE = "bigquery"

# Datasets to skip — add any internal / system datasets here
DATASET_EXCLUSIONS = {  # ← SUBSTITUTE: add datasets to exclude
    "_bqc_",
    "INFORMATION_SCHEMA",
}

# BigQuery type → Monte Carlo canonical type
BQ_TYPE_MAP: dict[str, str] = {
    "INT64": "INTEGER",
    "INTEGER": "INTEGER",
    "FLOAT64": "FLOAT",
    "FLOAT": "FLOAT",
    "BOOL": "BOOLEAN",
    "BOOLEAN": "BOOLEAN",
    "STRING": "VARCHAR",
    "BYTES": "BINARY",
    "DATE": "DATE",
    "DATETIME": "DATETIME",
    "TIMESTAMP": "TIMESTAMP",
    "TIME": "TIME",
    "NUMERIC": "DECIMAL",
    "BIGNUMERIC": "DECIMAL",
    "RECORD": "STRUCT",
    "STRUCT": "STRUCT",
    "REPEATED": "ARRAY",
    "JSON": "JSON",
    "GEOGRAPHY": "GEOGRAPHY",
}


def map_bq_type(bq_type: str) -> str:
    return BQ_TYPE_MAP.get(bq_type.upper(), bq_type.upper())


def _collect_assets(bq_client: bigquery.Client, project_id: str) -> list[dict]:
    """Collect table metadata from BigQuery and return as a list of dicts."""
    assets: list[dict] = []

    for dataset_item in bq_client.list_datasets():
        dataset_id = dataset_item.dataset_id

        if any(exc in dataset_id for exc in DATASET_EXCLUSIONS):
            log.info("Skipping dataset %s", dataset_id)
            continue

        dataset_ref = bq_client.dataset(dataset_id)

        for table_item in bq_client.list_tables(dataset_ref):
            table_ref = dataset_ref.table(table_item.table_id)
            table = bq_client.get_table(table_ref)

            fields = [
                {
                    "name": field.name,
                    "type": map_bq_type(field.field_type),
                    "description": field.description or None,
                }
                for field in table.schema
            ]

            asset = {
                "name": table.table_id,
                "database": project_id,  # ← SUBSTITUTE: use project or dataset as database
                "schema": dataset_id,
                "type": "VIEW" if table.table_type == "VIEW" else "TABLE",
                "description": table.description or None,
                "fields": fields,
                "volume": {
                    "row_count": table.num_rows,
                    "byte_count": table.num_bytes,
                },
                "freshness": {
                    "last_updated_time": table.modified.isoformat() if table.modified else None,
                },
            }
            assets.append(asset)
            log.info("Queued %s.%s.%s", project_id, dataset_id, table.table_id)

    return assets


def collect(
    project_id: str,
    output_file: str = "metadata_output.json",
) -> dict:
    """
    Connect to BigQuery, collect table metadata, and write a JSON manifest.

    Returns the manifest dict.
    """
    bq_client = bigquery.Client(project=project_id)  # ← SUBSTITUTE: adjust auth if needed

    log.info("Collecting metadata from project %s ...", project_id)
    assets = _collect_assets(bq_client, project_id)
    log.info("Collected %d asset(s).", len(assets))

    manifest = {
        "resource_type": RESOURCE_TYPE,
        "collected_at": datetime.now(timezone.utc).isoformat(),
        "assets": assets,
    }
    with open(output_file, "w") as fh:
        json.dump(manifest, fh, indent=2)
    log.info("Asset manifest written to %s", output_file)

    return manifest


def main() -> None:
    parser = argparse.ArgumentParser(
        description="Collect BigQuery metadata and write to a manifest file",
    )
    parser.add_argument("--project-id", default=os.getenv("BIGQUERY_PROJECT_ID"))  # ← SUBSTITUTE
    parser.add_argument("--output-file", default="metadata_output.json")
    args = parser.parse_args()

    missing = [k for k, v in vars(args).items() if v is None and k != "output_file"]
    if missing:
        parser.error(f"Missing required arguments/env vars: {missing}")

    collect(
        project_id=args.project_id,
        output_file=args.output_file,
    )


if __name__ == "__main__":
    main()

"""
BigQuery — Query Log Collection (collect only)
================================================
Collects completed job query logs from BigQuery job history and writes them to
a JSON manifest file for later push to Monte Carlo.

Can be run standalone via CLI or imported (use the ``collect()`` function).

Substitution points (search for "← SUBSTITUTE"):
  - BIGQUERY_PROJECT_ID   : GCP project ID to collect query logs from
  - GOOGLE_APPLICATION_CREDENTIALS : path to service-account JSON key file
  - LOOKBACK_HOURS        : how many hours back to collect (default 25, skip last 1 h)
  - STATEMENT_TYPE_FILTER : restrict to specific statement types, or leave empty for all
  - MAX_JOBS              : cap on number of jobs to collect per run

Prerequisites:
  pip install google-cloud-bigquery
"""

from __future__ import annotations

import argparse
import json
import logging
import os
from datetime import datetime, timedelta, timezone

from google.cloud import bigquery

logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s")
log = logging.getLogger(__name__)

LOG_TYPE = "bigquery"

# Collect jobs from [now - LOOKBACK_HOURS] to [now - LOOKBACK_LAG_HOURS].
# The lag avoids collecting in-flight jobs that have not yet completed.
LOOKBACK_HOURS: int = int(os.getenv("LOOKBACK_HOURS", "25"))        # ← SUBSTITUTE
LOOKBACK_LAG_HOURS: int = int(os.getenv("LOOKBACK_LAG_HOURS", "1")) # ← SUBSTITUTE

# Limit statement types — e.g. ["SELECT", "CREATE_TABLE_AS_SELECT", "INSERT"]
# Set to an empty list to collect all statement types.
STATEMENT_TYPE_FILTER: list[str] = []  # ← SUBSTITUTE

# Maximum number of jobs to collect in a single run to avoid runaway costs
MAX_JOBS: int = int(os.getenv("MAX_JOBS", "10000"))  # ← SUBSTITUTE


def _safe_isoformat(dt: datetime | None) -> str | None:
    if dt is None:
        return None
    if dt.tzinfo is None:
        dt = dt.replace(tzinfo=timezone.utc)
    return dt.isoformat()


def _collect_query_logs(
    bq_client: bigquery.Client,
    project_id: str,
    start_dt: datetime,
    end_dt: datetime,
) -> list[dict]:
    """Collect query logs from BigQuery job history and return as a list of dicts."""
    entries: list[dict] = []

    log.info(
        "Listing jobs for project=%s from %s to %s",
        project_id, start_dt.isoformat(), end_dt.isoformat(),
    )

    for job in bq_client.list_jobs(
        project=project_id,
        all_users=True,
        min_creation_time=start_dt,
        max_creation_time=end_dt,
    ):
        # Only process query jobs that have SQL text
        sql: str = getattr(job, "query", None) or ""
        if not sql.strip():
            continue

        statement_type: str = getattr(job, "statement_type", None) or ""
        if STATEMENT_TYPE_FILTER and statement_type not in STATEMENT_TYPE_FILTER:
            continue  # ← SUBSTITUTE: adjust filter as needed

        total_bytes_billed: int | None = getattr(job, "total_bytes_billed", None)

        entries.append(
            {
                "query_id": job.job_id,
                "query_text": sql,
                "start_time": _safe_isoformat(getattr(job, "created", None)),
                "end_time": _safe_isoformat(getattr(job, "ended", None)),
                "user": getattr(job, "user_email", None),
                "total_bytes_billed": total_bytes_billed,
                "statement_type": statement_type or None,
            }
        )

        if len(entries) >= MAX_JOBS:
            log.warning("Reached MAX_JOBS=%d — stopping early", MAX_JOBS)
            break

    return entries


def collect(
    project_id: str,
    lookback_hours: int = LOOKBACK_HOURS,
    lookback_lag_hours: int = LOOKBACK_LAG_HOURS,
    output_file: str = "query_logs_output.json",
) -> dict:
    """
    Connect to BigQuery, collect query logs, and write a JSON manifest.

    Returns the manifest dict.
    """
    bq_client = bigquery.Client(project=project_id)  # ← SUBSTITUTE: adjust auth if needed

    end_dt = datetime.now(timezone.utc) - timedelta(hours=lookback_lag_hours)
    start_dt = end_dt - timedelta(hours=lookback_hours)

    entries = _collect_query_logs(bq_client, project_id, start_dt, end_dt)
    log.info("Collected %d query log entries.", len(entries))

    manifest = {
        "log_type": LOG_TYPE,
        "collected_at": datetime.now(timezone.utc).isoformat(),
        "window_start": start_dt.isoformat(),
        "window_end": end_dt.isoformat(),
        "query_log_count": len(entries),
        "queries": entries,
    }
    with open(output_file, "w") as fh:
        json.dump(manifest, fh, indent=2)
    log.info("Query log manifest written to %s", output_file)

    return manifest


def main() -> None:
    parser = argparse.ArgumentParser(
        description="Collect BigQuery query logs and write to a manifest file",
    )
    parser.add_argument("--project-id", default=os.getenv("BIGQUERY_PROJECT_ID"))  # ← SUBSTITUTE
    parser.add_argument("--lookback-hours", type=int, default=LOOKBACK_HOURS)
    parser.add_argument("--lookback-lag-hours", type=int, default=LOOKBACK_LAG_HOURS)
    parser.add_argument("--output-file", default="query_logs_output.json")
    args = parser.parse_args()

    required = ["project_id"]
    missing = [k for k in required if getattr(args, k) is None]
    if missing:
        parser.error(f"Missing required arguments/env vars: {missing}")

    collect(
        project_id=args.project_id,
        lookback_hours=args.lookback_hours,
        lookback_lag_hours=args.lookback_lag_hours,
        output_file=args.output_file,
    )


if __name__ == "__main__":
    main()

"""
BigQuery — Lineage Push (push only)
====================================
Reads a manifest file produced by ``collect_lineage.py`` and pushes the lineage
events to Monte Carlo using the pycarlo push ingestion API.  Large payloads are
split into batches to stay under the 1 MB compressed limit.

Can be run standalone via CLI or imported (use the ``push()`` function).

Substitution points (search for "← SUBSTITUTE"):
  - MCD_INGEST_ID / MCD_INGEST_TOKEN : Monte Carlo API credentials
  - MCD_RESOURCE_UUID      : UUID of the BigQuery connection in Monte Carlo

Prerequisites:
  pip install pycarlo
"""

from __future__ import annotations

import argparse
import json
import logging
import os
from concurrent.futures import ThreadPoolExecutor, as_completed
from datetime import datetime, timezone

from pycarlo.core import Client, Session
from pycarlo.features.ingestion import IngestionService
from pycarlo.features.ingestion.models import (
    LineageAssetRef,
    LineageEvent,
)

logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s")
log = logging.getLogger(__name__)

RESOURCE_TYPE = "bigquery"

# Maximum events per batch — conservative default to keep compressed payload under 1 MB
# ← SUBSTITUTE: tune based on average edge complexity (number of sources per event)
_BATCH_SIZE = 500


def _make_ref(database: str, schema: str, table: str) -> LineageAssetRef:
    return LineageAssetRef(
        type="TABLE",
        name=table,
        database=database,
        schema=schema,
    )


def _build_events(edges: list[dict]) -> list[LineageEvent]:
    """Build LineageEvent objects from manifest edge dicts."""
    events = []
    for edge in edges:
        dest = edge["destination"]
        sources = edge.get("sources", [])
        if not sources:
            continue
        events.append(
            LineageEvent(
                destination=_make_ref(dest["database"], dest["schema"], dest["table"]),
                sources=[
                    _make_ref(s["database"], s["schema"], s["table"])
                    for s in sources
                ],
            )
        )
    return events


def push(
    input_file: str,
    resource_uuid: str,
    key_id: str,
    key_token: str,
    batch_size: int = _BATCH_SIZE,
    output_file: str = "lineage_push_result.json",
) -> dict:
    """
    Read a lineage manifest and push events to Monte Carlo in batches.

    Returns a result dict with invocation IDs for each batch.
    """
    with open(input_file) as fh:
        manifest = json.load(fh)

    edges = manifest.get("edges", [])
    resource_type = manifest.get("resource_type", RESOURCE_TYPE)
    events = _build_events(edges)
    log.info("Loaded %d lineage event(s) from %s", len(events), input_file)

    if not events:
        log.info("No lineage events to push.")
        push_result = {
            "resource_uuid": resource_uuid,
            "resource_type": resource_type,
            "invocation_ids": [],
            "pushed_at": datetime.now(timezone.utc).isoformat(),
            "total_events": 0,
            "batch_count": 0,
            "batch_size": batch_size,
        }
        with open(output_file, "w") as fh:
            json.dump(push_result, fh, indent=2)
        return push_result

    # Split into batches
    batches = []
    for i in range(0, len(events), batch_size):
        batches.append(events[i : i + batch_size])
    total_batches = len(batches)

    def _push_batch(batch: list, batch_num: int) -> str | None:
        """Push a single batch using a dedicated Session (thread-safe)."""
        log.info("Pushing batch %d/%d (%d events) ...", batch_num, total_batches, len(batch))
        client = Client(session=Session(mcd_id=key_id, mcd_token=key_token, scope="Ingestion"))
        service = IngestionService(mc_client=client)
        result = service.send_lineage(
            resource_uuid=resource_uuid,
            resource_type=resource_type,
            events=batch,
        )
        invocation_id = service.extract_invocation_id(result)
        if invocation_id:
            log.info("  Batch %d: invocation_id=%s", batch_num, invocation_id)
        return invocation_id

    # Push batches in parallel (each thread gets its own pycarlo Session)
    max_workers = min(4, total_batches)
    invocation_ids: list[str | None] = [None] * total_batches

    with ThreadPoolExecutor(max_workers=max_workers) as pool:
        futures = {
            pool.submit(_push_batch, batch, i + 1): i
            for i, batch in enumerate(batches)
        }
        for future in as_completed(futures):
            idx = futures[future]
            try:
                invocation_ids[idx] = future.result()
            except Exception as exc:
                log.error("ERROR pushing batch %d: %s", idx + 1, exc)
                raise

    log.info("All %d batches pushed (%d workers)", total_batches, max_workers)

    push_result = {
        "resource_uuid": resource_uuid,
        "resource_type": resource_type,
        "invocation_ids": invocation_ids,
        "pushed_at": datetime.now(timezone.utc).isoformat(),
        "total_events": len(events),
        "batch_count": total_batches,
        "batch_size": batch_size,
    }
    with open(output_file, "w") as fh:
        json.dump(push_result, fh, indent=2)
    log.info("Push result written to %s", output_file)

    return push_result


def main() -> None:
    parser = argparse.ArgumentParser(
        description="Push BigQuery lineage from a manifest to Monte Carlo",
    )
    parser.add_argument("--resource-uuid", default=os.getenv("MCD_RESOURCE_UUID"))
    parser.add_argument("--key-id", default=os.getenv("MCD_INGEST_ID"))
    parser.add_argument("--key-token", default=os.getenv("MCD_INGEST_TOKEN"))
    parser.add_argument("--input-file", default="lineage_output.json")
    parser.add_argument("--output-file", default="lineage_push_result.json")
    parser.add_argument(
        "--batch-size",
        type=int,
        default=_BATCH_SIZE,
        help=f"Max events per push batch (default: {_BATCH_SIZE})",
    )
    args = parser.parse_args()

    required = ["resource_uuid", "key_id", "key_token"]
    missing = [k for k in required if getattr(args, k) is None]
    if missing:
        parser.error(f"Missing required arguments/env vars: {missing}")

    push(
        input_file=args.input_file,
        resource_uuid=args.resource_uuid,
        key_id=args.key_id,
        key_token=args.key_token,
        batch_size=args.batch_size,
        output_file=args.output_file,
    )


if __name__ == "__main__":
    main()

"""
BigQuery — Metadata Push (push only)
=====================================
Reads a manifest file produced by ``collect_metadata.py`` and pushes the assets
to Monte Carlo using the pycarlo push ingestion API.  Large payloads are split
into batches to stay under the 1 MB compressed limit.

Can be run standalone via CLI or imported (use the ``push()`` function).

Substitution points (search for "← SUBSTITUTE"):
  - MCD_INGEST_ID / MCD_INGEST_TOKEN : Monte Carlo API credentials
  - MCD_RESOURCE_UUID      : UUID of the BigQuery connection in Monte Carlo

Prerequisites:
  pip install pycarlo
"""

from __future__ import annotations

import argparse
import json
import logging
import os
from concurrent.futures import ThreadPoolExecutor, as_completed
from datetime import datetime, timezone

from pycarlo.core import Client, Session
from pycarlo.features.ingestion import IngestionService
from pycarlo.features.ingestion.models import (
    AssetField,
    AssetFreshness,
    AssetMetadata,
    AssetVolume,
    RelationalAsset,
)

logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s")
log = logging.getLogger(__name__)

RESOURCE_TYPE = "bigquery"

# Maximum assets per batch — conservative default to keep compressed payload under 1 MB
# ← SUBSTITUTE: tune based on average asset size (fields per table, description length, etc.)
_BATCH_SIZE = 500


def _asset_from_dict(d: dict) -> RelationalAsset:
    """Reconstruct a RelationalAsset from a manifest dict entry."""
    fields = [
        AssetField(
            name=f["name"],
            type=f.get("type"),
            description=f.get("description"),
        )
        for f in d.get("fields", [])
    ]

    volume = None
    if d.get("volume"):
        volume = AssetVolume(
            row_count=d["volume"].get("row_count"),
            byte_count=d["volume"].get("byte_count"),
        )

    freshness = None
    if d.get("freshness"):
        freshness = AssetFreshness(
            last_update_time=d["freshness"].get("last_update_time"),
        )

    return RelationalAsset(
        type=d.get("type", "TABLE"),
        metadata=AssetMetadata(
            name=d["name"],
            database=d["database"],  # ← SUBSTITUTE: use project or dataset as database
            schema=d["schema"],
            description=d.get("description"),
        ),
        fields=fields,
        volume=volume,
        freshness=freshness,
    )


def push(
    input_file: str,
    resource_uuid: str,
    key_id: str,
    key_token: str,
    batch_size: int = _BATCH_SIZE,
    output_file: str = "metadata_push_result.json",
) -> dict:
    """
    Read a metadata manifest and push assets to Monte Carlo in batches.

    Returns a result dict with invocation IDs for each batch.
    """
    with open(input_file) as fh:
        manifest = json.load(fh)

    asset_dicts = manifest.get("assets", [])
    resource_type = manifest.get("resource_type", RESOURCE_TYPE)
    assets = [_asset_from_dict(d) for d in asset_dicts]
    log.info("Loaded %d asset(s) from %s", len(assets), input_file)

    # Split into batches
    batches = []
    for i in range(0, max(len(assets), 1), batch_size):
        batches.append(assets[i : i + batch_size])
    total_batches = len(batches)

    def _push_batch(batch: list, batch_num: int) -> str | None:
        """Push a single batch using a dedicated Session (thread-safe)."""
        client = Client(session=Session(mcd_id=key_id, mcd_token=key_token, scope="Ingestion"))
        service = IngestionService(mc_client=client)
        result = service.send_metadata(
            resource_uuid=resource_uuid,
            resource_type=resource_type,
            events=batch,
        )
        invocation_id = service.extract_invocation_id(result)
        log.info("Pushed batch %d/%d (%d assets) — invocation_id=%s", batch_num, total_batches, len(batch), invocation_id)
        return invocation_id

    # Push batches in parallel (each thread gets its own pycarlo Session)
    max_workers = min(4, total_batches)
    invocation_ids: list[str | None] = [None] * total_batches

    with ThreadPoolExecutor(max_workers=max_workers) as pool:
        futures = {
            pool.submit(_push_batch, batch, i + 1): i
            for i, batch in enumerate(batches)
        }
        for future in as_completed(futures):
            idx = futures[future]
            try:
                invocation_ids[idx] = future.result()
            except Exception as exc:
                log.error("ERROR pushing batch %d: %s", idx + 1, exc)
                raise

    log.info("All %d batches pushed (%d workers)", total_batches, max_workers)

    push_result = {
        "resource_uuid": resource_uuid,
        "resource_type": resource_type,
        "invocation_ids": invocation_ids,
        "pushed_at": datetime.now(timezone.utc).isoformat(),
        "total_assets": len(assets),
        "batch_count": total_batches,
        "batch_size": batch_size,
    }
    with open(output_file, "w") as fh:
        json.dump(push_result, fh, indent=2)
    log.info("Push result written to %s", output_file)

    return push_result


def main() -> None:
    parser = argparse.ArgumentParser(
        description="Push BigQuery metadata from a manifest to Monte Carlo",
    )
    parser.add_argument("--resource-uuid", default=os.getenv("MCD_RESOURCE_UUID"))
    parser.add_argument("--key-id", default=os.getenv("MCD_INGEST_ID"))
    parser.add_argument("--key-token", default=os.getenv("MCD_INGEST_TOKEN"))
    parser.add_argument("--input-file", default="metadata_output.json")
    parser.add_argument("--output-file", default="metadata_push_result.json")
    parser.add_argument(
        "--batch-size",
        type=int,
        default=_BATCH_SIZE,
        help=f"Max assets per push batch (default: {_BATCH_SIZE})",
    )
    args = parser.parse_args()

    required = ["resource_uuid", "key_id", "key_token"]
    missing = [k for k in required if getattr(args, k) is None]
    if missing:
        parser.error(f"Missing required arguments/env vars: {missing}")

    push(
        input_file=args.input_file,
        resource_uuid=args.resource_uuid,
        key_id=args.key_id,
        key_token=args.key_token,
        batch_size=args.batch_size,
        output_file=args.output_file,
    )


if __name__ == "__main__":
    main()

"""
BigQuery — Query Log Push (push only)
======================================
Reads a manifest file produced by ``collect_query_logs.py`` and pushes the query
log entries to Monte Carlo using the pycarlo push ingestion API.  Large payloads
are split into batches to stay under the 1 MB compressed limit.

Can be run standalone via CLI or imported (use the ``push()`` function).

Substitution points (search for "← SUBSTITUTE"):
  - MCD_INGEST_ID / MCD_INGEST_TOKEN : Monte Carlo API credentials
  - MCD_RESOURCE_UUID      : UUID of the BigQuery connection in Monte Carlo

Prerequisites:
  pip install pycarlo
"""

from __future__ import annotations

import argparse
import json
import logging
import os
from concurrent.futures import ThreadPoolExecutor, as_completed
from datetime import datetime, timezone

from dateutil.parser import isoparse
from pycarlo.core import Client, Session
from pycarlo.features.ingestion import IngestionService
from pycarlo.features.ingestion.models import QueryLogEntry

logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s")
log = logging.getLogger(__name__)

LOG_TYPE = "bigquery"

# Maximum entries per batch — conservative default to keep compressed payload under 1 MB.
# Query logs include full SQL text — keep batches small to stay under the 1 MB
# compressed payload limit.  50 entries can trigger 413 on active warehouses.
# ← SUBSTITUTE: tune based on average query length
_BATCH_SIZE = 100

# Truncate query_text longer than this to prevent 413 errors.
# Some SQL statements (e.g., generated by BI tools) can be 100KB+ and blow up
# compressed payloads even at small batch sizes.
_MAX_QUERY_TEXT_LEN = 10_000


def _build_query_log_entries(queries: list[dict]) -> list[QueryLogEntry]:
    """Convert manifest query dicts into QueryLogEntry objects."""
    entries = []
    truncated = 0
    for q in queries:
        query_text = q.get("query_text") or ""

        # Truncate very long SQL to prevent 413 Request Too Large
        if len(query_text) > _MAX_QUERY_TEXT_LEN:
            query_text = query_text[:_MAX_QUERY_TEXT_LEN] + "... [TRUNCATED]"
            truncated += 1

        extra = {}
        if q.get("total_bytes_billed") is not None:
            extra["total_bytes_billed"] = q["total_bytes_billed"]
        if q.get("statement_type") is not None:
            extra["statement_type"] = q["statement_type"]

        start_time = q.get("start_time")
        end_time = q.get("end_time")

        entry = QueryLogEntry(
            query_id=q.get("query_id"),
            query_text=query_text,
            start_time=isoparse(start_time) if start_time else None,
            end_time=isoparse(end_time) if end_time else None,
            user=q.get("user"),
            extra=extra or None,
        )
        entries.append(entry)
    if truncated:
        log.info("Truncated %d query text(s) exceeding %d chars", truncated, _MAX_QUERY_TEXT_LEN)
    return entries


def push(
    input_file: str,
    resource_uuid: str,
    key_id: str,
    key_token: str,
    batch_size: int = _BATCH_SIZE,
    output_file: str = "query_logs_push_result.json",
) -> dict:
    """
    Read a query log manifest and push entries to Monte Carlo in batches.

    Returns a result dict with invocation IDs for each batch.
    """
    with open(input_file) as fh:
        manifest = json.load(fh)

    queries = manifest.get("queries", [])
    log_type = manifest.get("log_type", LOG_TYPE)
    entries = _build_query_log_entries(queries)
    log.info("Loaded %d query log entry/entries from %s", len(entries), input_file)

    if not entries:
        log.info("No query log entries to push.")
        push_result = {
            "resource_uuid": resource_uuid,
            "log_type": log_type,
            "invocation_ids": [],
            "pushed_at": datetime.now(timezone.utc).isoformat(),
            "total_entries": 0,
            "batch_count": 0,
            "batch_size": batch_size,
        }
        with open(output_file, "w") as fh:
            json.dump(push_result, fh, indent=2)
        return push_result

    # Split into batches
    batches = []
    for i in range(0, len(entries), batch_size):
        batches.append(entries[i : i + batch_size])
    total_batches = len(batches)

    def _push_batch(batch: list, batch_num: int) -> str | None:
        """Push a single batch using a dedicated Session (thread-safe)."""
        client = Client(session=Session(mcd_id=key_id, mcd_token=key_token, scope="Ingestion"))
        service = IngestionService(mc_client=client)
        result = service.send_query_logs(
            resource_uuid=resource_uuid,
            log_type=log_type,
            events=batch,
        )
        invocation_id = service.extract_invocation_id(result)
        log.info("Pushed batch %d/%d (%d entries) — invocation_id=%s", batch_num, total_batches, len(batch), invocation_id)
        return invocation_id

    # Push batches in parallel (each thread gets its own pycarlo Session)
    max_workers = min(4, total_batches)
    invocation_ids: list[str | None] = [None] * total_batches

    with ThreadPoolExecutor(max_workers=max_workers) as pool:
        futures = {
            pool.submit(_push_batch, batch, i + 1): i
            for i, batch in enumerate(batches)
        }
        for future in as_completed(futures):
            idx = futures[future]
            try:
                invocation_ids[idx] = future.result()
            except Exception as exc:
                log.error("ERROR pushing batch %d: %s", idx + 1, exc)
                raise

    log.info("All %d batches pushed (%d workers)", total_batches, max_workers)

    push_result = {
        "resource_uuid": resource_uuid,
        "log_type": log_type,
        "invocation_ids": invocation_ids,
        "pushed_at": datetime.now(timezone.utc).isoformat(),
        "total_entries": len(entries),
        "batch_count": total_batches,
        "batch_size": batch_size,
    }
    with open(output_file, "w") as fh:
        json.dump(push_result, fh, indent=2)
    log.info("Push result written to %s", output_file)

    return push_result


def main() -> None:
    parser = argparse.ArgumentParser(
        description="Push BigQuery query logs from a manifest to Monte Carlo",
    )
    parser.add_argument("--resource-uuid", default=os.getenv("MCD_RESOURCE_UUID"))
    parser.add_argument("--key-id", default=os.getenv("MCD_INGEST_ID"))
    parser.add_argument("--key-token", default=os.getenv("MCD_INGEST_TOKEN"))
    parser.add_argument("--input-file", default="query_logs_output.json")
    parser.add_argument("--output-file", default="query_logs_push_result.json")
    parser.add_argument(
        "--batch-size",
        type=int,
        default=_BATCH_SIZE,
        help=f"Max entries per push batch (default: {_BATCH_SIZE})",
    )
    args = parser.parse_args()

    required = ["resource_uuid", "key_id", "key_token"]
    missing = [k for k in required if getattr(args, k) is None]
    if missing:
        parser.error(f"Missing required arguments/env vars: {missing}")

    push(
        input_file=args.input_file,
        resource_uuid=args.resource_uuid,
        key_id=args.key_id,
        key_token=args.key_token,
        batch_size=args.batch_size,
        output_file=args.output_file,
    )


if __name__ == "__main__":
    main()

"""
BigQuery Iceberg — Metadata Collect & Push (combined)
=====================================================
Convenience wrapper that runs collect_metadata.collect() followed by
push_metadata.push() in a single invocation. Supports
``--only-freshness-and-volume`` for fast periodic pushes.

Prerequisites:
  pip install google-cloud-bigquery pycarlo>=0.12.251
"""

from __future__ import annotations

import argparse
import os

from collect_metadata import collect
from push_metadata import push


def main() -> None:
    parser = argparse.ArgumentParser(
        description="Collect BigQuery Iceberg metadata and push to Monte Carlo",
    )
    # Collection args
    parser.add_argument("--project-id", default=os.getenv("BIGQUERY_PROJECT_ID"))
    parser.add_argument("--datasets", nargs="+", default=None)
    parser.add_argument("--tables", nargs="+", default=None)
    parser.add_argument(
        "--only-freshness-and-volume",
        action="store_true",
        help="Skip field/schema collection — only collect freshness and volume.",
    )
    parser.add_argument("--manifest-file", default="metadata_output.json")

    # Push args
    parser.add_argument("--resource-uuid", default=os.getenv("MCD_RESOURCE_UUID"))
    parser.add_argument("--key-id", default=os.getenv("MCD_INGEST_ID"))
    parser.add_argument("--key-token", default=os.getenv("MCD_INGEST_TOKEN"))
    parser.add_argument("--batch-size", type=int, default=500)
    parser.add_argument("--push-result-file", default="metadata_push_result.json")

    args = parser.parse_args()

    if not args.project_id:
        parser.error("--project-id or BIGQUERY_PROJECT_ID env var is required")
    required_push = ["resource_uuid", "key_id", "key_token"]
    missing = [k for k in required_push if getattr(args, k) is None]
    if missing:
        parser.error(f"Missing required push arguments/env vars: {missing}")

    collect(
        project_id=args.project_id,
        datasets=args.datasets,
        tables=args.tables,
        only_freshness_and_volume=args.only_freshness_and_volume,
        output_file=args.manifest_file,
    )

    push(
        input_file=args.manifest_file,
        resource_uuid=args.resource_uuid,
        key_id=args.key_id,
        key_token=args.key_token,
        batch_size=args.batch_size,
        output_file=args.push_result_file,
    )


if __name__ == "__main__":
    main()

"""
BigQuery Iceberg — Query Log Collect & Push (combined)
=====================================================
Convenience wrapper that runs collect_query_logs.collect() followed by
push_query_logs.push() in a single invocation.

Prerequisites:
  pip install google-cloud-bigquery pycarlo>=0.12.251 python-dateutil>=2.8.0
"""

from __future__ import annotations

import argparse
import os

from collect_query_logs import LOOKBACK_HOURS, LOOKBACK_LAG_HOURS, collect
from push_query_logs import push


def main() -> None:
    parser = argparse.ArgumentParser(
        description="Collect BigQuery query logs and push to Monte Carlo",
    )
    # Collection args
    parser.add_argument("--project-id", default=os.getenv("BIGQUERY_PROJECT_ID"))
    parser.add_argument("--lookback-hours", type=int, default=LOOKBACK_HOURS)
    parser.add_argument("--lookback-lag-hours", type=int, default=LOOKBACK_LAG_HOURS)
    parser.add_argument("--manifest-file", default="query_logs_output.json")

    # Push args
    parser.add_argument("--resource-uuid", default=os.getenv("MCD_RESOURCE_UUID"))
    parser.add_argument("--key-id", default=os.getenv("MCD_INGEST_ID"))
    parser.add_argument("--key-token", default=os.getenv("MCD_INGEST_TOKEN"))
    parser.add_argument("--batch-size", type=int, default=100)
    parser.add_argument("--push-result-file", default="query_logs_push_result.json")

    args = parser.parse_args()

    if not args.project_id:
        parser.error("--project-id or BIGQUERY_PROJECT_ID env var is required")
    required_push = ["resource_uuid", "key_id", "key_token"]
    missing = [k for k in required_push if getattr(args, k) is None]
    if missing:
        parser.error(f"Missing required push arguments/env vars: {missing}")

    collect(
        project_id=args.project_id,
        lookback_hours=args.lookback_hours,
        lookback_lag_hours=args.lookback_lag_hours,
        output_file=args.manifest_file,
    )

    push(
        input_file=args.manifest_file,
        resource_uuid=args.resource_uuid,
        key_id=args.key_id,
        key_token=args.key_token,
        batch_size=args.batch_size,
        output_file=args.push_result_file,
    )


if __name__ == "__main__":
    main()

"""
BigQuery Iceberg — Metadata Collection (collect only)
=====================================================
Collects table schemas, row counts, byte sizes, and freshness for BigQuery
Iceberg (BigLake-managed) tables using INFORMATION_SCHEMA.TABLE_STORAGE and
INFORMATION_SCHEMA.COLUMNS. Standard BigQuery collection uses __TABLES__ which
does not include Iceberg tables — this template fills that gap.

Can be run standalone via CLI or imported (use the ``collect()`` function).

Supports a ``--only-freshness-and-volume`` flag to skip the COLUMNS query for
fast periodic pushes after the initial full metadata push.

Substitution points (search for "← SUBSTITUTE"):
  - BIGQUERY_PROJECT_ID                : GCP project ID to collect from
  - GOOGLE_APPLICATION_CREDENTIALS     : path to service-account JSON key file
  - REGION                             : BigQuery region (default "us")

Prerequisites:
  pip install google-cloud-bigquery
"""

from __future__ import annotations

import argparse
import json
import logging
import os
from datetime import datetime, timezone

from google.cloud import bigquery

logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s")
log = logging.getLogger(__name__)

RESOURCE_TYPE = "bigquery"

# BigQuery type → Monte Carlo canonical type
BQ_TYPE_MAP: dict[str, str] = {
    "INT64": "INTEGER",
    "INTEGER": "INTEGER",
    "FLOAT64": "FLOAT",
    "FLOAT": "FLOAT",
    "BOOL": "BOOLEAN",
    "BOOLEAN": "BOOLEAN",
    "STRING": "VARCHAR",
    "BYTES": "BINARY",
    "DATE": "DATE",
    "DATETIME": "DATETIME",
    "TIMESTAMP": "TIMESTAMP",
    "TIME": "TIME",
    "NUMERIC": "DECIMAL",
    "BIGNUMERIC": "DECIMAL",
    "RECORD": "STRUCT",
    "STRUCT": "STRUCT",
    "REPEATED": "ARRAY",
    "JSON": "JSON",
    "GEOGRAPHY": "GEOGRAPHY",
}


def map_bq_type(bq_type: str) -> str:
    base = bq_type.split("(")[0].strip().upper()
    return BQ_TYPE_MAP.get(base, bq_type.upper())


def _fetch_iceberg_tables(
    client: bigquery.Client,
    project_id: str,
    datasets: list[str] | None = None,
    tables: list[str] | None = None,
) -> list[dict]:
    """Query TABLE_STORAGE for BigLake (Iceberg) tables."""
    conditions = [
        "managed_table_type = 'BIGLAKE'",
        "deleted = FALSE",
    ]
    if datasets:
        ds_list = ", ".join(f"'{d}'" for d in datasets)
        conditions.append(f"table_schema IN ({ds_list})")
    if tables:
        tbl_list = ", ".join(f"'{t}'" for t in tables)
        conditions.append(f"table_name IN ({tbl_list})")

    where = " AND ".join(conditions)
    query = f"""
        SELECT
            table_schema,
            table_name,
            total_rows,
            current_physical_bytes,
            storage_last_modified_time,
            creation_time
        FROM `{project_id}.region-us`.INFORMATION_SCHEMA.TABLE_STORAGE  -- ← SUBSTITUTE: change region if needed
        WHERE {where}
        ORDER BY table_schema, table_name
    """
    log.info("Querying TABLE_STORAGE for Iceberg tables ...")
    rows = list(client.query(query).result())
    log.info("Found %d Iceberg table(s).", len(rows))
    return [dict(row) for row in rows]


def _fetch_columns(
    client: bigquery.Client,
    project_id: str,
    dataset: str,
    table_name: str,
) -> list[dict]:
    """Fetch column metadata for a specific table."""
    query = f"""
        SELECT column_name, data_type, ordinal_position, is_nullable, column_default
        FROM `{project_id}.{dataset}.INFORMATION_SCHEMA.COLUMNS`
        WHERE table_name = '{table_name}'
        ORDER BY ordinal_position
    """
    return [
        {
            "name": row["column_name"],
            "type": map_bq_type(row["data_type"]),
        }
        for row in client.query(query).result()
    ]


def _resolve_freshness(row: dict) -> str:
    """Return the best available freshness timestamp as ISO8601.

    Uses storage_last_modified_time if Google has populated it (expected
    early April 2026). Falls back to current time with a warning.
    """
    if row.get("storage_last_modified_time"):
        return row["storage_last_modified_time"].isoformat()

    log.warning(
        "storage_last_modified_time is NULL for %s.%s — "
        "falling back to current time. Google's TABLE_STORAGE update "
        "for Iceberg tables may not have shipped yet.",
        row["table_schema"],
        row["table_name"],
    )
    return datetime.now(timezone.utc).isoformat()


def collect(
    project_id: str,
    datasets: list[str] | None = None,
    tables: list[str] | None = None,
    only_freshness_and_volume: bool = False,
    output_file: str = "metadata_output.json",
) -> dict:
    """Collect Iceberg table metadata and write a JSON manifest.

    When only_freshness_and_volume is True, skips the COLUMNS query and
    omits fields from the manifest. Use this for periodic hourly pushes
    after the initial full metadata push.
    """
    client = bigquery.Client(project=project_id)  # ← SUBSTITUTE: adjust auth if needed

    if only_freshness_and_volume:
        log.info("Running in freshness+volume only mode (skipping fields).")

    iceberg_tables = _fetch_iceberg_tables(client, project_id, datasets, tables)
    if not iceberg_tables:
        log.warning("No Iceberg tables found matching the criteria.")
        return {"resource_type": RESOURCE_TYPE, "assets": []}

    assets: list[dict] = []
    for row in iceberg_tables:
        dataset = row["table_schema"]
        name = row["table_name"]

        asset = {
            "name": name,
            "database": project_id,
            "schema": dataset,
            "type": "TABLE",
            "volume": {
                "row_count": row["total_rows"],
                "byte_count": row["current_physical_bytes"],
            },
            "freshness": {
                "last_updated_time": _resolve_freshness(row),
            },
        }

        if not only_freshness_and_volume:
            asset["description"] = None
            asset["fields"] = _fetch_columns(client, project_id, dataset, name)

        assets.append(asset)
        log.info(
            "Collected %s.%s.%s — rows=%s, bytes=%s",
            project_id, dataset, name,
            row["total_rows"], row["current_physical_bytes"],
        )

    manifest = {
        "resource_type": RESOURCE_TYPE,
        "collected_at": datetime.now(timezone.utc).isoformat(),
        "assets": assets,
    }
    with open(output_file, "w") as fh:
        json.dump(manifest, fh, indent=2)
    log.info("Manifest written to %s (%d assets)", output_file, len(assets))

    return manifest


def main() -> None:
    parser = argparse.ArgumentParser(
        description="Collect BigQuery Iceberg table metadata into a JSON manifest",
    )
    parser.add_argument(
        "--project-id",
        default=os.getenv("BIGQUERY_PROJECT_ID"),  # ← SUBSTITUTE
        help="GCP project ID (or set BIGQUERY_PROJECT_ID env var)",
    )
    parser.add_argument(
        "--datasets",
        nargs="+",
        default=None,
        help="Limit to specific dataset(s). Omit to scan all datasets.",
    )
    parser.add_argument(
        "--tables",
        nargs="+",
        default=None,
        help="Limit to specific table name(s) within the datasets.",
    )
    parser.add_argument(
        "--only-freshness-and-volume",
        action="store_true",
        help="Skip field/schema collection — only collect freshness and volume. "
             "Use for periodic hourly pushes after the initial full metadata push.",
    )
    parser.add_argument("--output-file", default="metadata_output.json")
    args = parser.parse_args()

    if not args.project_id:
        parser.error("--project-id or BIGQUERY_PROJECT_ID env var is required")

    collect(
        project_id=args.project_id,
        datasets=args.datasets,
        tables=args.tables,
        only_freshness_and_volume=args.only_freshness_and_volume,
        output_file=args.output_file,
    )


if __name__ == "__main__":
    main()

"""
BigQuery Iceberg — Query Log Collection (collect only)
======================================================
Queries the BigQuery Jobs API for completed query jobs within a time
window and writes a JSON manifest that can be fed to push_query_logs.py.

Can be run standalone via CLI or imported (use the ``collect()`` function).

Substitution points (search for "← SUBSTITUTE"):
  - BIGQUERY_PROJECT_ID                : GCP project ID to collect from
  - GOOGLE_APPLICATION_CREDENTIALS     : path to service-account JSON key file

Prerequisites:
  pip install google-cloud-bigquery
"""

from __future__ import annotations

import argparse
import json
import logging
import os
from datetime import datetime, timedelta, timezone

from google.cloud import bigquery

logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s")
log = logging.getLogger(__name__)

LOG_TYPE = "bigquery"

LOOKBACK_HOURS: int = int(os.getenv("LOOKBACK_HOURS", "25"))
LOOKBACK_LAG_HOURS: int = int(os.getenv("LOOKBACK_LAG_HOURS", "1"))
MAX_JOBS: int = int(os.getenv("MAX_JOBS", "10000"))

# Limit to specific statement types — empty list means collect all.
STATEMENT_TYPE_FILTER: list[str] = []


def _safe_isoformat(dt: datetime | None) -> str | None:
    if dt is None:
        return None
    if dt.tzinfo is None:
        dt = dt.replace(tzinfo=timezone.utc)
    return dt.isoformat()


def _collect_query_logs(
    bq_client: bigquery.Client,
    project_id: str,
    start_dt: datetime,
    end_dt: datetime,
) -> list[dict]:
    """Collect query logs from BigQuery job history."""
    entries: list[dict] = []

    log.info(
        "Listing jobs for project=%s from %s to %s",
        project_id, start_dt.isoformat(), end_dt.isoformat(),
    )

    for job in bq_client.list_jobs(
        project=project_id,
        all_users=True,
        min_creation_time=start_dt,
        max_creation_time=end_dt,
    ):
        sql: str = getattr(job, "query", None) or ""
        if not sql.strip():
            continue

        statement_type: str = getattr(job, "statement_type", None) or ""
        if STATEMENT_TYPE_FILTER and statement_type not in STATEMENT_TYPE_FILTER:
            continue

        entries.append({
            "query_id": job.job_id,
            "query_text": sql,
            "start_time": _safe_isoformat(getattr(job, "created", None)),
            "end_time": _safe_isoformat(getattr(job, "ended", None)),
            "user": getattr(job, "user_email", None),
            "total_bytes_billed": getattr(job, "total_bytes_billed", None),
            "statement_type": statement_type or None,
        })

        if len(entries) >= MAX_JOBS:
            log.warning("Reached MAX_JOBS=%d — stopping early", MAX_JOBS)
            break

    return entries


def collect(
    project_id: str,
    lookback_hours: int = LOOKBACK_HOURS,
    lookback_lag_hours: int = LOOKBACK_LAG_HOURS,
    output_file: str = "query_logs_output.json",
) -> dict:
    """Collect query logs and write a JSON manifest."""
    bq_client = bigquery.Client(project=project_id)

    end_dt = datetime.now(timezone.utc) - timedelta(hours=lookback_lag_hours)
    start_dt = end_dt - timedelta(hours=lookback_hours)

    entries = _collect_query_logs(bq_client, project_id, start_dt, end_dt)
    log.info("Collected %d query log entries.", len(entries))

    manifest = {
        "log_type": LOG_TYPE,
        "collected_at": datetime.now(timezone.utc).isoformat(),
        "window_start": start_dt.isoformat(),
        "window_end": end_dt.isoformat(),
        "query_log_count": len(entries),
        "queries": entries,
    }
    with open(output_file, "w") as fh:
        json.dump(manifest, fh, indent=2)
    log.info("Query log manifest written to %s", output_file)

    return manifest


def main() -> None:
    parser = argparse.ArgumentParser(
        description="Collect BigQuery query logs into a JSON manifest",
    )
    parser.add_argument(
        "--project-id",
        default=os.getenv("BIGQUERY_PROJECT_ID"),
        help="GCP project ID (or set BIGQUERY_PROJECT_ID env var)",
    )
    parser.add_argument("--lookback-hours", type=int, default=LOOKBACK_HOURS)
    parser.add_argument("--lookback-lag-hours", type=int, default=LOOKBACK_LAG_HOURS)
    parser.add_argument("--output-file", default="query_logs_output.json")
    args = parser.parse_args()

    if not args.project_id:
        parser.error("--project-id or BIGQUERY_PROJECT_ID env var is required")

    collect(
        project_id=args.project_id,
        lookback_hours=args.lookback_hours,
        lookback_lag_hours=args.lookback_lag_hours,
        output_file=args.output_file,
    )


if __name__ == "__main__":
    main()

"""
BigQuery Iceberg — Metadata Push (push only)
============================================
Reads a JSON manifest produced by collect_metadata.py and pushes table
metadata to Monte Carlo using the pycarlo SDK's IngestionService.

Can be run standalone via CLI or imported (use the ``push()`` function).

Substitution points (search for "← SUBSTITUTE"):
  - MCD_INGEST_ID      : Monte Carlo Ingestion API key ID
  - MCD_INGEST_TOKEN   : Monte Carlo Ingestion API key token
  - MCD_RESOURCE_UUID  : Monte Carlo warehouse resource UUID

Prerequisites:
  pip install pycarlo>=0.12.251
"""

from __future__ import annotations

import argparse
import json
import logging
import os
from concurrent.futures import ThreadPoolExecutor, as_completed
from datetime import datetime, timezone

from pycarlo.core import Client, Session
from pycarlo.features.ingestion import IngestionService
from pycarlo.features.ingestion.models import (
    AssetField,
    AssetFreshness,
    AssetMetadata,
    AssetVolume,
    RelationalAsset,
)

logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s")
log = logging.getLogger(__name__)

RESOURCE_TYPE = "bigquery"
_BATCH_SIZE = 500

_ENDPOINT = "https://integrations.getmontecarlo.com"


def _asset_from_dict(d: dict) -> RelationalAsset:
    """Reconstruct a RelationalAsset from a manifest dict entry."""
    fields = [
        AssetField(
            name=f["name"],
            type=f.get("type"),
            description=f.get("description"),
        )
        for f in d.get("fields", [])
    ]

    volume = None
    if d.get("volume"):
        volume = AssetVolume(
            row_count=d["volume"].get("row_count"),
            byte_count=d["volume"].get("byte_count"),
        )

    freshness = None
    if d.get("freshness") and d["freshness"].get("last_updated_time"):
        freshness = AssetFreshness(
            last_update_time=d["freshness"]["last_updated_time"],
        )

    return RelationalAsset(
        type=d.get("type", "TABLE"),
        metadata=AssetMetadata(
            name=d["name"],
            database=d["database"],
            schema=d["schema"],
            description=d.get("description"),
        ),
        fields=fields,
        volume=volume,
        freshness=freshness,
    )


def push(
    input_file: str,
    resource_uuid: str,
    key_id: str,
    key_token: str,
    batch_size: int = _BATCH_SIZE,
    output_file: str = "metadata_push_result.json",
) -> dict:
    """Read a metadata manifest and push assets to Monte Carlo in batches."""
    endpoint = _ENDPOINT
    log.info("Using endpoint: %s", endpoint)
    with open(input_file) as fh:
        manifest = json.load(fh)

    asset_dicts = manifest.get("assets", [])
    resource_type = manifest.get("resource_type", RESOURCE_TYPE)
    assets = [_asset_from_dict(d) for d in asset_dicts]
    log.info("Loaded %d asset(s) from %s", len(assets), input_file)

    batches = [assets[i : i + batch_size] for i in range(0, max(len(assets), 1), batch_size)]
    total_batches = len(batches)

    def _push_batch(batch: list[RelationalAsset], batch_num: int) -> str | None:
        client = Client(session=Session(
            mcd_id=key_id, mcd_token=key_token, scope="Ingestion", endpoint=endpoint,
        ))
        service = IngestionService(mc_client=client)
        result = service.send_metadata(
            resource_uuid=resource_uuid,
            resource_type=resource_type,
            events=batch,
        )
        invocation_id = service.extract_invocation_id(result)
        log.info(
            "Pushed batch %d/%d (%d assets) — invocation_id=%s",
            batch_num, total_batches, len(batch), invocation_id,
        )
        return invocation_id

    max_workers = min(4, total_batches)
    invocation_ids: list[str | None] = [None] * total_batches

    with ThreadPoolExecutor(max_workers=max_workers) as pool:
        futures = {
            pool.submit(_push_batch, batch, i + 1): i
            for i, batch in enumerate(batches)
        }
        for future in as_completed(futures):
            idx = futures[future]
            try:
                invocation_ids[idx] = future.result()
            except Exception as exc:
                log.error("ERROR pushing batch %d: %s", idx + 1, exc)
                raise

    log.info("All %d batch(es) pushed.", total_batches)

    push_result = {
        "resource_uuid": resource_uuid,
        "resource_type": resource_type,
        "invocation_ids": invocation_ids,
        "pushed_at": datetime.now(timezone.utc).isoformat(),
        "total_assets": len(assets),
        "batch_count": total_batches,
        "batch_size": batch_size,
    }
    with open(output_file, "w") as fh:
        json.dump(push_result, fh, indent=2)
    log.info("Push result written to %s", output_file)

    return push_result


def main() -> None:
    parser = argparse.ArgumentParser(
        description="Push BigQuery Iceberg metadata from a manifest to Monte Carlo",
    )
    parser.add_argument("--resource-uuid", default=os.getenv("MCD_RESOURCE_UUID"))
    parser.add_argument("--key-id", default=os.getenv("MCD_INGEST_ID"))
    parser.add_argument("--key-token", default=os.getenv("MCD_INGEST_TOKEN"))
    parser.add_argument("--input-file", default="metadata_output.json")
    parser.add_argument("--output-file", default="metadata_push_result.json")
    parser.add_argument(
        "--batch-size",
        type=int,
        default=_BATCH_SIZE,
        help=f"Max assets per push batch (default: {_BATCH_SIZE})",
    )
    args = parser.parse_args()

    required = ["resource_uuid", "key_id", "key_token"]
    missing = [k for k in required if getattr(args, k) is None]
    if missing:
        parser.error(f"Missing required arguments/env vars: {missing}")

    push(
        input_file=args.input_file,
        resource_uuid=args.resource_uuid,
        key_id=args.key_id,
        key_token=args.key_token,
        batch_size=args.batch_size,
        output_file=args.output_file,
    )


if __name__ == "__main__":
    main()

"""
BigQuery Iceberg — Query Log Push (push only)
=============================================
Reads a JSON manifest produced by collect_query_logs.py and pushes query
log entries to Monte Carlo using the pycarlo SDK's IngestionService.

Uses dateutil.isoparse() to convert ISO8601 strings back to datetime
objects (QueryLogEntry requires datetime, not str).

Can be run standalone via CLI or imported (use the ``push()`` function).

Substitution points (search for "← SUBSTITUTE"):
  - MCD_INGEST_ID      : Monte Carlo Ingestion API key ID
  - MCD_INGEST_TOKEN   : Monte Carlo Ingestion API key token
  - MCD_RESOURCE_UUID  : Monte Carlo warehouse resource UUID

Prerequisites:
  pip install pycarlo>=0.12.251 python-dateutil>=2.8.0
"""

from __future__ import annotations

import argparse
import json
import logging
import os
from concurrent.futures import ThreadPoolExecutor, as_completed
from datetime import datetime, timezone

from dateutil.parser import isoparse

from pycarlo.core import Client, Session
from pycarlo.features.ingestion import IngestionService
from pycarlo.features.ingestion.models import QueryLogEntry

logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s")
log = logging.getLogger(__name__)

LOG_TYPE = "bigquery"

# Query logs include full SQL text — keep batches small to stay under the
# 1 MB compressed payload limit.
_BATCH_SIZE = 100

# Truncate very long SQL to prevent 413 errors.
_MAX_QUERY_TEXT_LEN = 10_000

_ENDPOINT = "https://integrations.getmontecarlo.com"


def _build_query_log_entries(queries: list[dict]) -> list[QueryLogEntry]:
    """Convert manifest query dicts into QueryLogEntry objects."""
    entries = []
    truncated = 0
    for q in queries:
        query_text = q.get("query_text") or ""

        if len(query_text) > _MAX_QUERY_TEXT_LEN:
            query_text = query_text[:_MAX_QUERY_TEXT_LEN] + "... [TRUNCATED]"
            truncated += 1

        extra = {}
        if q.get("total_bytes_billed") is not None:
            extra["total_bytes_billed"] = q["total_bytes_billed"]
        if q.get("statement_type") is not None:
            extra["statement_type"] = q["statement_type"]

        start_time = q.get("start_time")
        end_time = q.get("end_time")

        entry = QueryLogEntry(
            query_id=q.get("query_id"),
            query_text=query_text,
            start_time=isoparse(start_time) if start_time else None,
            end_time=isoparse(end_time) if end_time else None,
            user=q.get("user"),
            extra=extra or None,
        )
        entries.append(entry)

    if truncated:
        log.info("Truncated %d query text(s) exceeding %d chars", truncated, _MAX_QUERY_TEXT_LEN)
    return entries


def push(
    input_file: str,
    resource_uuid: str,
    key_id: str,
    key_token: str,
    batch_size: int = _BATCH_SIZE,
    output_file: str = "query_logs_push_result.json",
) -> dict:
    """Read a query log manifest and push entries to Monte Carlo in batches."""
    endpoint = _ENDPOINT
    log.info("Using endpoint: %s", endpoint)

    with open(input_file) as fh:
        manifest = json.load(fh)

    queries = manifest.get("queries", [])
    log_type = manifest.get("log_type", LOG_TYPE)
    entries = _build_query_log_entries(queries)
    log.info("Loaded %d query log entry/entries from %s", len(entries), input_file)

    if not entries:
        log.info("No query log entries to push.")
        push_result = {
            "resource_uuid": resource_uuid,
            "log_type": log_type,
            "invocation_ids": [],
            "pushed_at": datetime.now(timezone.utc).isoformat(),
            "total_entries": 0,
            "batch_count": 0,
            "batch_size": batch_size,
        }
        with open(output_file, "w") as fh:
            json.dump(push_result, fh, indent=2)
        return push_result

    batches = [entries[i : i + batch_size] for i in range(0, len(entries), batch_size)]
    total_batches = len(batches)

    def _push_batch(batch: list[QueryLogEntry], batch_num: int) -> str | None:
        client = Client(session=Session(
            mcd_id=key_id, mcd_token=key_token, scope="Ingestion", endpoint=endpoint,
        ))
        service = IngestionService(mc_client=client)
        result = service.send_query_logs(
            resource_uuid=resource_uuid,
            log_type=log_type,
            events=batch,
        )
        invocation_id = service.extract_invocation_id(result)
        log.info(
            "Pushed batch %d/%d (%d entries) — invocation_id=%s",
            batch_num, total_batches, len(batch), invocation_id,
        )
        return invocation_id

    max_workers = min(4, total_batches)
    invocation_ids: list[str | None] = [None] * total_batches

    with ThreadPoolExecutor(max_workers=max_workers) as pool:
        futures = {
            pool.submit(_push_batch, batch, i + 1): i
            for i, batch in enumerate(batches)
        }
        for future in as_completed(futures):
            idx = futures[future]
            try:
                invocation_ids[idx] = future.result()
            except Exception as exc:
                log.error("ERROR pushing batch %d: %s", idx + 1, exc)
                raise

    log.info("All %d batch(es) pushed.", total_batches)

    push_result = {
        "resource_uuid": resource_uuid,
        "log_type": log_type,
        "invocation_ids": invocation_ids,
        "pushed_at": datetime.now(timezone.utc).isoformat(),
        "total_entries": len(entries),
        "batch_count": total_batches,
        "batch_size": batch_size,
    }
    with open(output_file, "w") as fh:
        json.dump(push_result, fh, indent=2)
    log.info("Push result written to %s", output_file)

    return push_result


def main() -> None:
    parser = argparse.ArgumentParser(
        description="Push BigQuery query logs from a manifest to Monte Carlo",
    )
    parser.add_argument("--resource-uuid", default=os.getenv("MCD_RESOURCE_UUID"))
    parser.add_argument("--key-id", default=os.getenv("MCD_INGEST_ID"))
    parser.add_argument("--key-token", default=os.getenv("MCD_INGEST_TOKEN"))
    parser.add_argument("--input-file", default="query_logs_output.json")
    parser.add_argument("--output-file", default="query_logs_push_result.json")
    parser.add_argument(
        "--batch-size",
        type=int,
        default=_BATCH_SIZE,
        help=f"Max entries per push batch (default: {_BATCH_SIZE})",
    )
    args = parser.parse_args()

    required = ["resource_uuid", "key_id", "key_token"]
    missing = [k for k in required if getattr(args, k) is None]
    if missing:
        parser.error(f"Missing required arguments/env vars: {missing}")

    push(
        input_file=args.input_file,
        resource_uuid=args.resource_uuid,
        key_id=args.key_id,
        key_token=args.key_token,
        batch_size=args.batch_size,
        output_file=args.output_file,
    )


if __name__ == "__main__":
    main()

"""
Databricks — Lineage Collect & Push (combined)
================================================
Collects table-level and (optionally) column-level lineage from Databricks Unity
Catalog system tables, then pushes them to Monte Carlo via the push ingestion API.

This script imports and calls collect() from collect_lineage and push() from
push_lineage, running both in sequence.

Substitution points (search for "← SUBSTITUTE"):
  - DATABRICKS_HOST       : workspace hostname
  - DATABRICKS_HTTP_PATH  : SQL warehouse HTTP path
  - DATABRICKS_TOKEN      : PAT or service-principal secret
  - LOOKBACK_DAYS         : how many days back to collect lineage (default 30)
  - MCD_INGEST_ID / MCD_INGEST_TOKEN : Monte Carlo API credentials
  - MCD_RESOURCE_UUID      : UUID of the Databricks connection in Monte Carlo
  - PUSH_BATCH_SIZE       : number of events per API call (default 500)

Use the --column-lineage flag to also push column-level lineage (disabled by default).

Prerequisites:
  pip install databricks-sql-connector pycarlo
"""

from __future__ import annotations

import argparse
import logging
import os

from collect_lineage import LOOKBACK_DAYS, collect
from push_lineage import DEFAULT_BATCH_SIZE, push

logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s")
log = logging.getLogger(__name__)


def main() -> None:
    parser = argparse.ArgumentParser(description="Collect and push Databricks lineage to Monte Carlo")
    parser.add_argument("--host", default=os.getenv("DATABRICKS_HOST"))           # ← SUBSTITUTE
    parser.add_argument("--http-path", default=os.getenv("DATABRICKS_HTTP_PATH")) # ← SUBSTITUTE
    parser.add_argument("--token", default=os.getenv("DATABRICKS_TOKEN"))         # ← SUBSTITUTE
    parser.add_argument("--resource-uuid", default=os.getenv("MCD_RESOURCE_UUID"))
    parser.add_argument("--key-id", default=os.getenv("MCD_INGEST_ID"))
    parser.add_argument("--key-token", default=os.getenv("MCD_INGEST_TOKEN"))
    parser.add_argument("--lookback-days", type=int, default=LOOKBACK_DAYS)
    parser.add_argument(
        "--column-lineage", action="store_true",
        help="Also collect column-level lineage (requires system.access.column_lineage access)",
    )
    parser.add_argument("--batch-size", type=int, default=DEFAULT_BATCH_SIZE)
    parser.add_argument("--manifest", default="manifest_lineage.json")
    args = parser.parse_args()

    required = ["host", "http_path", "token", "resource_uuid", "key_id", "key_token"]
    missing = [k for k in required if getattr(args, k) is None]
    if missing:
        parser.error(f"Missing required arguments/env vars: {missing}")

    log.info("Step 1: Collecting lineage …")
    collect(
        host=args.host,
        http_path=args.http_path,
        token=args.token,
        manifest_path=args.manifest,
        include_column_lineage=args.column_lineage,
        lookback_days=args.lookback_days,
    )

    log.info("Step 2: Pushing lineage to Monte Carlo …")
    push(
        manifest_path=args.manifest,
        resource_uuid=args.resource_uuid,
        key_id=args.key_id,
        key_token=args.key_token,
        batch_size=args.batch_size,
    )

    log.info("Done — collect and push complete.")


if __name__ == "__main__":
    main()

"""
Databricks — Metadata Collect & Push (combined)
=================================================
Collects table schemas, row counts, and byte sizes from Databricks Unity Catalog,
then pushes them to Monte Carlo via the push ingestion API.

This script imports and calls collect() from collect_metadata and push() from
push_metadata, running both in sequence.

Substitution points (search for "← SUBSTITUTE"):
  - DATABRICKS_HOST       : workspace hostname (e.g. adb-1234.azuredatabricks.net)
  - DATABRICKS_HTTP_PATH  : SQL warehouse HTTP path (e.g. /sql/1.0/warehouses/abc123)
  - DATABRICKS_TOKEN      : personal access token or service-principal secret
  - DATABRICKS_CATALOG    : catalog to collect from (default: "hive_metastore" or "main")
  - SCHEMA_EXCLUSIONS     : schemas to skip
  - MCD_INGEST_ID / MCD_INGEST_TOKEN : Monte Carlo API credentials
  - MCD_RESOURCE_UUID      : UUID of the Databricks connection in Monte Carlo
  - PUSH_BATCH_SIZE       : number of assets per API call (default 500)

Prerequisites:
  pip install databricks-sql-connector pycarlo
"""

from __future__ import annotations

import argparse
import logging
import os

from collect_metadata import collect
from push_metadata import DEFAULT_BATCH_SIZE, push

logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s")
log = logging.getLogger(__name__)


def main() -> None:
    parser = argparse.ArgumentParser(description="Collect and push Databricks metadata to Monte Carlo")
    parser.add_argument("--host", default=os.getenv("DATABRICKS_HOST"))           # ← SUBSTITUTE
    parser.add_argument("--http-path", default=os.getenv("DATABRICKS_HTTP_PATH")) # ← SUBSTITUTE
    parser.add_argument("--token", default=os.getenv("DATABRICKS_TOKEN"))         # ← SUBSTITUTE
    parser.add_argument("--catalog", default=os.getenv("DATABRICKS_CATALOG", "hive_metastore"))
    parser.add_argument("--resource-uuid", default=os.getenv("MCD_RESOURCE_UUID"))
    parser.add_argument("--key-id", default=os.getenv("MCD_INGEST_ID"))
    parser.add_argument("--key-token", default=os.getenv("MCD_INGEST_TOKEN"))
    parser.add_argument("--batch-size", type=int, default=DEFAULT_BATCH_SIZE)
    parser.add_argument("--manifest", default="manifest_metadata.json")
    args = parser.parse_args()

    required = ["host", "http_path", "token", "resource_uuid", "key_id", "key_token"]
    missing = [k for k in required if getattr(args, k) is None]
    if missing:
        parser.error(f"Missing required arguments/env vars: {missing}")

    log.info("Step 1: Collecting metadata …")
    collect(
        host=args.host,
        http_path=args.http_path,
        token=args.token,
        catalog=args.catalog,
        manifest_path=args.manifest,
    )

    log.info("Step 2: Pushing metadata to Monte Carlo …")
    push(
        manifest_path=args.manifest,
        resource_uuid=args.resource_uuid,
        key_id=args.key_id,
        key_token=args.key_token,
        batch_size=args.batch_size,
    )

    log.info("Done — collect and push complete.")


if __name__ == "__main__":
    main()

"""
Databricks — Query Log Collect & Push (combined)
==================================================
Collects finished query execution records from the Databricks system table
system.query.history and pushes them to Monte Carlo for query-pattern analysis,
lineage derivation, and usage attribution.

This script imports and calls collect() from collect_query_logs and push() from
push_query_logs, running both in sequence.

Substitution points (search for "← SUBSTITUTE"):
  - DATABRICKS_HOST       : workspace hostname
  - DATABRICKS_HTTP_PATH  : SQL warehouse HTTP path
  - DATABRICKS_TOKEN      : PAT or service-principal secret
  - LOOKBACK_HOURS        : hours back from [now - LAG_HOURS] to collect (default 25)
  - LOOKBACK_LAG_HOURS    : hours to lag behind now to avoid in-flight queries (default 1)
  - MAX_ROWS              : maximum query rows to collect per run (default 10000)
  - MCD_INGEST_ID / MCD_INGEST_TOKEN : Monte Carlo API credentials
  - MCD_RESOURCE_UUID      : UUID of the Databricks connection in Monte Carlo
  - PUSH_BATCH_SIZE       : number of entries per API call (default 250)

Prerequisites:
  pip install databricks-sql-connector pycarlo
"""

from __future__ import annotations

import argparse
import logging
import os

from collect_query_logs import LOOKBACK_HOURS, LOOKBACK_LAG_HOURS, MAX_ROWS, collect
from push_query_logs import DEFAULT_BATCH_SIZE, push

logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s")
log = logging.getLogger(__name__)


def main() -> None:
    parser = argparse.ArgumentParser(description="Collect and push Databricks query logs to Monte Carlo")
    parser.add_argument("--host", default=os.getenv("DATABRICKS_HOST"))           # ← SUBSTITUTE
    parser.add_argument("--http-path", default=os.getenv("DATABRICKS_HTTP_PATH")) # ← SUBSTITUTE
    parser.add_argument("--token", default=os.getenv("DATABRICKS_TOKEN"))         # ← SUBSTITUTE
    parser.add_argument("--resource-uuid", default=os.getenv("MCD_RESOURCE_UUID"))
    parser.add_argument("--key-id", default=os.getenv("MCD_INGEST_ID"))
    parser.add_argument("--key-token", default=os.getenv("MCD_INGEST_TOKEN"))
    parser.add_argument("--lookback-hours", type=int, default=LOOKBACK_HOURS)
    parser.add_argument("--lookback-lag-hours", type=int, default=LOOKBACK_LAG_HOURS)
    parser.add_argument("--max-rows", type=int, default=MAX_ROWS)
    parser.add_argument("--batch-size", type=int, default=DEFAULT_BATCH_SIZE)
    parser.add_argument("--manifest", default="manifest_query_logs.json")
    args = parser.parse_args()

    required = ["host", "http_path", "token", "resource_uuid", "key_id", "key_token"]
    missing = [k for k in required if getattr(args, k) is None]
    if missing:
        parser.error(f"Missing required arguments/env vars: {missing}")

    log.info("Step 1: Collecting query logs …")
    collect(
        host=args.host,
        http_path=args.http_path,
        token=args.token,
        manifest_path=args.manifest,
        lookback_hours=args.lookback_hours,
        lookback_lag_hours=args.lookback_lag_hours,
        max_rows=args.max_rows,
    )

    log.info("Step 2: Pushing query logs to Monte Carlo …")
    push(
        manifest_path=args.manifest,
        resource_uuid=args.resource_uuid,
        key_id=args.key_id,
        key_token=args.key_token,
        batch_size=args.batch_size,
    )

    log.info("Done — collect and push complete.")


if __name__ == "__main__":
    main()

"""
Databricks — Lineage Collection (collect-only)
================================================
Collects table-level and (optionally) column-level lineage from Databricks Unity
Catalog system tables (system.access.table_lineage and system.access.column_lineage).
No SQL parsing required — Databricks provides first-class lineage metadata.

Writes a JSON manifest file that can be consumed by push_lineage.py.

Substitution points (search for "← SUBSTITUTE"):
  - DATABRICKS_HOST       : workspace hostname
  - DATABRICKS_HTTP_PATH  : SQL warehouse HTTP path
  - DATABRICKS_TOKEN      : PAT or service-principal secret
  - LOOKBACK_DAYS         : how many days back to collect lineage (default 30)

Use the --column-lineage flag to also collect column-level lineage (disabled by default).

Prerequisites:
  pip install databricks-sql-connector
"""

from __future__ import annotations

import argparse
import json
import logging
import os
from datetime import datetime, timezone
from typing import Any

from databricks import sql

logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s")
log = logging.getLogger(__name__)

RESOURCE_TYPE = "databricks"
LOOKBACK_DAYS: int = int(os.getenv("LOOKBACK_DAYS", "30"))  # ← SUBSTITUTE


def _check_available_memory(min_gb: float = 2.0) -> None:
    """Warn if available memory is below the threshold."""
    try:
        if hasattr(os, "sysconf"):  # Linux / macOS
            page_size = os.sysconf("SC_PAGE_SIZE")
            avail_pages = os.sysconf("SC_AVPHYS_PAGES")
            avail_gb = (page_size * avail_pages) / (1024 ** 3)
        else:
            return  # Windows — skip check
    except (ValueError, OSError):
        return
    if avail_gb < min_gb:
        log.warning(
            "Only %.1f GB of memory available (minimum recommended: %.1f GB). "
            "Consider reducing the collection scope or increasing available memory.",
            avail_gb,
            min_gb,
        )


def _query(cursor: Any, sql_text: str) -> list[dict[str, Any]]:
    cursor.execute(sql_text)
    cols = [d[0] for d in cursor.description]
    rows = []
    while True:
        chunk = cursor.fetchmany(1000)
        if not chunk:
            break
        rows.extend(dict(zip(cols, row)) for row in chunk)
    return rows


def _parse_full_name(full_name: str) -> tuple[str, str, str]:
    """Split 'catalog.schema.table' into (catalog, schema, table)."""
    parts = (full_name or "").split(".")
    if len(parts) == 3:
        return parts[0], parts[1], parts[2]
    if len(parts) == 2:
        return "", parts[0], parts[1]
    return "", "", full_name


def collect_table_lineage(cursor: Any, lookback_days: int) -> list[dict[str, Any]]:
    rows = _query(
        cursor,
        f"""
        SELECT DISTINCT
            source_table_full_name,
            target_table_full_name,
            created_by,
            MAX(event_time) AS last_seen
        FROM system.access.table_lineage
        WHERE event_time >= DATEADD(DAY, -{lookback_days}, CURRENT_TIMESTAMP())
          AND source_table_full_name IS NOT NULL
          AND target_table_full_name IS NOT NULL
        GROUP BY source_table_full_name, target_table_full_name, created_by
        LIMIT 50000
        """,  # ← SUBSTITUTE: adjust lookback_days, LIMIT, or add catalog/schema filters
    )

    events: list[dict[str, Any]] = []
    for row in rows:
        src_catalog, src_schema, src_table = _parse_full_name(row["source_table_full_name"])
        dst_catalog, dst_schema, dst_table = _parse_full_name(row["target_table_full_name"])

        if not src_table or not dst_table:
            continue

        events.append({
            "sources": [{"database": src_catalog, "schema": src_schema, "asset_name": src_table}],
            "destination": {"database": dst_catalog, "schema": dst_schema, "asset_name": dst_table},
            "lineage_type": "table",
        })
    return events


def collect_column_lineage(cursor: Any, lookback_days: int) -> list[dict[str, Any]]:
    rows = _query(
        cursor,
        f"""
        SELECT DISTINCT
            source_table_full_name,
            source_column_name,
            target_table_full_name,
            target_column_name
        FROM system.access.column_lineage
        WHERE event_time >= DATEADD(DAY, -{lookback_days}, CURRENT_TIMESTAMP())
          AND source_table_full_name IS NOT NULL
          AND target_table_full_name IS NOT NULL
        LIMIT 50000
        """,  # ← SUBSTITUTE: adjust LIMIT or add catalog/schema filters if needed
    )

    # Group by destination table so we can build one event per destination
    grouped: dict[str, dict[str, Any]] = {}
    for row in rows:
        dst_key = row["target_table_full_name"]
        if dst_key not in grouped:
            grouped[dst_key] = {"dst_full": dst_key, "columns": []}
        grouped[dst_key]["columns"].append(row)

    events: list[dict[str, Any]] = []
    for dst_key, group in grouped.items():
        dst_catalog, dst_schema, dst_table = _parse_full_name(group["dst_full"])
        if not dst_table:
            continue

        col_fields: list[dict[str, Any]] = []
        for row in group["columns"]:
            src_catalog, src_schema, src_table = _parse_full_name(row["source_table_full_name"])
            col_fields.append({
                "destination_field": row["target_column_name"],
                "sources": [{
                    "database": src_catalog,
                    "schema": src_schema,
                    "asset_name": src_table,
                    "field": row["source_column_name"],
                }],
            })

        events.append({
            "sources": [],  # column lineage carries source refs inside col_fields
            "destination": {"database": dst_catalog, "schema": dst_schema, "asset_name": dst_table},
            "column_lineage": col_fields,
            "lineage_type": "column",
        })
    return events


def collect(
    host: str,
    http_path: str,
    token: str,
    manifest_path: str = "manifest_lineage.json",
    include_column_lineage: bool = False,
    lookback_days: int = LOOKBACK_DAYS,
) -> list[dict[str, Any]]:
    """Connect to Databricks, collect lineage, write a JSON manifest, and return events."""
    _check_available_memory(min_gb=2.0)
    collected_at = datetime.now(timezone.utc).isoformat()

    with sql.connect(
        server_hostname=host,    # ← SUBSTITUTE
        http_path=http_path,     # ← SUBSTITUTE
        access_token=token,      # ← SUBSTITUTE
    ) as conn:
        with conn.cursor() as cursor:
            table_events = collect_table_lineage(cursor, lookback_days)
            col_events = collect_column_lineage(cursor, lookback_days) if include_column_lineage else []

    all_events = table_events + col_events
    log.info(
        "Collected %d lineage events (%d table, %d column)",
        len(all_events), len(table_events), len(col_events),
    )

    manifest = {
        "resource_type": RESOURCE_TYPE,
        "collected_at": collected_at,
        "lookback_days": lookback_days,
        "table_lineage_events": len(table_events),
        "column_lineage_events": len(col_events),
        "events": all_events,
    }
    with open(manifest_path, "w") as fh:
        json.dump(manifest, fh, indent=2)
    log.info("Manifest written to %s (%d events)", manifest_path, len(all_events))

    return all_events


def main() -> None:
    parser = argparse.ArgumentParser(description="Collect Databricks lineage to a manifest file")
    parser.add_argument("--host", default=os.getenv("DATABRICKS_HOST"))           # ← SUBSTITUTE
    parser.add_argument("--http-path", default=os.getenv("DATABRICKS_HTTP_PATH")) # ← SUBSTITUTE
    parser.add_argument("--token", default=os.getenv("DATABRICKS_TOKEN"))         # ← SUBSTITUTE
    parser.add_argument("--lookback-days", type=int, default=LOOKBACK_DAYS)
    parser.add_argument(
        "--column-lineage", action="store_true",
        help="Also collect column-level lineage (requires system.access.column_lineage access)",
    )
    parser.add_argument("--manifest", default="manifest_lineage.json")
    args = parser.parse_args()

    required = ["host", "http_path", "token"]
    missing = [k for k in required if getattr(args, k) is None]
    if missing:
        parser.error(f"Missing required arguments/env vars: {missing}")

    collect(
        host=args.host,
        http_path=args.http_path,
        token=args.token,
        manifest_path=args.manifest,
        include_column_lineage=args.column_lineage,
        lookback_days=args.lookback_days,
    )


if __name__ == "__main__":
    main()

"""
Databricks — Metadata Collection (collect-only)
=================================================
Collects table schemas, row counts, and byte sizes from Databricks Unity Catalog
using INFORMATION_SCHEMA and DESCRIBE DETAIL, then writes a JSON manifest file
that can be consumed by push_metadata.py.

Substitution points (search for "← SUBSTITUTE"):
  - DATABRICKS_HOST       : workspace hostname (e.g. adb-1234.azuredatabricks.net)
  - DATABRICKS_HTTP_PATH  : SQL warehouse HTTP path (e.g. /sql/1.0/warehouses/abc123)
  - DATABRICKS_TOKEN      : personal access token or service-principal secret
  - DATABRICKS_CATALOG    : catalog to collect from (default: "hive_metastore" or "main")
  - SCHEMA_EXCLUSIONS     : schemas to skip

Prerequisites:
  pip install databricks-sql-connector
"""

from __future__ import annotations

import argparse
import json
import logging
import os
from datetime import datetime, timezone
from typing import Any

from databricks import sql

logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s")
log = logging.getLogger(__name__)

RESOURCE_TYPE = "databricks"

# Schemas to skip across all catalogs
SCHEMA_EXCLUSIONS: set[str] = {  # ← SUBSTITUTE: add any internal schemas to skip
    "information_schema",
    "__databricks_internal",
}


def _check_available_memory(min_gb: float = 2.0) -> None:
    """Warn if available memory is below the threshold."""
    try:
        if hasattr(os, "sysconf"):  # Linux / macOS
            page_size = os.sysconf("SC_PAGE_SIZE")
            avail_pages = os.sysconf("SC_AVPHYS_PAGES")
            avail_gb = (page_size * avail_pages) / (1024 ** 3)
        else:
            return  # Windows — skip check
    except (ValueError, OSError):
        return
    if avail_gb < min_gb:
        log.warning(
            "Only %.1f GB of memory available (minimum recommended: %.1f GB). "
            "Consider reducing the collection scope or increasing available memory.",
            avail_gb,
            min_gb,
        )


def _query(cursor: Any, sql_text: str, params: tuple | None = None) -> list[dict[str, Any]]:
    cursor.execute(sql_text, params)
    cols = [d[0] for d in cursor.description]
    rows = []
    while True:
        chunk = cursor.fetchmany(1000)
        if not chunk:
            break
        rows.extend(dict(zip(cols, row)) for row in chunk)
    return rows


def collect_tables(cursor: Any, catalog: str) -> list[dict[str, Any]]:
    return _query(
        cursor,
        f"""
        SELECT table_catalog, table_schema, table_name, table_type, comment
        FROM {catalog}.information_schema.tables
        WHERE table_schema NOT IN ({", ".join(f"'{s}'" for s in SCHEMA_EXCLUSIONS)})
        ORDER BY table_schema, table_name
        """,  # ← SUBSTITUTE: add additional WHERE filters if needed
    )


def collect_columns(cursor: Any, catalog: str, schema: str, table: str) -> list[dict[str, Any]]:
    return _query(
        cursor,
        f"""
        SELECT column_name, data_type, comment
        FROM {catalog}.information_schema.columns
        WHERE table_schema = '{schema}' AND table_name = '{table}'
        ORDER BY ordinal_position
        """,
    )


def collect_detail(cursor: Any, catalog: str, schema: str, table: str) -> dict[str, Any] | None:
    try:
        rows = _query(cursor, f"DESCRIBE DETAIL `{catalog}`.`{schema}`.`{table}`")
        return rows[0] if rows else None
    except Exception:
        log.debug("DESCRIBE DETAIL failed for %s.%s.%s", catalog, schema, table, exc_info=True)
        return None


def collect(
    host: str,
    http_path: str,
    token: str,
    catalog: str,
    manifest_path: str = "manifest_metadata.json",
) -> list[dict[str, Any]]:
    """Connect to Databricks, collect metadata, write a JSON manifest, and return the asset dicts.

    The manifest contains serialised asset dicts that push_metadata.py can read.
    """
    _check_available_memory(min_gb=2.0)
    collected_at = datetime.now(timezone.utc).isoformat()
    assets: list[dict[str, Any]] = []

    with sql.connect(
        server_hostname=host,    # ← SUBSTITUTE
        http_path=http_path,     # ← SUBSTITUTE
        access_token=token,      # ← SUBSTITUTE
    ) as conn:
        with conn.cursor() as cursor:
            tables = collect_tables(cursor, catalog)
            log.info("Found %d tables in catalog %s", len(tables), catalog)

            for row in tables:
                schema = row["table_schema"]
                table_name = row["table_name"]

                columns = collect_columns(cursor, catalog, schema, table_name)
                fields = [
                    {
                        "name": col["column_name"],
                        "type": col["data_type"].upper(),
                        "description": col.get("comment") or None,
                    }
                    for col in columns
                ]

                detail = collect_detail(cursor, catalog, schema, table_name)
                row_count: int | None = None
                byte_count: int | None = None
                last_updated: str | None = None
                if detail:
                    row_count = detail.get("numRows")
                    byte_count = detail.get("sizeInBytes")
                    last_modified = detail.get("lastModified")
                    if last_modified:
                        last_updated = (
                            last_modified.isoformat()
                            if hasattr(last_modified, "isoformat")
                            else str(last_modified)
                        )

                asset = {
                    "asset_name": table_name,
                    "database": catalog,    # ← SUBSTITUTE: use catalog as database
                    "schema": schema,
                    "asset_type": "VIEW" if row.get("table_type", "").upper() == "VIEW" else "TABLE",
                    "description": row.get("comment") or None,
                    "fields": fields,
                    "row_count": row_count,
                    "byte_count": byte_count,
                    "last_updated": last_updated,
                }
                assets.append(asset)
                log.info("Collected %s.%s.%s", catalog, schema, table_name)

    manifest = {
        "resource_type": RESOURCE_TYPE,
        "collected_at": collected_at,
        "catalog": catalog,
        "asset_count": len(assets),
        "assets": assets,
    }
    with open(manifest_path, "w") as fh:
        json.dump(manifest, fh, indent=2)
    log.info("Manifest written to %s (%d assets)", manifest_path, len(assets))

    return assets


def main() -> None:
    parser = argparse.ArgumentParser(description="Collect Databricks metadata to a manifest file")
    parser.add_argument("--host", default=os.getenv("DATABRICKS_HOST"))           # ← SUBSTITUTE
    parser.add_argument("--http-path", default=os.getenv("DATABRICKS_HTTP_PATH")) # ← SUBSTITUTE
    parser.add_argument("--token", default=os.getenv("DATABRICKS_TOKEN"))         # ← SUBSTITUTE
    parser.add_argument("--catalog", default=os.getenv("DATABRICKS_CATALOG", "hive_metastore"))
    parser.add_argument("--manifest", default="manifest_metadata.json")
    args = parser.parse_args()

    required = ["host", "http_path", "token"]
    missing = [k for k in required if getattr(args, k) is None]
    if missing:
        parser.error(f"Missing required arguments/env vars: {missing}")

    collect(
        host=args.host,
        http_path=args.http_path,
        token=args.token,
        catalog=args.catalog,
        manifest_path=args.manifest,
    )


if __name__ == "__main__":
    main()

"""
Databricks — Query Log Collection (collect-only)
==================================================
Collects finished query execution records from the Databricks system table
system.query.history and writes a JSON manifest file that can be consumed
by push_query_logs.py.

Substitution points (search for "← SUBSTITUTE"):
  - DATABRICKS_HOST       : workspace hostname
  - DATABRICKS_HTTP_PATH  : SQL warehouse HTTP path
  - DATABRICKS_TOKEN      : PAT or service-principal secret
  - LOOKBACK_HOURS        : hours back from [now - LAG_HOURS] to collect (default 25)
  - LOOKBACK_LAG_HOURS    : hours to lag behind now to avoid in-flight queries (default 1)
  - MAX_ROWS              : maximum query rows to collect per run (default 10000)

Prerequisites:
  pip install databricks-sql-connector
"""

from __future__ import annotations

import argparse
import json
import logging
import os
from datetime import datetime, timezone
from typing import Any

from databricks import sql

logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s")
log = logging.getLogger(__name__)

LOG_TYPE = "databricks"

LOOKBACK_HOURS: int = int(os.getenv("LOOKBACK_HOURS", "25"))        # ← SUBSTITUTE
LOOKBACK_LAG_HOURS: int = int(os.getenv("LOOKBACK_LAG_HOURS", "1")) # ← SUBSTITUTE
MAX_ROWS: int = int(os.getenv("MAX_ROWS", "10000"))                  # ← SUBSTITUTE

_QUERY_LOG_SQL = """\
SELECT
    statement_id       AS query_id,
    statement_text     AS query_text,
    start_time,
    end_time,
    executed_by        AS user_name,
    produced_rows      AS returned_rows,
    total_task_duration_ms,
    read_rows,
    read_bytes
FROM system.query.history
WHERE start_time >= DATEADD(HOUR, -{lookback_hours}, NOW())
  AND start_time <  DATEADD(HOUR, -{lag_hours}, NOW())
  AND status = 'FINISHED'
ORDER BY start_time
LIMIT {max_rows}
"""  # ← SUBSTITUTE: adjust status filter or add warehouse_id filter as needed


def _check_available_memory(min_gb: float = 2.0) -> None:
    """Warn if available memory is below the threshold."""
    try:
        if hasattr(os, "sysconf"):  # Linux / macOS
            page_size = os.sysconf("SC_PAGE_SIZE")
            avail_pages = os.sysconf("SC_AVPHYS_PAGES")
            avail_gb = (page_size * avail_pages) / (1024 ** 3)
        else:
            return  # Windows — skip check
    except (ValueError, OSError):
        return
    if avail_gb < min_gb:
        log.warning(
            "Only %.1f GB of memory available (minimum recommended: %.1f GB). "
            "Consider reducing the collection scope or increasing available memory.",
            avail_gb,
            min_gb,
        )


def _safe_isoformat(dt: Any) -> str | None:
    if dt is None:
        return None
    if hasattr(dt, "isoformat"):
        if dt.tzinfo is None:
            dt = dt.replace(tzinfo=timezone.utc)
        return dt.isoformat()
    return str(dt)


def _query(cursor: Any, sql_text: str) -> list[dict[str, Any]]:
    cursor.execute(sql_text)
    cols = [d[0] for d in cursor.description]
    rows = []
    while True:
        chunk = cursor.fetchmany(1000)
        if not chunk:
            break
        rows.extend(dict(zip(cols, row)) for row in chunk)
    return rows


def collect_query_logs(
    cursor: Any,
    lookback_hours: int,
    lag_hours: int,
    max_rows: int,
) -> list[dict[str, Any]]:
    rendered_sql = _QUERY_LOG_SQL.format(
        lookback_hours=lookback_hours + lag_hours,  # offset from NOW() to cover the window
        lag_hours=lag_hours,
        max_rows=max_rows,
    )
    rows = _query(cursor, rendered_sql)
    log.info("Retrieved %d query log rows from system.query.history", len(rows))

    entries: list[dict[str, Any]] = []
    for row in rows:
        query_text: str = row.get("query_text") or ""
        if not query_text.strip():
            continue  # ← SUBSTITUTE: decide whether to skip empty-text rows

        entry = {
            "query_id": row.get("query_id"),
            "query_text": query_text,
            "start_time": _safe_isoformat(row.get("start_time")),
            "end_time": _safe_isoformat(row.get("end_time")),
            "user": row.get("user_name"),
            "returned_rows": row.get("returned_rows"),
            "total_task_duration_ms": row.get("total_task_duration_ms"),
            "read_rows": row.get("read_rows"),
            "read_bytes": row.get("read_bytes"),
        }
        entries.append(entry)

    return entries


def collect(
    host: str,
    http_path: str,
    token: str,
    manifest_path: str = "manifest_query_logs.json",
    lookback_hours: int = LOOKBACK_HOURS,
    lookback_lag_hours: int = LOOKBACK_LAG_HOURS,
    max_rows: int = MAX_ROWS,
) -> list[dict[str, Any]]:
    """Connect to Databricks, collect query logs, write a JSON manifest, and return entries."""
    _check_available_memory(min_gb=2.0)
    collected_at = datetime.now(timezone.utc).isoformat()

    with sql.connect(
        server_hostname=host,    # ← SUBSTITUTE
        http_path=http_path,     # ← SUBSTITUTE
        access_token=token,      # ← SUBSTITUTE
    ) as conn:
        with conn.cursor() as cursor:
            entries = collect_query_logs(cursor, lookback_hours, lookback_lag_hours, max_rows)

    log.info("Collected %d query log entries", len(entries))

    manifest = {
        "log_type": LOG_TYPE,
        "collected_at": collected_at,
        "lookback_hours": lookback_hours,
        "lookback_lag_hours": lookback_lag_hours,
        "query_log_count": len(entries),
        "entries": entries,
    }
    with open(manifest_path, "w") as fh:
        json.dump(manifest, fh, indent=2)
    log.info("Manifest written to %s (%d entries)", manifest_path, len(entries))

    return entries


def main() -> None:
    parser = argparse.ArgumentParser(description="Collect Databricks query logs to a manifest file")
    parser.add_argument("--host", default=os.getenv("DATABRICKS_HOST"))           # ← SUBSTITUTE
    parser.add_argument("--http-path", default=os.getenv("DATABRICKS_HTTP_PATH")) # ← SUBSTITUTE
    parser.add_argument("--token", default=os.getenv("DATABRICKS_TOKEN"))         # ← SUBSTITUTE
    parser.add_argument("--lookback-hours", type=int, default=LOOKBACK_HOURS)
    parser.add_argument("--lookback-lag-hours", type=int, default=LOOKBACK_LAG_HOURS)
    parser.add_argument("--max-rows", type=int, default=MAX_ROWS)
    parser.add_argument("--manifest", default="manifest_query_logs.json")
    args = parser.parse_args()

    required = ["host", "http_path", "token"]
    missing = [k for k in required if getattr(args, k) is None]
    if missing:
        parser.error(f"Missing required arguments/env vars: {missing}")

    collect(
        host=args.host,
        http_path=args.http_path,
        token=args.token,
        manifest_path=args.manifest,
        lookback_hours=args.lookback_hours,
        lookback_lag_hours=args.lookback_lag_hours,
        max_rows=args.max_rows,
    )


if __name__ == "__main__":
    main()

"""
Databricks — Lineage Push (push-only)
=======================================
Reads a JSON manifest file produced by collect_lineage.py and pushes the lineage
events to Monte Carlo via the push ingestion API, with configurable batching to
keep compressed payloads under 1 MB.

Substitution points (search for "← SUBSTITUTE"):
  - MCD_INGEST_ID / MCD_INGEST_TOKEN : Monte Carlo API credentials
  - MCD_RESOURCE_UUID      : UUID of the Databricks connection in Monte Carlo
  - PUSH_BATCH_SIZE       : number of events per API call (default 500)

Prerequisites:
  pip install pycarlo
"""

from __future__ import annotations

import argparse
import json
import logging
import os
from concurrent.futures import ThreadPoolExecutor, as_completed
from datetime import datetime, timezone
from typing import Any

from pycarlo.core import Client, Session
from pycarlo.features.ingestion import IngestionService
from pycarlo.features.ingestion.models import (
    ColumnLineageField,
    ColumnLineageSourceField,
    LineageAssetRef,
    LineageEvent,
)

logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s")
log = logging.getLogger(__name__)

RESOURCE_TYPE = "databricks"
DEFAULT_BATCH_SIZE = 500  # ← SUBSTITUTE: conservative default to stay under 1 MB compressed


def _ref_from_dict(d: dict[str, Any]) -> LineageAssetRef:
    database = d.get("database", "")
    schema = d.get("schema", "")
    name = d["asset_name"]
    return LineageAssetRef(
        type="TABLE",
        name=name,
        database=database,
        schema=schema,
        asset_id=f"{database}__{schema}__{name}",
    )


def _event_from_dict(d: dict[str, Any]) -> LineageEvent:
    """Reconstruct a LineageEvent from a manifest dict."""
    sources = [_ref_from_dict(s) for s in d.get("sources", [])]
    destination = _ref_from_dict(d["destination"])

    fields: list[ColumnLineageField] | None = None
    if d.get("column_lineage"):
        fields = []
        for cl in d["column_lineage"]:
            src_fields = []
            for s in cl.get("sources", []):
                asset_id = f"{s.get('database', '')}__{s.get('schema', '')}__{s['asset_name']}"
                src_fields.append(
                    ColumnLineageSourceField(
                        asset_id=asset_id,
                        field_name=s["field"],
                    )
                )
            fields.append(
                ColumnLineageField(
                    name=cl["destination_field"],
                    source_fields=src_fields,
                )
            )

    return LineageEvent(
        sources=sources,
        destination=destination,
        fields=fields,
    )


def push(
    manifest_path: str,
    resource_uuid: str,
    key_id: str,
    key_token: str,
    batch_size: int = DEFAULT_BATCH_SIZE,
) -> dict[str, Any]:
    """Read a collect manifest and push lineage events to Monte Carlo in batches.

    Returns a summary dict with invocation IDs and counts.
    """
    with open(manifest_path) as fh:
        manifest = json.load(fh)

    event_dicts: list[dict[str, Any]] = manifest["events"]
    events = [_event_from_dict(d) for d in event_dicts]
    log.info("Loaded %d lineage events from %s", len(events), manifest_path)

    # Split into batches
    batches = []
    for i in range(0, len(events), batch_size):
        batches.append(events[i : i + batch_size])
    total_batches = len(batches)

    def _push_batch(batch: list, batch_num: int) -> str | None:
        """Push a single batch using a dedicated Session (thread-safe)."""
        log.info("Pushing batch %d/%d (%d events) ...", batch_num, total_batches, len(batch))
        client = Client(session=Session(mcd_id=key_id, mcd_token=key_token, scope="Ingestion"))
        service = IngestionService(mc_client=client)
        result = service.send_lineage(
            resource_uuid=resource_uuid,
            resource_type=RESOURCE_TYPE,
            events=batch,
        )
        invocation_id = service.extract_invocation_id(result)
        if invocation_id:
            log.info("Batch %d: invocation_id=%s", batch_num, invocation_id)
        return invocation_id

    # Push batches in parallel (each thread gets its own pycarlo Session)
    max_workers = min(4, total_batches)
    invocation_ids: list[str | None] = [None] * total_batches

    with ThreadPoolExecutor(max_workers=max_workers) as pool:
        futures = {
            pool.submit(_push_batch, batch, i + 1): i
            for i, batch in enumerate(batches)
        }
        for future in as_completed(futures):
            idx = futures[future]
            try:
                invocation_ids[idx] = future.result()
            except Exception as exc:
                log.error("ERROR pushing batch %d: %s", idx + 1, exc)
                raise

    log.info("All %d batches pushed (%d workers)", total_batches, max_workers)

    pushed_at = datetime.now(timezone.utc).isoformat()
    summary = {
        "resource_uuid": resource_uuid,
        "resource_type": RESOURCE_TYPE,
        "invocation_ids": invocation_ids,
        "pushed_at": pushed_at,
        "event_count": len(events),
        "batch_count": total_batches,
        "batch_size": batch_size,
        "lookback_days": manifest.get("lookback_days"),
        "table_lineage_events": manifest.get("table_lineage_events"),
        "column_lineage_events": manifest.get("column_lineage_events"),
    }

    push_manifest_path = manifest_path.replace(".json", "_push_result.json")
    with open(push_manifest_path, "w") as fh:
        json.dump(summary, fh, indent=2)
    log.info("Push result written to %s", push_manifest_path)

    return summary


def main() -> None:
    parser = argparse.ArgumentParser(description="Push Databricks lineage to Monte Carlo from manifest")
    parser.add_argument("--manifest", default="manifest_lineage.json")
    parser.add_argument("--resource-uuid", default=os.getenv("MCD_RESOURCE_UUID"))
    parser.add_argument("--key-id", default=os.getenv("MCD_INGEST_ID"))
    parser.add_argument("--key-token", default=os.getenv("MCD_INGEST_TOKEN"))
    parser.add_argument("--batch-size", type=int, default=DEFAULT_BATCH_SIZE)
    args = parser.parse_args()

    required = ["resource_uuid", "key_id", "key_token"]
    missing = [k for k in required if getattr(args, k) is None]
    if missing:
        parser.error(f"Missing required arguments/env vars: {missing}")

    push(
        manifest_path=args.manifest,
        resource_uuid=args.resource_uuid,
        key_id=args.key_id,
        key_token=args.key_token,
        batch_size=args.batch_size,
    )


if __name__ == "__main__":
    main()

"""
Databricks — Metadata Push (push-only)
========================================
Reads a JSON manifest file produced by collect_metadata.py and pushes the assets
to Monte Carlo via the push ingestion API, with configurable batching to keep
compressed payloads under 1 MB.

Substitution points (search for "← SUBSTITUTE"):
  - MCD_INGEST_ID / MCD_INGEST_TOKEN : Monte Carlo API credentials
  - MCD_RESOURCE_UUID      : UUID of the Databricks connection in Monte Carlo
  - PUSH_BATCH_SIZE       : number of assets per API call (default 500)

Prerequisites:
  pip install pycarlo
"""

from __future__ import annotations

import argparse
import json
import logging
import os
from concurrent.futures import ThreadPoolExecutor, as_completed
from datetime import datetime, timezone
from typing import Any

from pycarlo.core import Client, Session
from pycarlo.features.ingestion import IngestionService
from pycarlo.features.ingestion.models import (
    AssetField,
    AssetFreshness,
    AssetMetadata,
    AssetVolume,
    RelationalAsset,
)

logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s")
log = logging.getLogger(__name__)

RESOURCE_TYPE = "databricks"
DEFAULT_BATCH_SIZE = 500  # ← SUBSTITUTE: conservative default to stay under 1 MB compressed


def _asset_from_dict(d: dict[str, Any]) -> RelationalAsset:
    """Reconstruct a RelationalAsset from a manifest dict."""
    fields = [
        AssetField(
            name=f["name"],
            type=f.get("type"),
            description=f.get("description"),
        )
        for f in d.get("fields", [])
    ]

    volume = None
    if d.get("row_count") is not None or d.get("byte_count") is not None:
        volume = AssetVolume(row_count=d.get("row_count"), byte_count=d.get("byte_count"))

    freshness = None
    if d.get("last_updated") is not None:
        freshness = AssetFreshness(last_update_time=d.get("last_updated"))

    return RelationalAsset(
        type=d.get("asset_type", "TABLE"),
        metadata=AssetMetadata(
            name=d["asset_name"],
            database=d["database"],    # ← SUBSTITUTE: use catalog as database
            schema=d["schema"],
            description=d.get("description"),
        ),
        fields=fields,
        volume=volume,
        freshness=freshness,
    )


def push(
    manifest_path: str,
    resource_uuid: str,
    key_id: str,
    key_token: str,
    batch_size: int = DEFAULT_BATCH_SIZE,
) -> dict[str, Any]:
    """Read a collect manifest and push assets to Monte Carlo in batches.

    Returns a summary dict with invocation IDs and counts.
    """
    with open(manifest_path) as fh:
        manifest = json.load(fh)

    asset_dicts: list[dict[str, Any]] = manifest["assets"]
    assets = [_asset_from_dict(d) for d in asset_dicts]
    log.info("Loaded %d assets from %s", len(assets), manifest_path)

    # Split into batches
    batches = []
    for i in range(0, max(len(assets), 1), batch_size):
        batches.append(assets[i : i + batch_size])
    total_batches = len(batches)

    def _push_batch(batch: list, batch_num: int) -> str | None:
        """Push a single batch using a dedicated Session (thread-safe)."""
        client = Client(session=Session(mcd_id=key_id, mcd_token=key_token, scope="Ingestion"))
        service = IngestionService(mc_client=client)
        result = service.send_metadata(
            resource_uuid=resource_uuid,
            resource_type=RESOURCE_TYPE,
            events=batch,
        )
        invocation_id = service.extract_invocation_id(result)
        log.info("Pushed batch %d/%d (%d assets) — invocation_id=%s", batch_num, total_batches, len(batch), invocation_id)
        return invocation_id

    # Push batches in parallel (each thread gets its own pycarlo Session)
    max_workers = min(4, total_batches)
    invocation_ids: list[str | None] = [None] * total_batches

    with ThreadPoolExecutor(max_workers=max_workers) as pool:
        futures = {
            pool.submit(_push_batch, batch, i + 1): i
            for i, batch in enumerate(batches)
        }
        for future in as_completed(futures):
            idx = futures[future]
            try:
                invocation_ids[idx] = future.result()
            except Exception as exc:
                log.error("ERROR pushing batch %d: %s", idx + 1, exc)
                raise

    log.info("All %d batches pushed (%d workers)", total_batches, max_workers)

    pushed_at = datetime.now(timezone.utc).isoformat()
    summary = {
        "resource_uuid": resource_uuid,
        "resource_type": RESOURCE_TYPE,
        "invocation_ids": invocation_ids,
        "pushed_at": pushed_at,
        "asset_count": len(assets),
        "batch_count": total_batches,
        "batch_size": batch_size,
        "catalog": manifest.get("catalog"),
    }

    # Write push result alongside the collect manifest
    push_manifest_path = manifest_path.replace(".json", "_push_result.json")
    with open(push_manifest_path, "w") as fh:
        json.dump(summary, fh, indent=2)
    log.info("Push result written to %s", push_manifest_path)

    return summary


def main() -> None:
    parser = argparse.ArgumentParser(description="Push Databricks metadata to Monte Carlo from manifest")
    parser.add_argument("--manifest", default="manifest_metadata.json")
    parser.add_argument("--resource-uuid", default=os.getenv("MCD_RESOURCE_UUID"))
    parser.add_argument("--key-id", default=os.getenv("MCD_INGEST_ID"))
    parser.add_argument("--key-token", default=os.getenv("MCD_INGEST_TOKEN"))
    parser.add_argument("--batch-size", type=int, default=DEFAULT_BATCH_SIZE)
    args = parser.parse_args()

    required = ["resource_uuid", "key_id", "key_token"]
    missing = [k for k in required if getattr(args, k) is None]
    if missing:
        parser.error(f"Missing required arguments/env vars: {missing}")

    push(
        manifest_path=args.manifest,
        resource_uuid=args.resource_uuid,
        key_id=args.key_id,
        key_token=args.key_token,
        batch_size=args.batch_size,
    )


if __name__ == "__main__":
    main()

"""
Databricks — Query Log Push (push-only)
=========================================
Reads a JSON manifest file produced by collect_query_logs.py and pushes the query
log entries to Monte Carlo via the push ingestion API, with configurable batching
to keep compressed payloads under 1 MB.

Substitution points (search for "← SUBSTITUTE"):
  - MCD_INGEST_ID / MCD_INGEST_TOKEN : Monte Carlo API credentials
  - MCD_RESOURCE_UUID      : UUID of the Databricks connection in Monte Carlo
  - PUSH_BATCH_SIZE       : number of entries per API call (default 100)

Prerequisites:
  pip install pycarlo
"""

from __future__ import annotations

import argparse
import json
import logging
import os
from concurrent.futures import ThreadPoolExecutor, as_completed
from datetime import datetime, timezone
from typing import Any

from dateutil.parser import isoparse
from pycarlo.core import Client, Session
from pycarlo.features.ingestion import IngestionService
from pycarlo.features.ingestion.models import QueryLogEntry

logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s")
log = logging.getLogger(__name__)

LOG_TYPE = "databricks"
DEFAULT_BATCH_SIZE = 100  # ← SUBSTITUTE: conservative default to stay under 1 MB compressed

# Truncate query_text longer than this to prevent 413 errors.
# Some SQL statements (e.g., generated by BI tools) can be 100KB+ and blow up
# compressed payloads even at small batch sizes.
_MAX_QUERY_TEXT_LEN = 10_000


def _build_query_log_entries(entry_dicts: list[dict[str, Any]]) -> list[QueryLogEntry]:
    """Convert manifest query dicts into QueryLogEntry objects."""
    entries = []
    truncated = 0
    for d in entry_dicts:
        query_text = d.get("query_text") or ""

        # Truncate very long SQL to prevent 413 Request Too Large
        if len(query_text) > _MAX_QUERY_TEXT_LEN:
            query_text = query_text[:_MAX_QUERY_TEXT_LEN] + "... [TRUNCATED]"
            truncated += 1

        extra = {}
        if d.get("total_task_duration_ms") is not None:
            extra["total_task_duration_ms"] = d["total_task_duration_ms"]
        if d.get("read_rows") is not None:
            extra["read_rows"] = d["read_rows"]
        if d.get("read_bytes") is not None:
            extra["read_bytes"] = d["read_bytes"]

        start_time = d.get("start_time")
        end_time = d.get("end_time")

        entries.append(
            QueryLogEntry(
                query_id=d.get("query_id"),
                query_text=query_text,
                start_time=isoparse(start_time) if start_time else None,
                end_time=isoparse(end_time) if end_time else None,
                user=d.get("user"),
                returned_rows=d.get("returned_rows"),
                extra=extra or None,
            )
        )
    if truncated:
        log.info("Truncated %d query text(s) exceeding %d chars", truncated, _MAX_QUERY_TEXT_LEN)
    return entries


def push(
    manifest_path: str,
    resource_uuid: str,
    key_id: str,
    key_token: str,
    batch_size: int = DEFAULT_BATCH_SIZE,
) -> dict[str, Any]:
    """Read a collect manifest and push query log entries to Monte Carlo in batches.

    Returns a summary dict with invocation IDs and counts.
    """
    with open(manifest_path) as fh:
        manifest = json.load(fh)

    entry_dicts: list[dict[str, Any]] = manifest["entries"]
    entries = _build_query_log_entries(entry_dicts)
    log.info("Loaded %d query log entries from %s", len(entries), manifest_path)

    if not entries:
        log.info("No query log entries to push.")
        summary = {
            "resource_uuid": resource_uuid,
            "log_type": LOG_TYPE,
            "invocation_ids": [],
            "pushed_at": datetime.now(timezone.utc).isoformat(),
            "query_log_count": 0,
            "batch_count": 0,
            "batch_size": batch_size,
        }
        push_manifest_path = manifest_path.replace(".json", "_push_result.json")
        with open(push_manifest_path, "w") as fh:
            json.dump(summary, fh, indent=2)
        return summary

    # Split into batches
    batches = []
    for i in range(0, len(entries), batch_size):
        batches.append(entries[i : i + batch_size])
    total_batches = len(batches)

    def _push_batch(batch: list, batch_num: int) -> str | None:
        """Push a single batch using a dedicated Session (thread-safe)."""
        client = Client(session=Session(mcd_id=key_id, mcd_token=key_token, scope="Ingestion"))
        service = IngestionService(mc_client=client)
        result = service.send_query_logs(
            resource_uuid=resource_uuid,
            log_type=LOG_TYPE,
            events=batch,
        )
        invocation_id = service.extract_invocation_id(result)
        log.info("Pushed batch %d/%d (%d entries) — invocation_id=%s", batch_num, total_batches, len(batch), invocation_id)
        return invocation_id

    # Push batches in parallel (each thread gets its own pycarlo Session)
    max_workers = min(4, total_batches)
    invocation_ids: list[str | None] = [None] * total_batches

    with ThreadPoolExecutor(max_workers=max_workers) as pool:
        futures = {
            pool.submit(_push_batch, batch, i + 1): i
            for i, batch in enumerate(batches)
        }
        for future in as_completed(futures):
            idx = futures[future]
            try:
                invocation_ids[idx] = future.result()
            except Exception as exc:
                log.error("ERROR pushing batch %d: %s", idx + 1, exc)
                raise

    log.info("All %d batches pushed (%d workers)", total_batches, max_workers)

    pushed_at = datetime.now(timezone.utc).isoformat()
    summary = {
        "resource_uuid": resource_uuid,
        "log_type": LOG_TYPE,
        "invocation_ids": invocation_ids,
        "pushed_at": pushed_at,
        "query_log_count": len(entries),
        "batch_count": total_batches,
        "batch_size": batch_size,
        "lookback_hours": manifest.get("lookback_hours"),
        "lookback_lag_hours": manifest.get("lookback_lag_hours"),
    }

    push_manifest_path = manifest_path.replace(".json", "_push_result.json")
    with open(push_manifest_path, "w") as fh:
        json.dump(summary, fh, indent=2)
    log.info("Push result written to %s", push_manifest_path)

    return summary


def main() -> None:
    parser = argparse.ArgumentParser(description="Push Databricks query logs to Monte Carlo from manifest")
    parser.add_argument("--manifest", default="manifest_query_logs.json")
    parser.add_argument("--resource-uuid", default=os.getenv("MCD_RESOURCE_UUID"))
    parser.add_argument("--key-id", default=os.getenv("MCD_INGEST_ID"))
    parser.add_argument("--key-token", default=os.getenv("MCD_INGEST_TOKEN"))
    parser.add_argument("--batch-size", type=int, default=DEFAULT_BATCH_SIZE)
    args = parser.parse_args()

    required = ["resource_uuid", "key_id", "key_token"]
    missing = [k for k in required if getattr(args, k) is None]
    if missing:
        parser.error(f"Missing required arguments/env vars: {missing}")

    push(
        manifest_path=args.manifest,
        resource_uuid=args.resource_uuid,
        key_id=args.key_id,
        key_token=args.key_token,
        batch_size=args.batch_size,
    )


if __name__ == "__main__":
    main()

#!/usr/bin/env python3
"""
Extract Hive lineage from a local log file and push it to Monte Carlo in one step.

Thin wrapper that calls ``collect()`` from ``collect_lineage`` followed by
``push()`` from ``push_lineage``, then writes the final manifest (with
``resource_uuid`` and ``invocation_id``) to ``--output-file``.

Substitution points
-------------------
- MCD_INGEST_ID    (env) / --key-id        (CLI) : Monte Carlo ingestion key ID
- MCD_INGEST_TOKEN (env) / --key-token      (CLI) : Monte Carlo ingestion key token
- MCD_RESOURCE_UUID    (env) / --resource-uuid  (CLI) : MC resource UUID for this connection
- --log-file                                         : path to local HiveServer2 log

Prerequisites
-------------
    pip install pycarlo python-dotenv

Usage (table-level):
    python collect_and_push_lineage.py \\
        --key-id  <MCD_INGEST_ID> \\
        --key-token <MCD_INGEST_TOKEN> \\
        --resource-uuid <MCD_RESOURCE_UUID> \\
        --log-file /tmp/root/hive.log

Usage (column-level):
    python collect_and_push_lineage.py ... --column-lineage
"""

import argparse
import json
import os

from collect_lineage import collect
from push_lineage import DEFAULT_BATCH_SIZE, DEFAULT_TIMEOUT_SECONDS, push


def main() -> None:
    parser = argparse.ArgumentParser(
        description="Extract Hive lineage from a local log file and push to Monte Carlo",
    )
    # Collect args
    parser.add_argument(
        "--log-file",
        default="/tmp/root/hive.log",
        help="Path to local HiveServer2 log file (default: /tmp/root/hive.log)",  # ← SUBSTITUTE: your log path
    )
    # Push / MC args
    parser.add_argument(
        "--key-id",
        default=os.environ.get("MCD_INGEST_ID"),
        help="Monte Carlo ingestion key ID (env: MCD_INGEST_ID)",
    )
    parser.add_argument(
        "--key-token",
        default=os.environ.get("MCD_INGEST_TOKEN"),
        help="Monte Carlo ingestion key token (env: MCD_INGEST_TOKEN)",
    )
    parser.add_argument(
        "--resource-uuid",
        default=os.environ.get("MCD_RESOURCE_UUID"),
        help="Monte Carlo resource UUID for this Hive connection (env: MCD_RESOURCE_UUID)",
    )
    parser.add_argument(
        "--column-lineage",
        action="store_true",
        help="Push column-level lineage instead of table-level",
    )
    parser.add_argument(
        "--output-file",
        default="lineage_output.json",
        help="Path to write the lineage manifest (default: lineage_output.json)",
    )
    parser.add_argument(
        "--batch-size",
        type=int,
        default=DEFAULT_BATCH_SIZE,
        metavar="N",
        help=f"Max events per POST (default: {DEFAULT_BATCH_SIZE})",
    )
    parser.add_argument(
        "--timeout",
        type=int,
        default=DEFAULT_TIMEOUT_SECONDS,
        metavar="SEC",
        help=f"HTTP timeout per request in seconds (default: {DEFAULT_TIMEOUT_SECONDS})",
    )
    args = parser.parse_args()

    if not args.key_id or not args.key_token:
        parser.error("--key-id and --key-token are required (or set MCD_INGEST_ID / MCD_INGEST_TOKEN)")
    if not args.resource_uuid:
        parser.error("--resource-uuid is required (or set MCD_RESOURCE_UUID)")

    manifest = collect(log_file=args.log_file)

    if not manifest["edges"]:
        print("No lineage edges detected — no CTAS or INSERT INTO ... SELECT patterns found.")
        return

    push(
        manifest=manifest,
        resource_uuid=args.resource_uuid,
        key_id=args.key_id,
        key_token=args.key_token,
        column_lineage=args.column_lineage,
        batch_size=args.batch_size,
        timeout_seconds=args.timeout,
    )

    with open(args.output_file, "w") as fh:
        json.dump(manifest, fh, indent=2)
    print(f"Lineage manifest written to {args.output_file}")
    print("Done.")


if __name__ == "__main__":
    main()

#!/usr/bin/env python3
"""
Collect Hive table metadata and push it to Monte Carlo in one step.

Thin wrapper that calls ``collect()`` from ``collect_metadata`` followed by
``push()`` from ``push_metadata``, then writes the final manifest (with
``resource_uuid`` and ``invocation_id``) to ``--output-file``.

Substitution points
-------------------
- HIVE_HOST           (env) / --hive-host      (CLI) : HiveServer2 hostname
- MCD_INGEST_ID    (env) / --key-id         (CLI) : Monte Carlo ingestion key ID
- MCD_INGEST_TOKEN (env) / --key-token      (CLI) : Monte Carlo ingestion key token
- MCD_RESOURCE_UUID    (env) / --resource-uuid  (CLI) : MC resource UUID for this connection

Prerequisites
-------------
    pip install pycarlo pyhive python-dotenv

Usage
-----
    python collect_and_push_metadata.py \\
        --key-id  <MCD_INGEST_ID> \\
        --key-token <MCD_INGEST_TOKEN> \\
        --resource-uuid <MCD_RESOURCE_UUID> \\
        --hive-host <HIVESERVER2_HOSTNAME>
"""

import argparse
import json
import os

from collect_metadata import collect
from push_metadata import DEFAULT_BATCH_SIZE, DEFAULT_TIMEOUT_SECONDS, push


def main() -> None:
    parser = argparse.ArgumentParser(
        description="Collect Hive table metadata and push to Monte Carlo",
    )
    # Hive / collect args
    parser.add_argument(
        "--hive-host",
        default=os.environ.get("HIVE_HOST"),
        help="HiveServer2 hostname (env: HIVE_HOST)",  # ← SUBSTITUTE: your EMR master DNS or Hive host
    )
    parser.add_argument(
        "--hive-port",
        type=int,
        default=10000,
        help="HiveServer2 port (default: 10000)",  # ← SUBSTITUTE if your cluster uses a non-standard port
    )
    # Push / MC args
    parser.add_argument(
        "--key-id",
        default=os.environ.get("MCD_INGEST_ID"),
        help="Monte Carlo ingestion key ID (env: MCD_INGEST_ID)",  # ← SUBSTITUTE env var name if different
    )
    parser.add_argument(
        "--key-token",
        default=os.environ.get("MCD_INGEST_TOKEN"),
        help="Monte Carlo ingestion key token (env: MCD_INGEST_TOKEN)",  # ← SUBSTITUTE env var name if different
    )
    parser.add_argument(
        "--resource-uuid",
        default=os.environ.get("MCD_RESOURCE_UUID"),
        required=False,
        help="Monte Carlo resource UUID for this Hive connection (env: MCD_RESOURCE_UUID)",
    )
    parser.add_argument(
        "--output-file",
        default="metadata_output.json",
        help="Path to write the output manifest (default: metadata_output.json)",
    )
    parser.add_argument(
        "--batch-size",
        type=int,
        default=DEFAULT_BATCH_SIZE,
        metavar="N",
        help=f"Max assets per POST (default: {DEFAULT_BATCH_SIZE})",
    )
    parser.add_argument(
        "--timeout",
        type=int,
        default=DEFAULT_TIMEOUT_SECONDS,
        metavar="SEC",
        help=f"HTTP timeout per request in seconds (default: {DEFAULT_TIMEOUT_SECONDS})",
    )
    args = parser.parse_args()

    if not args.hive_host:
        parser.error("--hive-host is required (or set HIVE_HOST)")
    if not args.key_id or not args.key_token:
        parser.error("--key-id and --key-token are required (or set MCD_INGEST_ID / MCD_INGEST_TOKEN)")
    if not args.resource_uuid:
        parser.error("--resource-uuid is required (or set MCD_RESOURCE_UUID)")

    manifest = collect(
        hive_host=args.hive_host,
        hive_port=args.hive_port,
    )

    push(
        manifest=manifest,
        resource_uuid=args.resource_uuid,
        key_id=args.key_id,
        key_token=args.key_token,
        batch_size=args.batch_size,
        timeout_seconds=args.timeout,
    )

    with open(args.output_file, "w") as fh:
        json.dump(manifest, fh, indent=2)
    print(f"Manifest written to {args.output_file}")
    print("Done.")


if __name__ == "__main__":
    main()

#!/usr/bin/env python3
"""
Collect Hive query logs from a local log file and push them to Monte Carlo
in one step.

Thin wrapper that calls ``collect()`` from ``collect_query_logs`` followed by
``push()`` from ``push_query_logs``, then writes the final manifest (with
``resource_uuid`` and ``invocation_id``) to ``--output-file``.

Substitution points
-------------------
- MCD_INGEST_ID    (env) / --key-id        (CLI) : Monte Carlo ingestion key ID
- MCD_INGEST_TOKEN (env) / --key-token      (CLI) : Monte Carlo ingestion key token
- MCD_RESOURCE_UUID    (env) / --resource-uuid  (CLI) : MC resource UUID (optional for query logs)
- --log-file                  path to local HiveServer2 log (default: /tmp/root/hive.log)
- --op-logs-dir               optional directory of per-query <queryId>.log files

Prerequisites
-------------
    pip install pycarlo python-dateutil python-dotenv

Usage
-----
    python collect_and_push_query_logs.py \\
        --key-id  <MCD_INGEST_ID> \\
        --key-token <MCD_INGEST_TOKEN> \\
        --resource-uuid <MCD_RESOURCE_UUID> \\
        --log-file /tmp/root/hive.log \\
        [--op-logs-dir /var/log/hive/operation_logs]
"""

import argparse
import json
import os

from collect_query_logs import collect
from push_query_logs import DEFAULT_BATCH_SIZE, DEFAULT_TIMEOUT_SECONDS, push


def main() -> None:
    parser = argparse.ArgumentParser(
        description="Collect Hive query logs from a local log file and push to Monte Carlo",
    )
    # Collect args
    parser.add_argument(
        "--log-file",
        default="/tmp/root/hive.log",
        help="Path to local HiveServer2 log file (default: /tmp/root/hive.log)",  # ← SUBSTITUTE: your log path
    )
    parser.add_argument(
        "--op-logs-dir",
        default=None,
        help=(
            "Directory containing per-query Hive operation logs (<queryId>.log). "
            "When provided, returned_rows is populated from SelectOperator RECORDS_OUT counts."
        ),
        # ← SUBSTITUTE: e.g. /var/log/hive/operation_logs or wherever Hive writes op logs
    )
    # Push / MC args
    parser.add_argument(
        "--key-id",
        default=os.environ.get("MCD_INGEST_ID"),
        help="Monte Carlo ingestion key ID (env: MCD_INGEST_ID)",
    )
    parser.add_argument(
        "--key-token",
        default=os.environ.get("MCD_INGEST_TOKEN"),
        help="Monte Carlo ingestion key token (env: MCD_INGEST_TOKEN)",
    )
    parser.add_argument(
        "--resource-uuid",
        default=os.environ.get("MCD_RESOURCE_UUID"),
        help="Monte Carlo resource UUID (optional for query logs) (env: MCD_RESOURCE_UUID)",
    )
    parser.add_argument(
        "--output-file",
        default="query_logs_output.json",
        help="Path to write the output manifest (default: query_logs_output.json)",
    )
    parser.add_argument(
        "--batch-size",
        type=int,
        default=DEFAULT_BATCH_SIZE,
        metavar="N",
        help=f"Max events per POST (default: {DEFAULT_BATCH_SIZE})",
    )
    parser.add_argument(
        "--timeout",
        type=int,
        default=DEFAULT_TIMEOUT_SECONDS,
        metavar="SEC",
        help=f"HTTP timeout per request in seconds (default: {DEFAULT_TIMEOUT_SECONDS})",
    )
    args = parser.parse_args()

    if not args.key_id or not args.key_token:
        parser.error("--key-id and --key-token are required (or set MCD_INGEST_ID / MCD_INGEST_TOKEN)")

    manifest = collect(log_file=args.log_file, op_logs_dir=args.op_logs_dir)

    push(
        manifest=manifest,
        key_id=args.key_id,
        key_token=args.key_token,
        resource_uuid=args.resource_uuid,
        batch_size=args.batch_size,
        timeout_seconds=args.timeout,
    )

    with open(args.output_file, "w") as fh:
        json.dump(manifest, fh, indent=2)
    print(f"Query log manifest written to {args.output_file}")
    print("Done.")


if __name__ == "__main__":
    main()

#!/usr/bin/env python3
"""
Extract table and column lineage from a local HiveServer2 log file — collection only.

Reads a plain-text Hive log file (not compressed), extracts SQL query blocks
from "Executing command" / "Starting command" entries, detects CTAS and
INSERT INTO ... SELECT patterns to build lineage edges, then writes a JSON
manifest file.

Can be run standalone via CLI or imported (use the ``collect()`` function).

Substitution points
-------------------
- --log-file  path to local HiveServer2 log (default: /tmp/root/hive.log)

Prerequisites
-------------
    pip install python-dotenv

Usage
-----
    python collect_lineage.py \\
        --log-file /tmp/root/hive.log \\
        --output-file lineage_output.json
"""

from __future__ import annotations

import argparse
import json
import re
from dataclasses import dataclass, field
from datetime import datetime, timezone

# ← SUBSTITUTE: set RESOURCE_TYPE to match your Monte Carlo connection type
RESOURCE_TYPE = "data-lake"

# Regex for CTAS: CREATE TABLE [IF NOT EXISTS] db.table AS SELECT ... FROM db.table
_CTAS_RE = re.compile(
    r"CREATE\s+TABLE\s+(?:IF\s+NOT\s+EXISTS\s+)?"
    r"(?P<dest_db>\w+)\.(?P<dest_table>\w+)"
    r".*?AS\s+SELECT\s+(?P<select_cols>.+?)\s+FROM\s+(?P<src_db>\w+)\.(?P<src_table>\w+)",
    re.IGNORECASE | re.DOTALL,
)

# Regex for INSERT INTO/OVERWRITE db.table SELECT ... FROM db.table
_INSERT_RE = re.compile(
    r"INSERT\s+(?:INTO|OVERWRITE)\s+(?:TABLE\s+)?(?P<dest_db>\w+)\.(?P<dest_table>\w+)"
    r".*?SELECT\s+(?P<select_cols>.+?)\s+FROM\s+(?P<src_db>\w+)\.(?P<src_table>\w+)",
    re.IGNORECASE | re.DOTALL,
)

# Regex to detect additional JOIN sources beyond the primary FROM clause
_JOIN_RE = re.compile(r"JOIN\s+(?P<src_db>\w+)\.(?P<src_table>\w+)", re.IGNORECASE)

# Simple column alias extraction: [alias.]col [AS dest]
_COL_RE = re.compile(r"(?:(\w+)\.)?(\w+)(?:\s+AS\s+(\w+))?", re.IGNORECASE)

# Hive string literals — strip before scanning so words inside 'status' AS ...
# are not treated as column refs
_STR_LITERAL_RE = re.compile(r"'(?:''|[^'])*'")

# ROW_NUMBER() OVER (...) AS alias — whole expression has no single source column;
# removing it avoids bogus tokens in col_mappings
_WINDOW_AS_ALIAS_RE = re.compile(
    r"\b(?:ROW_NUMBER|RANK|DENSE_RANK|NTILE)\s*\(\s*\)\s+OVER\s*\([^)]*\)\s+AS\s+\w+",
    re.IGNORECASE,
)

# Regex to pull query text out of Hive log "Executing/Starting command" lines
_COMMAND_START_RE = re.compile(
    r"(?:Executing|Starting)\s+command\(queryId=\S*\):\s+(?P<query>.+?)(?=\n\d{4}-\d{2}-\d{2}|\Z)",
    re.DOTALL,
)

# Tokens that are almost never real column names — SQL keywords, functions, casts, etc.
_SQL_SCAN_NOISE = frozenset(
    {
        "ROW_NUMBER", "RANK", "DENSE_RANK", "NTILE", "OVER", "PARTITION",
        "ORDER", "BY", "CASE", "WHEN", "THEN", "ELSE", "END", "AND", "OR",
        "NOT", "IN", "IS", "DISTINCT", "CAST", "CONVERT", "CURRENT_TIMESTAMP",
        "CURRENT_DATE", "TRUE", "FALSE", "NULL", "BETWEEN", "LIKE", "EXISTS",
        "ASC", "DESC", "LIMIT", "OFFSET", "GROUP", "HAVING", "UNION", "ALL",
        "INNER", "LEFT", "RIGHT", "FULL", "OUTER", "CROSS", "JOIN", "ON",
        "WHERE", "SELECT", "FROM", "AS", "STRING", "BIGINT", "INT", "SMALLINT",
        "TINYINT", "DOUBLE", "FLOAT", "REAL", "DECIMAL", "BOOLEAN", "DATE",
        "TIMESTAMP", "VARCHAR", "CHAR", "BINARY", "ARRAY", "MAP", "STRUCT",
        "SUM", "AVG", "COUNT", "MIN", "MAX", "STDDEV", "VARIANCE", "VAR_POP",
        "COALESCE", "IF", "SUBSTRING", "YEAR", "MONTH", "DAY", "LEAD", "LAG",
        "FIRST_VALUE", "LAST_VALUE",
    }
)


@dataclass
class _LineageEdge:
    dest_db: str
    dest_table: str
    sources: list[tuple[str, str]] = field(default_factory=list)
    # col_mappings: (dest_col, src_table, src_col)
    col_mappings: list[tuple[str, str, str]] = field(default_factory=list)


def _prepare_select_for_col_scan(select_clause: str) -> str:
    """Remove literals and window headers so _COL_RE sees fewer false positives."""
    s = _STR_LITERAL_RE.sub(" ", select_clause)
    s = _WINDOW_AS_ALIAS_RE.sub(" ", s)
    return s


def _dedupe_col_mappings(mappings: list[tuple[str, str, str]]) -> list[tuple[str, str, str]]:
    seen: set[tuple[str, str, str]] = set()
    out: list[tuple[str, str, str]] = []
    for t in mappings:
        if t in seen:
            continue
        seen.add(t)
        out.append(t)
    return out


def _extract_query_blocks(log_text: str) -> list[str]:
    """Extract individual SQL query strings from a Hive log file."""
    return [m.group("query").strip() for m in _COMMAND_START_RE.finditer(log_text)]


def _parse_select_cols(select_clause: str, src_table: str) -> list[tuple[str, str, str]]:
    """
    Lightweight column mapping: for each `alias.col AS dest` or `col AS dest`
    in the SELECT clause, return (dest_col, src_table, src_col).

    Strips string literals and window function headers first to reduce false
    positives, and filters out SQL keywords/noise tokens.
    """
    prepared = _prepare_select_for_col_scan(select_clause)
    mappings = []
    for m in _COL_RE.finditer(prepared):
        src_col = m.group(2)
        dest_col = m.group(3) or src_col
        if src_col.upper() in ("FROM", "SELECT", "WHERE", "JOIN", "ON", "AS", "*"):
            continue
        if src_col.upper() in _SQL_SCAN_NOISE or dest_col.upper() in _SQL_SCAN_NOISE:
            continue
        # After stripping 'literal' AS col, we get " AS col" — skip bare (col, col) with no source expr.
        if dest_col == src_col:
            prefix = prepared[: m.start()].rstrip()
            if prefix.upper().endswith("AS"):
                continue
        mappings.append((dest_col, src_table, src_col))
    return _dedupe_col_mappings(mappings)


def _parse_edges(queries: list[str]) -> list[_LineageEdge]:
    """Parse SQL query strings into _LineageEdge objects."""
    edges: dict[str, _LineageEdge] = {}

    for sql in queries:
        # Strip string literals to avoid false table/column matches inside quoted strings
        sql_clean = re.sub(r"\s+", " ", _STR_LITERAL_RE.sub(" ", sql)).strip()

        for pattern in (_CTAS_RE, _INSERT_RE):
            m = pattern.search(sql_clean)
            if not m:
                continue

            dest_db = m.group("dest_db").lower()
            dest_table = m.group("dest_table").lower()
            src_db = m.group("src_db").lower()
            src_table = m.group("src_table").lower()
            select_cols = m.group("select_cols")

            key = f"{dest_db}.{dest_table}"
            if key not in edges:
                edges[key] = _LineageEdge(dest_db=dest_db, dest_table=dest_table)

            edge = edges[key]
            src_pair = (src_db, src_table)
            if src_pair not in edge.sources:
                edge.sources.append(src_pair)

            # Pick up additional JOIN sources
            for jm in _JOIN_RE.finditer(sql_clean):
                jp = (jm.group("src_db").lower(), jm.group("src_table").lower())
                if jp not in edge.sources:
                    edge.sources.append(jp)

            edge.col_mappings.extend(_parse_select_cols(select_cols, src_table))
            break  # matched one pattern, move to next query

    # Deduplicate column mappings per edge (same INSERT may appear many times in HS2 logs)
    for e in edges.values():
        e.col_mappings = _dedupe_col_mappings(e.col_mappings)

    return list(edges.values())


def collect(log_file: str) -> dict:
    """
    Parse lineage edges from a HiveServer2 log file and return a manifest dict.

    Args:
        log_file: Path to a local HiveServer2 log file.

    Returns:
        Manifest dict with keys: resource_type, collected_at, edges.
        Each edge has destination, sources, and col_mappings lists.
    """
    print(f"Reading Hive log file: {log_file} ...")
    with open(log_file, errors="replace") as fh:
        log_text = fh.read()

    queries = _extract_query_blocks(log_text)
    print(f"  Extracted {len(queries)} query block(s).")

    edges = _parse_edges(queries)
    print(f"  Parsed {len(edges)} lineage edge(s).")

    manifest = {
        "resource_type": RESOURCE_TYPE,
        "collected_at": datetime.now(tz=timezone.utc).isoformat(),
        "edges": [
            {
                "destination": {"database": e.dest_db, "table": e.dest_table},
                "sources": [{"database": sdb, "table": stbl} for sdb, stbl in e.sources],
                "col_mappings": [
                    {"dest_col": dc, "src_table": st, "src_col": sc}
                    for dc, st, sc in e.col_mappings
                ],
            }
            for e in edges
        ],
    }
    return manifest


def main() -> None:
    parser = argparse.ArgumentParser(
        description="Extract Hive lineage from a local log file and write a JSON manifest",
    )
    parser.add_argument(
        "--log-file",
        default="/tmp/root/hive.log",
        help="Path to local HiveServer2 log file (default: /tmp/root/hive.log)",  # ← SUBSTITUTE: your log path
    )
    parser.add_argument(
        "--output-file",
        default="lineage_output.json",
        help="Path to write the lineage manifest (default: lineage_output.json)",
    )
    args = parser.parse_args()

    manifest = collect(log_file=args.log_file)

    if not manifest["edges"]:
        print("No lineage edges detected — no CTAS or INSERT INTO ... SELECT patterns found.")
        return

    with open(args.output_file, "w") as fh:
        json.dump(manifest, fh, indent=2)
    print(f"Lineage manifest written to {args.output_file}")
    print("Done.")


if __name__ == "__main__":
    main()

#!/usr/bin/env python3
"""
Collect table metadata from a Hive Metastore — collection only.

Connects to HiveServer2 (default port 10000), discovers all databases and
tables via SHOW DATABASES / SHOW TABLES, reads schema and table statistics
via DESCRIBE FORMATTED, then writes a JSON manifest file.

Can be run standalone via CLI or imported (use the ``collect()`` function).

Substitution points
-------------------
- HIVE_HOST         (env) / --hive-host   (CLI) : HiveServer2 hostname
- HIVE_PORT         (env) / --hive-port   (CLI) : HiveServer2 port (default 10000)

Prerequisites
-------------
    pip install pyhive python-dotenv

Usage
-----
    python collect_metadata.py \\
        --hive-host <HIVESERVER2_HOSTNAME> \\
        --output-file metadata_output.json
"""

import argparse
import json
import os
import re
from datetime import datetime, timezone

from pyhive import hive


def _check_available_memory(min_gb: float = 2.0) -> None:
    """Warn if available memory is below the threshold."""
    try:
        if hasattr(os, "sysconf"):  # Linux / macOS
            page_size = os.sysconf("SC_PAGE_SIZE")
            avail_pages = os.sysconf("SC_AVPHYS_PAGES")
            avail_gb = (page_size * avail_pages) / (1024 ** 3)
        else:
            return  # Windows — skip check
    except (ValueError, OSError):
        return
    if avail_gb < min_gb:
        print(
            f"WARNING: Only {avail_gb:.1f} GB of memory available "
            f"(minimum recommended: {min_gb:.1f} GB). "
            f"Consider reducing the number of databases/tables or increasing available memory."
        )

# ← SUBSTITUTE: set RESOURCE_TYPE to match your Monte Carlo connection type
RESOURCE_TYPE = "data-lake"

# Map Hive native types to SQL-standard uppercase types expected by Monte Carlo
_HIVE_TYPE_MAP: dict[str, str] = {
    "tinyint": "TINYINT",
    "smallint": "SMALLINT",
    "int": "INTEGER",
    "integer": "INTEGER",
    "bigint": "BIGINT",
    "float": "FLOAT",
    "double": "DOUBLE",
    "double precision": "DOUBLE",
    "decimal": "DECIMAL",
    "numeric": "DECIMAL",
    "boolean": "BOOLEAN",
    "string": "VARCHAR",
    "varchar": "VARCHAR",
    "char": "CHAR",
    "binary": "BINARY",
    "timestamp": "TIMESTAMP",
    "date": "DATE",
    "interval": "INTERVAL",
    "array": "ARRAY",
    "map": "MAP",
    "struct": "STRUCT",
    "uniontype": "UNION",
}

# ← SUBSTITUTE: add any internal table name prefixes you want to skip
_INTERNAL_TABLE_PREFIXES = ("tmp_", "__", "hive_")


def _normalize_hive_type(hive_type: str) -> str:
    """Uppercase and normalize a Hive type string to a SQL-standard form.

    Parametrized types like ``decimal(10,2)`` or ``varchar(255)`` keep their
    suffix; the base type is mapped through ``_HIVE_TYPE_MAP``.
    """
    lower = hive_type.lower().strip()
    base = lower.split("(")[0].strip()
    suffix = hive_type[len(base):].strip()  # preserve original params, e.g. decimal(10,2)
    return _HIVE_TYPE_MAP.get(base, base.upper()) + suffix


def _connect(host: str, port: int) -> hive.Connection:
    # ← SUBSTITUTE: update username/auth if your cluster requires Kerberos or LDAP
    return hive.connect(host=host, port=port, username="hadoop", auth="NONE")


def _fetch_rows(cursor, query: str) -> list[tuple]:
    """Execute a query and fetch results in memory-safe chunks."""
    cursor.execute(query)
    rows: list[tuple] = []
    while True:
        chunk = cursor.fetchmany(1000)
        if not chunk:
            break
        rows.extend(chunk)
    return rows


def _parse_describe_formatted(rows: list[tuple]) -> dict:
    """
    Parse DESCRIBE FORMATTED <db>.<table> output into a structured dict:
      columns, row_count, total_size, last_modified, description, created_on
    """
    result: dict = {
        "columns": [],
        "row_count": None,
        "total_size": None,
        "last_modified": None,
        "description": None,
        "created_on": None,
    }
    in_col_info = False
    in_table_info = False

    for row in rows:
        col_name = (row[0] or "").strip()
        data_type = (row[1] or "").strip()
        comment = (row[2] or "").strip() if len(row) > 2 else ""

        if col_name.startswith("# col_name"):
            in_col_info = True
            in_table_info = False
            continue
        if col_name.startswith("# Detailed Table Information"):
            in_col_info = False
            in_table_info = True
            continue
        if col_name.startswith("#"):
            in_col_info = False
            continue

        if in_col_info and col_name and data_type:
            result["columns"].append(
                {
                    "name": col_name,
                    "type": _normalize_hive_type(data_type),
                    "description": comment or None,
                }
            )

        if in_table_info:
            # Table Parameters rows have an empty col_name; key is in data_type, value in comment
            param_key = data_type.strip() if not col_name else col_name.strip().rstrip(":")
            param_val = (comment.strip() if not col_name else data_type.strip()) or ""

            if re.search(r"numRows", param_key, re.IGNORECASE):
                try:
                    result["row_count"] = int(param_val)
                except (ValueError, TypeError):
                    pass
            elif re.search(r"totalSize", param_key, re.IGNORECASE):
                try:
                    result["total_size"] = int(param_val)
                except (ValueError, TypeError):
                    pass
            elif re.search(r"last_modified_time", param_key, re.IGNORECASE):
                try:
                    result["last_modified"] = datetime.fromtimestamp(
                        int(param_val), tz=timezone.utc
                    ).isoformat()
                except (ValueError, TypeError):
                    pass
            elif re.search(r"^CreateTime", param_key):
                # e.g. "Wed Mar 18 20:15:40 UTC 2026"
                try:
                    result["created_on"] = datetime.strptime(
                        param_val, "%a %b %d %H:%M:%S %Z %Y"
                    ).replace(tzinfo=timezone.utc).isoformat()
                except (ValueError, TypeError):
                    pass
            elif param_key == "comment" and not result["description"] and param_val:
                result["description"] = param_val

    return result


def collect(
    hive_host: str,
    hive_port: int = 10000,
) -> dict:
    """
    Connect to HiveServer2, discover all databases and tables, and return a
    manifest dict with collected asset metadata.

    Args:
        hive_host: HiveServer2 hostname.
        hive_port: HiveServer2 port (default 10000).

    Returns:
        Manifest dict with keys: resource_type, collected_at, assets.
    """
    _check_available_memory()
    print(f"Connecting to HiveServer2 at {hive_host}:{hive_port} ...")
    conn = _connect(hive_host, hive_port)
    cursor = conn.cursor()
    assets: list[dict] = []

    print("Collecting table metadata ...")
    databases = [row[0] for row in _fetch_rows(cursor, "SHOW DATABASES")]
    print(f"  Found databases: {databases}")

    for db in databases:
        # ← SUBSTITUTE: add any system databases you want to skip
        if db in ("information_schema",):
            continue

        tables = _fetch_rows(cursor, f"SHOW TABLES IN {db}")
        table_names = [row[0] for row in tables]
        print(f"  {db}: {len(table_names)} table(s)")

        for table in table_names:
            if any(table.startswith(p) for p in _INTERNAL_TABLE_PREFIXES):
                continue

            try:
                desc_rows = _fetch_rows(cursor, f"DESCRIBE FORMATTED {db}.{table}")
            except Exception as exc:
                print(f"    WARNING: could not describe {db}.{table}: {exc}")
                continue

            info = _parse_describe_formatted(desc_rows)

            row_count = info["row_count"] if info["row_count"] and info["row_count"] > 0 else None
            byte_count = info["total_size"] if info["total_size"] and info["total_size"] > 0 else None

            assets.append(
                {
                    "database": db,
                    "schema": db,
                    "name": table,
                    "description": info["description"],
                    "created_on": info["created_on"],
                    "row_count": row_count,
                    "byte_count": byte_count,
                    "last_modified": info["last_modified"],
                    "fields": [
                        {"name": col["name"], "type": col["type"], "description": col["description"]}
                        for col in info["columns"]
                    ],
                }
            )
            print(
                f"    + {db}.{table} ({len(info['columns'])} columns, "
                f"desc={info['description']!r}, created={info['created_on']})"
            )

    cursor.close()
    conn.close()
    print(f"\nCollected {len(assets)} table(s).")

    manifest = {
        "resource_type": RESOURCE_TYPE,
        "collected_at": datetime.now(tz=timezone.utc).isoformat(),
        "assets": assets,
    }
    return manifest


def main() -> None:
    parser = argparse.ArgumentParser(
        description="Collect Hive table metadata and write a JSON manifest",
    )
    parser.add_argument(
        "--hive-host",
        default=os.environ.get("HIVE_HOST"),
        help="HiveServer2 hostname (env: HIVE_HOST)",  # ← SUBSTITUTE: your EMR master DNS or Hive host
    )
    parser.add_argument(
        "--hive-port",
        type=int,
        default=10000,
        help="HiveServer2 port (default: 10000)",  # ← SUBSTITUTE if your cluster uses a non-standard port
    )
    parser.add_argument(
        "--output-file",
        default="metadata_output.json",
        help="Path to write the output manifest (default: metadata_output.json)",
    )
    args = parser.parse_args()

    if not args.hive_host:
        parser.error("--hive-host is required (or set HIVE_HOST)")

    manifest = collect(
        hive_host=args.hive_host,
        hive_port=args.hive_port,
    )

    with open(args.output_file, "w") as fh:
        json.dump(manifest, fh, indent=2)
    print(f"Asset manifest written to {args.output_file}")
    print("Done.")


if __name__ == "__main__":
    main()

#!/usr/bin/env python3
"""
Collect Hive query logs from a local HiveServer2 log file — collection only.

Parses a plain-text HiveServer2 log for "Executing/Starting command" entries
to extract query text, query ID, start time and end time.  Optionally reads
per-query operation logs to populate ``returned_rows`` from SelectOperator
``RECORDS_OUT`` counters.  Deduplicates entries by query ID.

Can be run standalone via CLI or imported (use the ``collect()`` function).

Substitution points
-------------------
- --log-file       path to local HiveServer2 log (default: /tmp/root/hive.log)
- --op-logs-dir    optional directory of per-query <queryId>.log files

Prerequisites
-------------
    pip install python-dateutil python-dotenv

Usage
-----
    python collect_query_logs.py \\
        --log-file /tmp/root/hive.log \\
        [--op-logs-dir /var/log/hive/operation_logs] \\
        --output-file query_logs_output.json
"""

from __future__ import annotations

import argparse
import json
import re
from datetime import datetime, timezone
from io import StringIO
from pathlib import Path

from dateutil.parser import isoparse

# NOTE: the normalizer requires "hive-s3" — do not change to "hive" or "data-lake"
LOG_TYPE = "hive-s3"

# Matches the start of a new query block in the Hive log
_COMMAND_START_RE = re.compile(
    r"(Executing|Starting)\s+command\(queryId=(?P<query_id>\S*)\):\s+(?P<command>.*)$"
)

# Extracts returned row counts from per-query Hive operation logs
_RECORDS_OUT_RE = re.compile(r"RECORDS_OUT_OPERATOR_SEL_\d+:(\d+)")


def _parse_log_entries(log_text: str) -> list[dict]:
    """
    Parse a HiveServer2 log file and return a list of dicts:
      query_id, start_time (datetime), end_time (datetime), query (str)

    Each timestamped "Executing/Starting command" line starts a new entry.
    The previous entry's end_time is set to the timestamp of the next line.
    """
    entries = []
    query = ""
    query_id = ""
    start_time: datetime | None = None
    last_timestamp: datetime | None = None

    for line in StringIO(log_text):
        parts = line.split()
        if not parts:
            continue

        try:
            timestamp = isoparse(parts[0])
            if not timestamp.tzinfo:
                timestamp = timestamp.replace(tzinfo=timezone.utc)
        except ValueError:
            # Continuation line for a multi-line query
            if query:
                query += "\n" + line.rstrip()
            continue

        command_start = _COMMAND_START_RE.search(line)
        if command_start:
            # Emit the previous entry before starting a new one
            if query and start_time:
                entries.append(
                    {
                        "query_id": query_id,
                        "start_time": start_time,
                        "end_time": timestamp,
                        "query": query,
                    }
                )
            query_id = command_start.group("query_id")
            start_time = timestamp
            query = command_start.group("command").strip()
        elif query and start_time:
            # A timestamped non-command line closes the current entry
            entries.append(
                {
                    "query_id": query_id,
                    "start_time": start_time,
                    "end_time": timestamp,
                    "query": query,
                }
            )
            query = ""
            query_id = ""
            start_time = None

        last_timestamp = timestamp

    # Flush any trailing entry
    if query and start_time:
        end_time = last_timestamp or start_time
        entries.append(
            {
                "query_id": query_id,
                "start_time": start_time,
                "end_time": end_time,
                "query": query,
            }
        )

    return entries


def _load_returned_rows(op_logs_dir: str) -> dict[str, int]:
    """
    Scan a directory of per-query Hive operation logs (named <queryId>.log) and
    return a mapping of query_id -> rows returned.

    The row count is taken from the last RECORDS_OUT_OPERATOR_SEL_N value in
    each file, which reflects the final number of rows delivered to the client.
    """
    rows_by_id: dict[str, int] = {}
    for log_file in Path(op_logs_dir).glob("*.log"):
        query_id = log_file.stem
        last_count: int | None = None
        try:
            text = log_file.read_text(errors="replace")
        except OSError:
            continue
        for m in _RECORDS_OUT_RE.finditer(text):
            last_count = int(m.group(1))
        if last_count is not None:
            rows_by_id[query_id] = last_count
    return rows_by_id


def _build_query_log_entries(
    raw_entries: list[dict],
    rows_by_id: dict[str, int] | None = None,
) -> list[dict]:
    """
    Deduplicate raw log entries by query_id and enrich with returned_rows.

    Returns plain dicts so that ``push_query_logs.py`` can reconstruct
    QueryLogEntry objects from the JSON manifest.
    """
    seen: set[str] = set()
    entries = []
    for r in raw_entries:
        qid = r["query_id"]
        if qid and qid in seen:
            continue
        if qid:
            seen.add(qid)

        returned_rows: int | None = rows_by_id.get(qid) if rows_by_id and qid else None

        entries.append(
            {
                "query_id": qid or None,
                "start_time": r["start_time"].isoformat(),
                "end_time": r["end_time"].isoformat(),
                "query_text": r["query"],
                "user": "hadoop",  # ← SUBSTITUTE: set the user appropriate for your cluster
                "returned_rows": returned_rows,
            }
        )
    return entries


def collect(
    log_file: str,
    op_logs_dir: str | None = None,
) -> dict:
    """
    Parse query log entries from a HiveServer2 log file and return a manifest dict.

    Args:
        log_file: Path to a local HiveServer2 log file.
        op_logs_dir: Optional directory containing per-query operation logs
                     (<queryId>.log). When provided, returned_rows is populated
                     from SelectOperator RECORDS_OUT counts.

    Returns:
        Manifest dict with keys: log_type, collected_at, entry_count,
        window_start, window_end, queries.
    """
    print(f"Reading Hive log file: {log_file} ...")
    with open(log_file, errors="replace") as fh:
        log_text = fh.read()

    raw_entries = _parse_log_entries(log_text)
    print(f"  Parsed {len(raw_entries)} query log entry/entries.")

    if not raw_entries:
        print("No query log entries found.")
        return {
            "log_type": LOG_TYPE,
            "collected_at": datetime.now(tz=timezone.utc).isoformat(),
            "entry_count": 0,
            "window_start": None,
            "window_end": None,
            "queries": [],
        }

    rows_by_id: dict[str, int] | None = None
    if op_logs_dir:
        rows_by_id = _load_returned_rows(op_logs_dir)
        print(f"  Loaded row counts for {len(rows_by_id)} query/queries from {op_logs_dir}")

    queries = _build_query_log_entries(raw_entries, rows_by_id)

    start_times = [r["start_time"] for r in raw_entries]
    end_times = [r["end_time"] for r in raw_entries]

    manifest = {
        "log_type": LOG_TYPE,
        "collected_at": datetime.now(tz=timezone.utc).isoformat(),
        "entry_count": len(queries),
        "window_start": min(start_times).isoformat() if start_times else None,
        "window_end": max(end_times).isoformat() if end_times else None,
        "queries": [
            {
                "query_id": q["query_id"],
                "start_time": q["start_time"],
                "end_time": q["end_time"],
                "query": q["query_text"],
                "user": q["user"],
                "returned_rows": q["returned_rows"],
            }
            for q in queries
        ],
    }
    return manifest


def main() -> None:
    parser = argparse.ArgumentParser(
        description="Collect Hive query logs from a local log file and write a JSON manifest",
    )
    parser.add_argument(
        "--log-file",
        default="/tmp/root/hive.log",
        help="Path to local HiveServer2 log file (default: /tmp/root/hive.log)",  # ← SUBSTITUTE: your log path
    )
    parser.add_argument(
        "--op-logs-dir",
        default=None,
        help=(
            "Directory containing per-query Hive operation logs (<queryId>.log). "
            "When provided, returned_rows is populated from SelectOperator RECORDS_OUT counts."
        ),
        # ← SUBSTITUTE: e.g. /var/log/hive/operation_logs or wherever Hive writes op logs
    )
    parser.add_argument(
        "--output-file",
        default="query_logs_output.json",
        help="Path to write the output manifest (default: query_logs_output.json)",
    )
    args = parser.parse_args()

    manifest = collect(log_file=args.log_file, op_logs_dir=args.op_logs_dir)

    with open(args.output_file, "w") as fh:
        json.dump(manifest, fh, indent=2)
    print(f"Query log manifest written to {args.output_file}")
    print("Done.")


if __name__ == "__main__":
    main()

#!/usr/bin/env python3
"""
Push a collected Hive lineage manifest to Monte Carlo — push only.

Reads a JSON manifest produced by ``collect_lineage.py``, builds LineageEvent
objects (table-level or column-level), and calls ``send_lineage`` in batches.
The manifest is updated in-place with ``resource_uuid`` and ``invocation_id``
after a successful push.

Can be run standalone via CLI or imported (use the ``push()`` function).

Substitution points
-------------------
- MCD_INGEST_ID    (env) / --key-id        (CLI) : Monte Carlo ingestion key ID
- MCD_INGEST_TOKEN (env) / --key-token      (CLI) : Monte Carlo ingestion key token
- MCD_RESOURCE_UUID    (env) / --resource-uuid  (CLI) : MC resource UUID for this connection

Prerequisites
-------------
    pip install pycarlo python-dotenv

Usage (table-level):
    python push_lineage.py \\
        --key-id  <MCD_INGEST_ID> \\
        --key-token <MCD_INGEST_TOKEN> \\
        --resource-uuid <MCD_RESOURCE_UUID> \\
        --input-file lineage_output.json

Usage (column-level):
    python push_lineage.py ... --column-lineage
"""

import argparse
import json
import os
from concurrent.futures import ThreadPoolExecutor, as_completed

from pycarlo.core import Client, Session
from pycarlo.features.ingestion import IngestionService
from pycarlo.features.ingestion.models import (
    ColumnLineageField,
    ColumnLineageSourceField,
    LineageAssetRef,
    LineageEvent,
)

# ← SUBSTITUTE: set RESOURCE_TYPE to match your Monte Carlo connection type
RESOURCE_TYPE = "data-lake"

# ← SUBSTITUTE: default batch size for lineage push (events per request)
DEFAULT_BATCH_SIZE = 500

# ← SUBSTITUTE: HTTP timeout for MC ingestion requests (seconds)
DEFAULT_TIMEOUT_SECONDS = 120


def _build_table_lineage(edges_data: list[dict]) -> list[LineageEvent]:
    """Build table-level LineageEvent objects from raw edge dicts."""
    events = []
    for edge in edges_data:
        sources = edge.get("sources", [])
        if not sources:
            continue
        dest = edge["destination"]
        events.append(
            LineageEvent(
                destination=LineageAssetRef(
                    type="TABLE",
                    name=dest["table"],
                    database=dest["database"],
                    schema=dest["database"],
                ),
                sources=[
                    LineageAssetRef(
                        type="TABLE",
                        name=src["table"],
                        database=src["database"],
                        schema=src["database"],
                    )
                    for src in sources
                ],
            )
        )
    return events


def _build_column_lineage(edges_data: list[dict]) -> list[LineageEvent]:
    """Build column-level LineageEvent objects from raw edge dicts."""
    events = []
    for edge in edges_data:
        sources = edge.get("sources", [])
        if not sources:
            continue

        dest = edge["destination"]
        dest_asset_id = f"{dest['database']}__{dest['table']}"
        source_asset_ids = {
            (src["database"], src["table"]): f"{src['database']}__{src['table']}"
            for src in sources
        }

        col_fields: dict[str, ColumnLineageField] = {}
        for mapping in edge.get("col_mappings", []):
            dest_col = mapping["dest_col"]
            src_table = mapping["src_table"]
            src_col = mapping["src_col"]
            # Find the matching source db for this src_table
            src_db = next(
                (src["database"] for src in sources if src["table"] == src_table),
                dest["database"],
            )
            src_aid = source_asset_ids.get((src_db, src_table), f"{src_db}__{src_table}")
            if dest_col not in col_fields:
                col_fields[dest_col] = ColumnLineageField(name=dest_col, source_fields=[])
            col_fields[dest_col].source_fields.append(
                ColumnLineageSourceField(asset_id=src_aid, field_name=src_col)
            )

        events.append(
            LineageEvent(
                destination=LineageAssetRef(
                    type="TABLE",
                    name=dest["table"],
                    database=dest["database"],
                    schema=dest["database"],
                    asset_id=dest_asset_id,
                ),
                sources=[
                    LineageAssetRef(
                        type="TABLE",
                        name=src["table"],
                        database=src["database"],
                        schema=src["database"],
                        asset_id=source_asset_ids[(src["database"], src["table"])],
                    )
                    for src in sources
                ],
                fields=list(col_fields.values()) if col_fields else None,
            )
        )
    return events


def push(
    manifest: dict,
    resource_uuid: str,
    key_id: str,
    key_token: str,
    column_lineage: bool = False,
    batch_size: int = DEFAULT_BATCH_SIZE,
    timeout_seconds: int = DEFAULT_TIMEOUT_SECONDS,
) -> str | None:
    """
    Push collected lineage to Monte Carlo and update the manifest in-place.

    Events are sent in batches of ``batch_size`` (default 500) to avoid
    oversized payloads.  Supports both table-level and column-level lineage.

    Args:
        manifest: Dict loaded from a ``collect_lineage.py`` output file.
        resource_uuid: MC resource UUID for this Hive connection.
        key_id: MC ingestion key ID.
        key_token: MC ingestion key token.
        column_lineage: When True, push column-level lineage; otherwise table-level.
        batch_size: Events per POST request (default 500).
        timeout_seconds: HTTP timeout per request (default 120).

    Returns:
        The last invocation ID string if returned by MC, otherwise None.
    """
    resource_type = manifest.get("resource_type", RESOURCE_TYPE)
    edges_data = manifest.get("edges", [])

    if column_lineage:
        events = _build_column_lineage(edges_data)
        label = "column-level"
    else:
        events = _build_table_lineage(edges_data)
        label = "table-level"

    print(f"Loaded {len(events)} {label} lineage event(s) from manifest")

    if not events:
        print("No lineage events to push.")
        manifest["resource_uuid"] = resource_uuid
        manifest["invocation_id"] = None
        return None

    # Split into batches
    batch_list = []
    for i in range(0, len(events), batch_size):
        batch_list.append(events[i : i + batch_size])
    total_batches = len(batch_list)

    def _push_batch(batch: list, batch_num: int) -> str | None:
        """Push a single batch using a dedicated Session (thread-safe)."""
        print(f"  Pushing batch {batch_num}/{total_batches} ({len(batch)} events) ...")
        client = Client(session=Session(mcd_id=key_id, mcd_token=key_token, scope="Ingestion"))
        service = IngestionService(mc_client=client)
        result = service.send_lineage(
            resource_uuid=resource_uuid,
            resource_type=resource_type,
            events=batch,
        )
        invocation_id = service.extract_invocation_id(result)
        if invocation_id:
            print(f"    Batch {batch_num}: invocation_id={invocation_id}")
        return invocation_id

    # Push batches in parallel (each thread gets its own pycarlo Session)
    max_workers = min(4, total_batches)
    invocation_ids: list[str | None] = [None] * total_batches

    with ThreadPoolExecutor(max_workers=max_workers) as pool:
        futures = {
            pool.submit(_push_batch, batch, i + 1): i
            for i, batch in enumerate(batch_list)
        }
        for future in as_completed(futures):
            idx = futures[future]
            try:
                invocation_ids[idx] = future.result()
            except Exception as exc:
                print(f"    ERROR pushing batch {idx + 1}: {exc}")
                raise

    print(f"  All {total_batches} batches pushed ({max_workers} workers)")

    manifest["resource_uuid"] = resource_uuid
    manifest["invocation_id"] = invocation_ids[-1] if invocation_ids else None
    if len([i for i in invocation_ids if i]) > 1:
        manifest["invocation_ids"] = invocation_ids
    elif "invocation_ids" in manifest:
        del manifest["invocation_ids"]

    return manifest.get("invocation_id")


def main() -> None:
    parser = argparse.ArgumentParser(
        description="Push a collected Hive lineage manifest to Monte Carlo",
    )
    parser.add_argument(
        "--key-id",
        default=os.environ.get("MCD_INGEST_ID"),
        help="Monte Carlo ingestion key ID (env: MCD_INGEST_ID)",
    )
    parser.add_argument(
        "--key-token",
        default=os.environ.get("MCD_INGEST_TOKEN"),
        help="Monte Carlo ingestion key token (env: MCD_INGEST_TOKEN)",
    )
    parser.add_argument(
        "--resource-uuid",
        default=os.environ.get("MCD_RESOURCE_UUID"),
        help="Monte Carlo resource UUID for this Hive connection (env: MCD_RESOURCE_UUID)",
    )
    parser.add_argument(
        "--input-file",
        default="lineage_output.json",
        help="Path to the JSON manifest written by collect_lineage.py (default: lineage_output.json)",
    )
    parser.add_argument(
        "--column-lineage",
        action="store_true",
        help="Push column-level lineage instead of table-level",
    )
    parser.add_argument(
        "--batch-size",
        type=int,
        default=DEFAULT_BATCH_SIZE,
        metavar="N",
        help=f"Max events per POST (default: {DEFAULT_BATCH_SIZE})",
    )
    parser.add_argument(
        "--timeout",
        type=int,
        default=DEFAULT_TIMEOUT_SECONDS,
        metavar="SEC",
        help=f"HTTP timeout per request in seconds (default: {DEFAULT_TIMEOUT_SECONDS})",
    )
    args = parser.parse_args()

    if not args.key_id or not args.key_token:
        parser.error("--key-id and --key-token are required (or set MCD_INGEST_ID / MCD_INGEST_TOKEN)")
    if not args.resource_uuid:
        parser.error("--resource-uuid is required (or set MCD_RESOURCE_UUID)")

    with open(args.input_file) as fh:
        manifest = json.load(fh)

    push(
        manifest=manifest,
        resource_uuid=args.resource_uuid,
        key_id=args.key_id,
        key_token=args.key_token,
        column_lineage=args.column_lineage,
        batch_size=args.batch_size,
        timeout_seconds=args.timeout,
    )

    with open(args.input_file, "w") as fh:
        json.dump(manifest, fh, indent=2)
    print(f"Manifest updated in-place: {args.input_file}")
    print("Done.")


if __name__ == "__main__":
    main()

#!/usr/bin/env python3
"""
Push a collected Hive metadata manifest to Monte Carlo — push only.

Reads a JSON manifest produced by ``collect_metadata.py``, builds
RelationalAsset objects, and calls ``send_metadata`` in batches.  The manifest
is updated in-place with ``resource_uuid`` and ``invocation_id`` after a
successful push.

Can be run standalone via CLI or imported (use the ``push()`` function).

Substitution points
-------------------
- MCD_INGEST_ID    (env) / --key-id        (CLI) : Monte Carlo ingestion key ID
- MCD_INGEST_TOKEN (env) / --key-token      (CLI) : Monte Carlo ingestion key token
- MCD_RESOURCE_UUID    (env) / --resource-uuid  (CLI) : MC resource UUID for this connection

Prerequisites
-------------
    pip install pycarlo python-dotenv

Usage
-----
    python push_metadata.py \\
        --key-id  <MCD_INGEST_ID> \\
        --key-token <MCD_INGEST_TOKEN> \\
        --resource-uuid <MCD_RESOURCE_UUID> \\
        --input-file metadata_output.json
"""

import argparse
import json
import os
from concurrent.futures import ThreadPoolExecutor, as_completed
from datetime import datetime, timezone

from pycarlo.core import Client, Session
from pycarlo.features.ingestion import IngestionService
from pycarlo.features.ingestion.models import (
    AssetField,
    AssetFreshness,
    AssetMetadata,
    AssetVolume,
    RelationalAsset,
)

# ← SUBSTITUTE: default batch size for metadata push (assets per request)
DEFAULT_BATCH_SIZE = 500

# ← SUBSTITUTE: HTTP timeout for MC ingestion requests (seconds)
DEFAULT_TIMEOUT_SECONDS = 120


def _build_assets(manifest: dict) -> list[RelationalAsset]:
    """Rebuild RelationalAsset objects from a collected metadata manifest."""
    assets = []
    for a in manifest.get("assets", []):
        fields = [
            AssetField(
                name=f["name"],
                type=f["type"],
                description=f.get("description"),
            )
            for f in a.get("fields", [])
        ]

        volume = None
        row_count = a.get("row_count")
        byte_count = a.get("byte_count")
        if row_count or byte_count:
            volume = AssetVolume(
                row_count=row_count if row_count and row_count > 0 else None,
                byte_count=byte_count if byte_count and byte_count > 0 else None,
            )

        freshness = None
        last_modified = a.get("last_modified")
        if last_modified:
            freshness = AssetFreshness(last_update_time=last_modified)

        assets.append(
            RelationalAsset(
                type="TABLE",
                metadata=AssetMetadata(
                    name=a["name"],
                    database=a["database"],
                    schema=a["schema"],
                    description=a.get("description"),
                    created_on=a.get("created_on"),
                ),
                fields=fields,
                volume=volume,
                freshness=freshness,
            )
        )
    return assets


def push(
    manifest: dict,
    resource_uuid: str,
    key_id: str,
    key_token: str,
    batch_size: int = DEFAULT_BATCH_SIZE,
    timeout_seconds: int = DEFAULT_TIMEOUT_SECONDS,
) -> str | None:
    """
    Push collected metadata to Monte Carlo and update the manifest in-place.

    Assets are sent in batches of ``batch_size`` (default 500) to avoid
    oversized payloads.  The manifest is enriched with ``resource_uuid``
    and the last ``invocation_id`` from the response.

    Args:
        manifest: Dict loaded from a ``collect_metadata.py`` output file.
        resource_uuid: MC resource UUID for this Hive connection.
        key_id: MC ingestion key ID.
        key_token: MC ingestion key token.
        batch_size: Assets per POST request (default 500).
        timeout_seconds: HTTP timeout per request (default 120).

    Returns:
        The last invocation ID string if returned by MC, otherwise None.
    """
    resource_type = manifest.get("resource_type", "data-lake")

    assets = _build_assets(manifest)
    n = len(assets)

    print(f"Loaded {n} asset(s) from manifest")

    # Split into batches
    batch_list = []
    for i in range(0, max(n, 1), batch_size):
        batch_list.append(assets[i : i + batch_size])
    total_batches = len(batch_list)

    def _push_batch(batch: list, batch_num: int) -> str | None:
        """Push a single batch using a dedicated Session (thread-safe)."""
        client = Client(session=Session(mcd_id=key_id, mcd_token=key_token, scope="Ingestion"))
        service = IngestionService(mc_client=client)
        result = service.send_metadata(
            resource_uuid=resource_uuid,
            resource_type=resource_type,
            events=batch,
        )
        invocation_id = service.extract_invocation_id(result)
        print(f"  Pushed batch {batch_num}/{total_batches} ({len(batch)} assets) — invocation_id={invocation_id}")
        return invocation_id

    # Push batches in parallel (each thread gets its own pycarlo Session)
    max_workers = min(4, total_batches)
    invocation_ids: list[str | None] = [None] * total_batches

    with ThreadPoolExecutor(max_workers=max_workers) as pool:
        futures = {
            pool.submit(_push_batch, batch, i + 1): i
            for i, batch in enumerate(batch_list)
        }
        for future in as_completed(futures):
            idx = futures[future]
            try:
                invocation_ids[idx] = future.result()
            except Exception as exc:
                print(f"    ERROR pushing batch {idx + 1}: {exc}")
                raise

    print(f"  All {total_batches} batches pushed ({max_workers} workers)")

    manifest["resource_uuid"] = resource_uuid
    manifest["invocation_id"] = invocation_ids[-1] if invocation_ids else None
    if len([i for i in invocation_ids if i]) > 1:
        manifest["invocation_ids"] = invocation_ids
    elif "invocation_ids" in manifest:
        del manifest["invocation_ids"]

    return manifest.get("invocation_id")


def main() -> None:
    parser = argparse.ArgumentParser(
        description="Push a collected Hive metadata manifest to Monte Carlo",
    )
    parser.add_argument(
        "--key-id",
        default=os.environ.get("MCD_INGEST_ID"),
        help="Monte Carlo ingestion key ID (env: MCD_INGEST_ID)",  # ← SUBSTITUTE env var name if different
    )
    parser.add_argument(
        "--key-token",
        default=os.environ.get("MCD_INGEST_TOKEN"),
        help="Monte Carlo ingestion key token (env: MCD_INGEST_TOKEN)",  # ← SUBSTITUTE env var name if different
    )
    parser.add_argument(
        "--resource-uuid",
        default=os.environ.get("MCD_RESOURCE_UUID"),
        required=False,
        help="Monte Carlo resource UUID for this Hive connection (env: MCD_RESOURCE_UUID)",
    )
    parser.add_argument(
        "--input-file",
        default="metadata_output.json",
        help="Path to the JSON manifest written by collect_metadata.py (default: metadata_output.json)",
    )
    parser.add_argument(
        "--batch-size",
        type=int,
        default=DEFAULT_BATCH_SIZE,
        metavar="N",
        help=f"Max assets per POST (default: {DEFAULT_BATCH_SIZE})",
    )
    parser.add_argument(
        "--timeout",
        type=int,
        default=DEFAULT_TIMEOUT_SECONDS,
        metavar="SEC",
        help=f"HTTP timeout per request in seconds (default: {DEFAULT_TIMEOUT_SECONDS})",
    )
    args = parser.parse_args()

    if not args.key_id or not args.key_token:
        parser.error("--key-id and --key-token are required (or set MCD_INGEST_ID / MCD_INGEST_TOKEN)")
    if not args.resource_uuid:
        parser.error("--resource-uuid is required (or set MCD_RESOURCE_UUID)")

    with open(args.input_file) as fh:
        manifest = json.load(fh)

    push(
        manifest=manifest,
        resource_uuid=args.resource_uuid,
        key_id=args.key_id,
        key_token=args.key_token,
        batch_size=args.batch_size,
        timeout_seconds=args.timeout,
    )

    with open(args.input_file, "w") as fh:
        json.dump(manifest, fh, indent=2)
    print(f"Manifest updated in-place: {args.input_file}")
    print("Done.")


if __name__ == "__main__":
    main()

#!/usr/bin/env python3
"""
Push a collected Hive query log manifest to Monte Carlo — push only.

Reads a JSON manifest produced by ``collect_query_logs.py``, builds
QueryLogEntry objects, and calls ``send_query_logs`` in batches.  The manifest
is updated in-place with ``resource_uuid`` and ``invocation_id`` after a
successful push.

Can be run standalone via CLI or imported (use the ``push()`` function).

Substitution points
-------------------
- MCD_INGEST_ID    (env) / --key-id        (CLI) : Monte Carlo ingestion key ID
- MCD_INGEST_TOKEN (env) / --key-token      (CLI) : Monte Carlo ingestion key token
- MCD_RESOURCE_UUID    (env) / --resource-uuid  (CLI) : MC resource UUID (optional for query logs)

Prerequisites
-------------
    pip install pycarlo python-dateutil python-dotenv

Usage
-----
    python push_query_logs.py \\
        --key-id  <MCD_INGEST_ID> \\
        --key-token <MCD_INGEST_TOKEN> \\
        --resource-uuid <MCD_RESOURCE_UUID> \\
        --input-file query_logs_output.json
"""

import argparse
import json
import os
from concurrent.futures import ThreadPoolExecutor, as_completed
from datetime import datetime, timezone

from dateutil.parser import isoparse

from pycarlo.core import Client, Session
from pycarlo.features.ingestion import IngestionService
from pycarlo.features.ingestion.models import QueryLogEntry

# ← SUBSTITUTE: default batch size for query log push (events per request)
# Query logs include full SQL text — keep batches small to stay under the 1 MB
# compressed payload limit.  50 entries can trigger 413 on active warehouses.
DEFAULT_BATCH_SIZE = 100

# ← SUBSTITUTE: HTTP timeout for MC ingestion requests (seconds)
DEFAULT_TIMEOUT_SECONDS = 120

# Truncate query_text longer than this to prevent 413 errors.
# Some SQL statements (e.g., generated by BI tools) can be 100KB+ and blow up
# compressed payloads even at small batch sizes.
_MAX_QUERY_TEXT_LEN = 10_000


def _build_events(manifest: dict) -> list[QueryLogEntry]:
    """
    Rebuild QueryLogEntry objects from a collected query log manifest.

    ISO timestamp strings are parsed back to datetime.  Entries are
    deduplicated by query_id.
    """
    seen: set[str] = set()
    events = []
    truncated = 0
    for q in manifest.get("queries", []):
        qid = q.get("query_id")
        if qid and qid in seen:
            continue
        if qid:
            seen.add(qid)

        start_time = isoparse(q["start_time"])
        if not start_time.tzinfo:
            start_time = start_time.replace(tzinfo=timezone.utc)

        end_time = isoparse(q["end_time"])
        if not end_time.tzinfo:
            end_time = end_time.replace(tzinfo=timezone.utc)

        query_text = q.get("query") or ""

        # Truncate very long SQL to prevent 413 Request Too Large
        if len(query_text) > _MAX_QUERY_TEXT_LEN:
            query_text = query_text[:_MAX_QUERY_TEXT_LEN] + "... [TRUNCATED]"
            truncated += 1

        events.append(
            QueryLogEntry(
                start_time=start_time,
                end_time=end_time,
                query_text=query_text,
                query_id=qid or None,
                user=q.get("user", "hadoop"),  # ← SUBSTITUTE: set the user appropriate for your cluster
                returned_rows=q.get("returned_rows"),
            )
        )
    if truncated:
        print(f"  Truncated {truncated} query text(s) exceeding {_MAX_QUERY_TEXT_LEN} chars")
    return events


def push(
    manifest: dict,
    key_id: str,
    key_token: str,
    resource_uuid: str | None = None,
    batch_size: int = DEFAULT_BATCH_SIZE,
    timeout_seconds: int = DEFAULT_TIMEOUT_SECONDS,
) -> str | None:
    """
    Push collected query logs to Monte Carlo and update the manifest in-place.

    Events are sent in batches of ``batch_size`` (default 100) to avoid
    oversized payloads.

    Args:
        manifest: Dict loaded from a ``collect_query_logs.py`` output file.
        key_id: MC ingestion key ID.
        key_token: MC ingestion key token.
        resource_uuid: Optional MC resource UUID.
        batch_size: Events per POST request (default 100).
        timeout_seconds: HTTP timeout per request (default 120).

    Returns:
        The last invocation ID string if returned by MC, otherwise None.
    """
    log_type = manifest.get("log_type", "hive-s3")

    events = _build_events(manifest)
    n = len(events)
    print(f"Loaded {n} query log entry/entries from manifest")

    if not events:
        print("No query log entries to push.")
        manifest["log_type"] = log_type
        if resource_uuid is not None:
            manifest["resource_uuid"] = resource_uuid
        manifest["invocation_id"] = None
        return None

    # Split into batches
    batch_list = []
    for i in range(0, n, batch_size):
        batch_list.append(events[i : i + batch_size])
    total_batches = len(batch_list)

    def _push_batch(batch: list, batch_num: int) -> str | None:
        """Push a single batch using a dedicated Session (thread-safe)."""
        client = Client(session=Session(mcd_id=key_id, mcd_token=key_token, scope="Ingestion"))
        service = IngestionService(mc_client=client)
        result = service.send_query_logs(
            resource_uuid=resource_uuid,
            log_type=log_type,
            events=batch,
        )
        invocation_id = service.extract_invocation_id(result)
        print(f"  Pushed batch {batch_num}/{total_batches} ({len(batch)} entries) — invocation_id={invocation_id}")
        return invocation_id

    # Push batches in parallel (each thread gets its own pycarlo Session)
    max_workers = min(4, total_batches)
    invocation_ids: list[str | None] = [None] * total_batches

    with ThreadPoolExecutor(max_workers=max_workers) as pool:
        futures = {
            pool.submit(_push_batch, batch, i + 1): i
            for i, batch in enumerate(batch_list)
        }
        for future in as_completed(futures):
            idx = futures[future]
            try:
                invocation_ids[idx] = future.result()
            except Exception as exc:
                print(f"    ERROR pushing batch {idx + 1}: {exc}")
                raise

    print(f"  All {total_batches} batches pushed ({max_workers} workers)")

    manifest["log_type"] = log_type
    if resource_uuid is not None:
        manifest["resource_uuid"] = resource_uuid
    manifest["invocation_id"] = invocation_ids[-1] if invocation_ids else None
    if len([i for i in invocation_ids if i]) > 1:
        manifest["invocation_ids"] = invocation_ids
    elif "invocation_ids" in manifest:
        del manifest["invocation_ids"]

    return manifest.get("invocation_id")


def main() -> None:
    parser = argparse.ArgumentParser(
        description="Push a collected Hive query log manifest to Monte Carlo",
    )
    parser.add_argument(
        "--key-id",
        default=os.environ.get("MCD_INGEST_ID"),
        help="Monte Carlo ingestion key ID (env: MCD_INGEST_ID)",
    )
    parser.add_argument(
        "--key-token",
        default=os.environ.get("MCD_INGEST_TOKEN"),
        help="Monte Carlo ingestion key token (env: MCD_INGEST_TOKEN)",
    )
    parser.add_argument(
        "--resource-uuid",
        default=os.environ.get("MCD_RESOURCE_UUID"),
        help="Monte Carlo resource UUID (optional for query logs) (env: MCD_RESOURCE_UUID)",
    )
    parser.add_argument(
        "--input-file",
        default="query_logs_output.json",
        help="Path to the JSON manifest written by collect_query_logs.py (default: query_logs_output.json)",
    )
    parser.add_argument(
        "--batch-size",
        type=int,
        default=DEFAULT_BATCH_SIZE,
        metavar="N",
        help=f"Max events per POST (default: {DEFAULT_BATCH_SIZE})",
    )
    parser.add_argument(
        "--timeout",
        type=int,
        default=DEFAULT_TIMEOUT_SECONDS,
        metavar="SEC",
        help=f"HTTP timeout per request in seconds (default: {DEFAULT_TIMEOUT_SECONDS})",
    )
    args = parser.parse_args()

    if not args.key_id or not args.key_token:
        parser.error("--key-id and --key-token are required (or set MCD_INGEST_ID / MCD_INGEST_TOKEN)")

    with open(args.input_file) as fh:
        manifest = json.load(fh)

    push(
        manifest=manifest,
        key_id=args.key_id,
        key_token=args.key_token,
        resource_uuid=args.resource_uuid,
        batch_size=args.batch_size,
        timeout_seconds=args.timeout,
    )

    with open(args.input_file, "w") as fh:
        json.dump(manifest, fh, indent=2)
    print(f"Manifest updated in-place: {args.input_file}")
    print("Done.")


if __name__ == "__main__":
    main()

"""
Redshift — Lineage Collect & Push (combined)
==============================================
Collects table-level lineage from Redshift by parsing query history, then pushes
the derived lineage events to Monte Carlo via the push ingestion API.

This script imports and calls collect() from collect_lineage and push() from
push_lineage, running both in sequence.

Substitution points (search for "← SUBSTITUTE"):
  - REDSHIFT_HOST / REDSHIFT_DB / REDSHIFT_USER / REDSHIFT_PASSWORD : connection
  - LOOKBACK_HOURS    : how far back to scan query history (default 24 h)
  - MCD_INGEST_ID / MCD_INGEST_TOKEN : Monte Carlo API credentials
  - MCD_RESOURCE_UUID  : UUID of the Redshift connection in Monte Carlo
  - PUSH_BATCH_SIZE   : number of events per API call (default 500)

Prerequisites:
  pip install psycopg2-binary pycarlo
"""

from __future__ import annotations

import argparse
import logging
import os

from collect_lineage import LOOKBACK_HOURS, collect
from push_lineage import DEFAULT_BATCH_SIZE, push

logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s")
log = logging.getLogger(__name__)


def main() -> None:
    parser = argparse.ArgumentParser(description="Collect and push Redshift lineage to Monte Carlo")
    parser.add_argument("--host", default=os.getenv("REDSHIFT_HOST"))         # ← SUBSTITUTE
    parser.add_argument("--db", default=os.getenv("REDSHIFT_DB"))             # ← SUBSTITUTE
    parser.add_argument("--user", default=os.getenv("REDSHIFT_USER"))         # ← SUBSTITUTE
    parser.add_argument("--password", default=os.getenv("REDSHIFT_PASSWORD")) # ← SUBSTITUTE
    parser.add_argument("--port", type=int, default=int(os.getenv("REDSHIFT_PORT", "5439")))
    parser.add_argument("--resource-uuid", default=os.getenv("MCD_RESOURCE_UUID"))
    parser.add_argument("--key-id", default=os.getenv("MCD_INGEST_ID"))
    parser.add_argument("--key-token", default=os.getenv("MCD_INGEST_TOKEN"))
    parser.add_argument("--lookback-hours", type=int, default=LOOKBACK_HOURS)
    parser.add_argument("--batch-size", type=int, default=DEFAULT_BATCH_SIZE)
    parser.add_argument("--manifest", default="manifest_lineage.json")
    args = parser.parse_args()

    required = ["host", "db", "user", "password", "resource_uuid", "key_id", "key_token"]
    missing = [k for k in required if getattr(args, k) is None]
    if missing:
        parser.error(f"Missing required arguments/env vars: {missing}")

    log.info("Step 1: Collecting lineage …")
    collect(
        host=args.host,
        db=args.db,
        user=args.user,
        password=args.password,
        manifest_path=args.manifest,
        port=args.port,
        lookback_hours=args.lookback_hours,
    )

    log.info("Step 2: Pushing lineage to Monte Carlo …")
    push(
        manifest_path=args.manifest,
        resource_uuid=args.resource_uuid,
        key_id=args.key_id,
        key_token=args.key_token,
        batch_size=args.batch_size,
    )

    log.info("Done — collect and push complete.")


if __name__ == "__main__":
    main()

"""
Redshift — Metadata Collect & Push (combined)
===============================================
Collects table schemas, row counts, and byte sizes from Amazon Redshift,
then pushes them to Monte Carlo via the push ingestion API.

This script imports and calls collect() from collect_metadata and push() from
push_metadata, running both in sequence.

Substitution points (search for "← SUBSTITUTE"):
  - REDSHIFT_HOST     : Redshift cluster endpoint or serverless workgroup endpoint
  - REDSHIFT_DB       : database name to connect to
  - REDSHIFT_USER     : database user (or IAM role user)
  - REDSHIFT_PASSWORD : database password
  - DB_EXCLUSIONS     : databases to skip
  - SCHEMA_EXCLUSIONS : schemas to skip in every database
  - MCD_INGEST_ID / MCD_INGEST_TOKEN : Monte Carlo API credentials
  - MCD_RESOURCE_UUID  : UUID of the Redshift connection in Monte Carlo
  - PUSH_BATCH_SIZE   : number of assets per API call (default 500)

Prerequisites:
  pip install psycopg2-binary pycarlo
"""

from __future__ import annotations

import argparse
import logging
import os

from collect_metadata import collect
from push_metadata import DEFAULT_BATCH_SIZE, push

logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s")
log = logging.getLogger(__name__)


def main() -> None:
    parser = argparse.ArgumentParser(description="Collect and push Redshift metadata to Monte Carlo")
    parser.add_argument("--host", default=os.getenv("REDSHIFT_HOST"))         # ← SUBSTITUTE
    parser.add_argument("--db", default=os.getenv("REDSHIFT_DB"))             # ← SUBSTITUTE
    parser.add_argument("--user", default=os.getenv("REDSHIFT_USER"))         # ← SUBSTITUTE
    parser.add_argument("--password", default=os.getenv("REDSHIFT_PASSWORD")) # ← SUBSTITUTE
    parser.add_argument("--port", type=int, default=int(os.getenv("REDSHIFT_PORT", "5439")))
    parser.add_argument("--resource-uuid", default=os.getenv("MCD_RESOURCE_UUID"))
    parser.add_argument("--key-id", default=os.getenv("MCD_INGEST_ID"))
    parser.add_argument("--key-token", default=os.getenv("MCD_INGEST_TOKEN"))
    parser.add_argument("--batch-size", type=int, default=DEFAULT_BATCH_SIZE)
    parser.add_argument("--manifest", default="manifest_metadata.json")
    args = parser.parse_args()

    required = ["host", "db", "user", "password", "resource_uuid", "key_id", "key_token"]
    missing = [k for k in required if getattr(args, k) is None]
    if missing:
        parser.error(f"Missing required arguments/env vars: {missing}")

    log.info("Step 1: Collecting metadata …")
    collect(
        host=args.host,
        db=args.db,
        user=args.user,
        password=args.password,
        manifest_path=args.manifest,
        port=args.port,
    )

    log.info("Step 2: Pushing metadata to Monte Carlo …")
    push(
        manifest_path=args.manifest,
        resource_uuid=args.resource_uuid,
        key_id=args.key_id,
        key_token=args.key_token,
        batch_size=args.batch_size,
    )

    log.info("Done — collect and push complete.")


if __name__ == "__main__":
    main()

"""
Redshift — Query Log Collect & Push (combined)
================================================
Collects completed query execution records from Redshift using sys_query_history
and sys_querytext, then pushes them to Monte Carlo for query-pattern analysis,
lineage derivation, and usage attribution.

This script imports and calls collect() from collect_query_logs and push() from
push_query_logs, running both in sequence.

Substitution points (search for "← SUBSTITUTE"):
  - REDSHIFT_HOST / REDSHIFT_DB / REDSHIFT_USER / REDSHIFT_PASSWORD : connection
  - LOOKBACK_HOURS    : hours back from [now - LAG_HOURS] to collect (default 25)
  - LOOKBACK_LAG_HOURS: lag behind now to avoid in-flight queries (default 1)
  - BATCH_SIZE        : number of query_ids to fetch texts for in one SQL call
  - MAX_QUERIES       : maximum query rows to process per run
  - MCD_INGEST_ID / MCD_INGEST_TOKEN : Monte Carlo API credentials
  - MCD_RESOURCE_UUID  : UUID of the Redshift connection in Monte Carlo
  - PUSH_BATCH_SIZE   : number of entries per API call (default 250)

Prerequisites:
  pip install psycopg2-binary pycarlo
"""

from __future__ import annotations

import argparse
import logging
import os

from collect_query_logs import BATCH_SIZE, LOOKBACK_HOURS, LOOKBACK_LAG_HOURS, MAX_QUERIES, collect
from push_query_logs import DEFAULT_BATCH_SIZE, push

logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s")
log = logging.getLogger(__name__)


def main() -> None:
    parser = argparse.ArgumentParser(description="Collect and push Redshift query logs to Monte Carlo")
    parser.add_argument("--host", default=os.getenv("REDSHIFT_HOST"))         # ← SUBSTITUTE
    parser.add_argument("--db", default=os.getenv("REDSHIFT_DB"))             # ← SUBSTITUTE
    parser.add_argument("--user", default=os.getenv("REDSHIFT_USER"))         # ← SUBSTITUTE
    parser.add_argument("--password", default=os.getenv("REDSHIFT_PASSWORD")) # ← SUBSTITUTE
    parser.add_argument("--port", type=int, default=int(os.getenv("REDSHIFT_PORT", "5439")))
    parser.add_argument("--resource-uuid", default=os.getenv("MCD_RESOURCE_UUID"))
    parser.add_argument("--key-id", default=os.getenv("MCD_INGEST_ID"))
    parser.add_argument("--key-token", default=os.getenv("MCD_INGEST_TOKEN"))
    parser.add_argument("--lookback-hours", type=int, default=LOOKBACK_HOURS)
    parser.add_argument("--lookback-lag-hours", type=int, default=LOOKBACK_LAG_HOURS)
    parser.add_argument("--batch-size", type=int, default=BATCH_SIZE)
    parser.add_argument("--max-queries", type=int, default=MAX_QUERIES)
    parser.add_argument("--push-batch-size", type=int, default=DEFAULT_BATCH_SIZE)
    parser.add_argument("--manifest", default="manifest_query_logs.json")
    args = parser.parse_args()

    required = ["host", "db", "user", "password", "resource_uuid", "key_id", "key_token"]
    missing = [k for k in required if getattr(args, k) is None]
    if missing:
        parser.error(f"Missing required arguments/env vars: {missing}")

    log.info("Step 1: Collecting query logs …")
    collect(
        host=args.host,
        db=args.db,
        user=args.user,
        password=args.password,
        manifest_path=args.manifest,
        port=args.port,
        lookback_hours=args.lookback_hours,
        lookback_lag_hours=args.lookback_lag_hours,
        batch_size=args.batch_size,
        max_queries=args.max_queries,
    )

    log.info("Step 2: Pushing query logs to Monte Carlo …")
    push(
        manifest_path=args.manifest,
        resource_uuid=args.resource_uuid,
        key_id=args.key_id,
        key_token=args.key_token,
        batch_size=args.push_batch_size,
    )

    log.info("Done — collect and push complete.")


if __name__ == "__main__":
    main()

"""
Redshift — Lineage Collection (collect-only)
==============================================
Collects table-level lineage from Redshift by fetching recent successful query
history from sys_query_history + sys_querytext and parsing CREATE TABLE AS SELECT
(CTAS) and INSERT INTO SELECT patterns to derive source->destination relationships.

Writes a JSON manifest file that can be consumed by push_lineage.py.

Substitution points (search for "← SUBSTITUTE"):
  - REDSHIFT_HOST / REDSHIFT_DB / REDSHIFT_USER / REDSHIFT_PASSWORD : connection
  - LOOKBACK_HOURS    : how far back to scan query history (default 24 h)

Prerequisites:
  pip install psycopg2-binary
"""

from __future__ import annotations

import argparse
import json
import logging
import os
import re
from datetime import datetime, timezone
from typing import Any

import psycopg2

logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s")
log = logging.getLogger(__name__)

RESOURCE_TYPE = "redshift"
LOOKBACK_HOURS: int = int(os.getenv("LOOKBACK_HOURS", "24"))  # ← SUBSTITUTE


def _check_available_memory(min_gb: float = 2.0) -> None:
    """Warn if available memory is below the threshold."""
    try:
        if hasattr(os, "sysconf"):  # Linux / macOS
            page_size = os.sysconf("SC_PAGE_SIZE")
            avail_pages = os.sysconf("SC_AVPHYS_PAGES")
            avail_gb = (page_size * avail_pages) / (1024 ** 3)
        else:
            return  # Windows — skip check
    except (ValueError, OSError):
        return
    if avail_gb < min_gb:
        log.warning(
            "Only %.1f GB of memory available (minimum recommended: %.1f GB). "
            "Consider reducing the collection scope or increasing available memory.",
            avail_gb,
            min_gb,
        )


# Regex: CTAS — CREATE [OR REPLACE] TABLE <dest> AS SELECT
_CTAS_RE = re.compile(
    r"CREATE\s+(?:OR\s+REPLACE\s+)?(?:TABLE|VIEW)\s+(?P<dest>\"?[\w.\"]+\"?)\s*(?:\([^)]*\))?\s*AS\s+SELECT\b",
    re.IGNORECASE | re.DOTALL,
)
# Regex: INSERT INTO <dest> … SELECT
_INSERT_RE = re.compile(
    r"INSERT\s+INTO\s+(?P<dest>\"?[\w.\"]+\"?)\s.*?SELECT\b",
    re.IGNORECASE | re.DOTALL,
)
# Matches any schema.table or database.schema.table reference in the query
_TABLE_REF_RE = re.compile(r'"?([\w]+)"?\."?([\w]+)"?(?:\."?([\w]+)"?)?', re.IGNORECASE)


def _clean_name(name: str) -> str:
    return name.strip('"').strip()


def _parse_ref(ref: str) -> tuple[str, str, str]:
    """Parse 'db.schema.table' or 'schema.table' -> (database, schema, table)."""
    parts = [_clean_name(p) for p in ref.split(".")]
    if len(parts) == 3:
        return parts[0], parts[1], parts[2]
    if len(parts) == 2:
        return "", parts[0], parts[1]
    return "", "", parts[0]


def _dictfetch(cursor: Any, sql: str, params: tuple | None = None) -> list[dict[str, Any]]:
    cursor.execute(sql, params)
    cols = [d.name for d in cursor.description]
    rows = []
    while True:
        chunk = cursor.fetchmany(1000)
        if not chunk:
            break
        rows.extend(dict(zip(cols, row)) for row in chunk)
    return rows


def fetch_query_texts(cursor: Any, lookback_hours: int) -> list[str]:
    """Assemble full query texts from sys_query_history + sys_querytext."""
    rows = _dictfetch(
        cursor,
        f"""
        SELECT
            sq.query_id,
            LISTAGG(
                CASE WHEN LEN(st.text) <= 200 THEN st.text ELSE LEFT(st.text, 200) END,
                ''
            ) WITHIN GROUP (ORDER BY st.sequence) AS full_text
        FROM sys_query_history sq
        JOIN sys_querytext st ON sq.query_id = st.query_id
        WHERE sq.start_time >= DATEADD(hour, -{lookback_hours}, GETDATE())
          AND sq.status = 'success'
        GROUP BY sq.query_id
        LIMIT 50000
        """,  # ← SUBSTITUTE: adjust lookback_hours, LIMIT, or add user/database filters
    )
    return [r["full_text"] for r in rows if r.get("full_text")]


def parse_lineage_from_sql(sql_text: str) -> list[dict[str, Any]]:
    events: list[dict[str, Any]] = []

    dest_match = _CTAS_RE.search(sql_text) or _INSERT_RE.search(sql_text)
    if not dest_match:
        return events

    dest_raw = dest_match.group("dest")
    dest_db, dest_schema, dest_table = _parse_ref(dest_raw)
    if not dest_table:
        return events

    # Find all schema.table refs in the query, excluding the destination
    source_refs: list[str] = []
    for m in _TABLE_REF_RE.finditer(sql_text):
        if m.group(3):
            ref = f"{m.group(1)}.{m.group(2)}.{m.group(3)}"
        else:
            ref = f"{m.group(1)}.{m.group(2)}"

        db, schema, table = _parse_ref(ref)
        if not table or (db == dest_db and schema == dest_schema and table == dest_table):
            continue
        source_refs.append(ref)

    if not source_refs:
        return events

    # Deduplicate sources while preserving order
    seen: set[str] = set()
    sources: list[dict[str, str]] = []
    for ref in source_refs:
        if ref not in seen:
            seen.add(ref)
            db, schema, table = _parse_ref(ref)
            sources.append({"database": db, "schema": schema, "asset_name": table})

    events.append({
        "sources": sources,
        "destination": {"database": dest_db, "schema": dest_schema, "asset_name": dest_table},
    })
    return events


def collect(
    host: str,
    db: str,
    user: str,
    password: str,
    manifest_path: str = "manifest_lineage.json",
    port: int = 5439,
    lookback_hours: int = LOOKBACK_HOURS,
) -> list[dict[str, Any]]:
    """Connect to Redshift, collect lineage, write a JSON manifest, and return events."""
    _check_available_memory()
    collected_at = datetime.now(timezone.utc).isoformat()

    conn = psycopg2.connect(
        host=host, port=port, dbname=db, user=user, password=password, connect_timeout=30,
    )
    try:
        with conn.cursor() as cursor:
            query_texts = fetch_query_texts(cursor, lookback_hours)
    finally:
        conn.close()

    log.info("Parsing lineage from %d query texts …", len(query_texts))
    all_events: list[dict[str, Any]] = []
    for sql_text in query_texts:
        all_events.extend(parse_lineage_from_sql(sql_text))

    log.info("Collected %d lineage events", len(all_events))

    manifest = {
        "resource_type": RESOURCE_TYPE,
        "collected_at": collected_at,
        "lookback_hours": lookback_hours,
        "queries_scanned": len(query_texts),
        "lineage_event_count": len(all_events),
        "events": all_events,
    }
    with open(manifest_path, "w") as fh:
        json.dump(manifest, fh, indent=2)
    log.info("Manifest written to %s (%d events)", manifest_path, len(all_events))

    return all_events


def main() -> None:
    parser = argparse.ArgumentParser(description="Collect Redshift lineage to a manifest file")
    parser.add_argument("--host", default=os.getenv("REDSHIFT_HOST"))         # ← SUBSTITUTE
    parser.add_argument("--db", default=os.getenv("REDSHIFT_DB"))             # ← SUBSTITUTE
    parser.add_argument("--user", default=os.getenv("REDSHIFT_USER"))         # ← SUBSTITUTE
    parser.add_argument("--password", default=os.getenv("REDSHIFT_PASSWORD")) # ← SUBSTITUTE
    parser.add_argument("--port", type=int, default=int(os.getenv("REDSHIFT_PORT", "5439")))
    parser.add_argument("--lookback-hours", type=int, default=LOOKBACK_HOURS)
    parser.add_argument("--manifest", default="manifest_lineage.json")
    args = parser.parse_args()

    required = ["host", "db", "user", "password"]
    missing = [k for k in required if getattr(args, k) is None]
    if missing:
        parser.error(f"Missing required arguments/env vars: {missing}")

    collect(
        host=args.host,
        db=args.db,
        user=args.user,
        password=args.password,
        manifest_path=args.manifest,
        port=args.port,
        lookback_hours=args.lookback_hours,
    )


if __name__ == "__main__":
    main()

"""
Redshift — Metadata Collection (collect-only)
===============================================
Collects table schemas, row counts, and byte sizes from Amazon Redshift using
SVV system views, then writes a JSON manifest file that can be consumed by
push_metadata.py.

Substitution points (search for "← SUBSTITUTE"):
  - REDSHIFT_HOST     : Redshift cluster endpoint or serverless workgroup endpoint
  - REDSHIFT_DB       : database name to connect to
  - REDSHIFT_USER     : database user (or IAM role user)
  - REDSHIFT_PASSWORD : database password
  - DB_EXCLUSIONS     : databases to skip
  - SCHEMA_EXCLUSIONS : schemas to skip in every database

Prerequisites:
  pip install psycopg2-binary
"""

from __future__ import annotations

import argparse
import json
import logging
import os
from datetime import datetime, timezone
from typing import Any

import psycopg2
import psycopg2.extras

logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s")
log = logging.getLogger(__name__)

RESOURCE_TYPE = "redshift"

DB_EXCLUSIONS: set[str] = {"dev", "padb_harvest"}  # ← SUBSTITUTE: add internal databases

SCHEMA_EXCLUSIONS: set[str] = {  # ← SUBSTITUTE: add internal schemas
    "information_schema",
    "pg_catalog",
    "pg_internal",
    "catalog_history",
}


def _check_available_memory(min_gb: float = 2.0) -> None:
    """Warn if available memory is below the threshold."""
    try:
        if hasattr(os, "sysconf"):  # Linux / macOS
            page_size = os.sysconf("SC_PAGE_SIZE")
            avail_pages = os.sysconf("SC_AVPHYS_PAGES")
            avail_gb = (page_size * avail_pages) / (1024 ** 3)
        else:
            return  # Windows — skip check
    except (ValueError, OSError):
        return
    if avail_gb < min_gb:
        log.warning(
            "Only %.1f GB of memory available (minimum recommended: %.1f GB). "
            "Consider reducing the collection scope or increasing available memory.",
            avail_gb,
            min_gb,
        )


def _dictfetch(cursor: Any, sql: str, params: tuple | None = None) -> list[dict[str, Any]]:
    cursor.execute(sql, params)
    cols = [d.name for d in cursor.description]
    rows = []
    while True:
        chunk = cursor.fetchmany(1000)
        if not chunk:
            break
        rows.extend(dict(zip(cols, row)) for row in chunk)
    return rows


def collect_databases(cursor: Any) -> list[str]:
    rows = _dictfetch(
        cursor,
        "SELECT database_name FROM svv_redshift_databases ORDER BY database_name",
    )
    return [r["database_name"] for r in rows if r["database_name"] not in DB_EXCLUSIONS]


def collect_tables(cursor: Any, db: str) -> list[dict[str, Any]]:
    schema_list = ", ".join(f"'{s}'" for s in SCHEMA_EXCLUSIONS)
    return _dictfetch(
        cursor,
        f"""
        SELECT
            database      AS db,
            schema,
            "table"       AS table_name,
            "rows"        AS row_count,
            size * 1024 * 1024 AS byte_count
        FROM svv_table_info
        WHERE database = %s
          AND schema NOT IN ({schema_list})
        ORDER BY schema, "table"
        """,  # ← SUBSTITUTE: add additional WHERE clauses to narrow scope
        (db,),
    )


def collect_columns(cursor: Any, db: str, schema: str, table: str) -> list[dict[str, Any]]:
    return _dictfetch(
        cursor,
        """
        SELECT column_name, data_type, remarks AS comment
        FROM svv_columns
        WHERE table_catalog = %s
          AND table_schema  = %s
          AND table_name    = %s
        ORDER BY ordinal_position
        """,
        (db, schema, table),
    )


def collect(
    host: str,
    db: str,
    user: str,
    password: str,
    manifest_path: str = "manifest_metadata.json",
    port: int = 5439,
) -> list[dict[str, Any]]:
    """Connect to Redshift, collect metadata, write a JSON manifest, and return asset dicts."""
    _check_available_memory()
    collected_at = datetime.now(timezone.utc).isoformat()
    assets: list[dict[str, Any]] = []

    conn = psycopg2.connect(
        host=host,          # ← SUBSTITUTE
        port=port,
        dbname=db,          # ← SUBSTITUTE
        user=user,          # ← SUBSTITUTE
        password=password,  # ← SUBSTITUTE
        connect_timeout=30,
    )
    try:
        with conn.cursor() as cursor:
            databases = collect_databases(cursor)
            log.info("Found databases: %s", databases)

            for database in databases:
                tables = collect_tables(cursor, database)
                log.info("Database %s — %d tables", database, len(tables))

                for t in tables:
                    schema = t["schema"]
                    table_name = t["table_name"]

                    columns = collect_columns(cursor, database, schema, table_name)
                    fields = [
                        {
                            "name": col["column_name"],
                            "type": col["data_type"].upper(),
                            "description": col.get("comment") or None,
                        }
                        for col in columns
                    ]

                    asset = {
                        "asset_name": table_name,
                        "database": database,   # ← SUBSTITUTE: use database as top-level namespace
                        "schema": schema,
                        "asset_type": "TABLE",
                        "fields": fields,
                        "row_count": t.get("row_count"),
                        "byte_count": t.get("byte_count"),
                    }
                    assets.append(asset)
                    log.info("Collected %s.%s.%s", database, schema, table_name)
    finally:
        conn.close()

    manifest = {
        "resource_type": RESOURCE_TYPE,
        "collected_at": collected_at,
        "asset_count": len(assets),
        "assets": assets,
    }
    with open(manifest_path, "w") as fh:
        json.dump(manifest, fh, indent=2)
    log.info("Manifest written to %s (%d assets)", manifest_path, len(assets))

    return assets


def main() -> None:
    parser = argparse.ArgumentParser(description="Collect Redshift metadata to a manifest file")
    parser.add_argument("--host", default=os.getenv("REDSHIFT_HOST"))         # ← SUBSTITUTE
    parser.add_argument("--db", default=os.getenv("REDSHIFT_DB"))             # ← SUBSTITUTE
    parser.add_argument("--user", default=os.getenv("REDSHIFT_USER"))         # ← SUBSTITUTE
    parser.add_argument("--password", default=os.getenv("REDSHIFT_PASSWORD")) # ← SUBSTITUTE
    parser.add_argument("--port", type=int, default=int(os.getenv("REDSHIFT_PORT", "5439")))
    parser.add_argument("--manifest", default="manifest_metadata.json")
    args = parser.parse_args()

    required = ["host", "db", "user", "password"]
    missing = [k for k in required if getattr(args, k) is None]
    if missing:
        parser.error(f"Missing required arguments/env vars: {missing}")

    collect(
        host=args.host,
        db=args.db,
        user=args.user,
        password=args.password,
        manifest_path=args.manifest,
        port=args.port,
    )


if __name__ == "__main__":
    main()

"""
Redshift — Query Log Collection (collect-only)
================================================
Collects completed query execution records from Redshift using sys_query_history
and sys_querytext (modern RA3/serverless), assembles full SQL text from
multi-row text chunks, and writes a JSON manifest file that can be consumed
by push_query_logs.py.

Substitution points (search for "← SUBSTITUTE"):
  - REDSHIFT_HOST / REDSHIFT_DB / REDSHIFT_USER / REDSHIFT_PASSWORD : connection
  - LOOKBACK_HOURS    : hours back from [now - LAG_HOURS] to collect (default 25)
  - LOOKBACK_LAG_HOURS: lag behind now to avoid in-flight queries (default 1)
  - BATCH_SIZE        : number of query_ids to fetch texts for in one SQL call
  - MAX_QUERIES       : maximum query rows to process per run

Prerequisites:
  pip install psycopg2-binary
"""

from __future__ import annotations

import argparse
import json
import logging
import os
from datetime import datetime, timezone
from typing import Any

import psycopg2

logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s")
log = logging.getLogger(__name__)

LOG_TYPE = "redshift"

LOOKBACK_HOURS: int = int(os.getenv("LOOKBACK_HOURS", "25"))        # ← SUBSTITUTE
LOOKBACK_LAG_HOURS: int = int(os.getenv("LOOKBACK_LAG_HOURS", "1")) # ← SUBSTITUTE
BATCH_SIZE: int = int(os.getenv("BATCH_SIZE", "200"))               # ← SUBSTITUTE
MAX_QUERIES: int = int(os.getenv("MAX_QUERIES", "10000"))           # ← SUBSTITUTE


def _check_available_memory(min_gb: float = 2.0) -> None:
    """Warn if available memory is below the threshold."""
    try:
        if hasattr(os, "sysconf"):  # Linux / macOS
            page_size = os.sysconf("SC_PAGE_SIZE")
            avail_pages = os.sysconf("SC_AVPHYS_PAGES")
            avail_gb = (page_size * avail_pages) / (1024 ** 3)
        else:
            return  # Windows — skip check
    except (ValueError, OSError):
        return
    if avail_gb < min_gb:
        log.warning(
            "Only %.1f GB of memory available (minimum recommended: %.1f GB). "
            "Consider reducing the collection scope or increasing available memory.",
            avail_gb,
            min_gb,
        )


def _dictfetch(cursor: Any, sql: str, params: tuple | None = None) -> list[dict[str, Any]]:
    cursor.execute(sql, params)
    cols = [d.name for d in cursor.description]
    rows = []
    while True:
        chunk = cursor.fetchmany(1000)
        if not chunk:
            break
        rows.extend(dict(zip(cols, row)) for row in chunk)
    return rows


def _safe_isoformat(dt: Any) -> str | None:
    if dt is None:
        return None
    if hasattr(dt, "isoformat"):
        if dt.tzinfo is None:
            dt = dt.replace(tzinfo=timezone.utc)
        return dt.isoformat()
    return str(dt)


def fetch_query_metadata(
    cursor: Any,
    lookback_hours: int,
    lag_hours: int,
    max_queries: int,
) -> list[dict[str, Any]]:
    """Fetch query execution metadata from sys_query_history."""
    return _dictfetch(
        cursor,
        f"""
        SELECT
            query_id,
            start_time,
            end_time,
            status,
            user_id,
            database_name,
            elapsed_time
        FROM sys_query_history
        WHERE start_time >= DATEADD(hour, -{lookback_hours}, GETDATE())
          AND start_time <  DATEADD(hour, -{lag_hours},      GETDATE())
          AND status = 'success'
        ORDER BY start_time
        LIMIT {max_queries}
        """,  # ← SUBSTITUTE: add AND database_name = 'mydb' to narrow scope
    )


def fetch_query_texts_batch(cursor: Any, query_ids: list[int]) -> dict[int, str]:
    """Batch-fetch and assemble multi-row query texts for a list of query_ids."""
    if not query_ids:
        return {}

    # Build a VALUES list for the IN clause to avoid large parameter arrays
    id_list = ", ".join(str(qid) for qid in query_ids)
    rows = _dictfetch(
        cursor,
        f"""
        SELECT
            query_id,
            LISTAGG(
                CASE WHEN LEN(text) <= 200 THEN text ELSE LEFT(text, 200) END,
                ''
            ) WITHIN GROUP (ORDER BY sequence) AS query_text
        FROM sys_querytext
        WHERE query_id IN ({id_list})
        GROUP BY query_id
        """,
    )
    return {r["query_id"]: r["query_text"] for r in rows if r.get("query_text")}


def collect(
    host: str,
    db: str,
    user: str,
    password: str,
    manifest_path: str = "manifest_query_logs.json",
    port: int = 5439,
    lookback_hours: int = LOOKBACK_HOURS,
    lookback_lag_hours: int = LOOKBACK_LAG_HOURS,
    batch_size: int = BATCH_SIZE,
    max_queries: int = MAX_QUERIES,
) -> list[dict[str, Any]]:
    """Connect to Redshift, collect query logs, write a JSON manifest, and return entries."""
    _check_available_memory()
    collected_at = datetime.now(timezone.utc).isoformat()

    conn = psycopg2.connect(
        host=host, port=port, dbname=db, user=user, password=password, connect_timeout=30,
    )
    try:
        with conn.cursor() as cursor:
            query_meta = fetch_query_metadata(cursor, lookback_hours, lookback_lag_hours, max_queries)
            log.info("Retrieved %d query metadata rows", len(query_meta))

            # Batch-fetch texts to avoid enormous single queries
            query_ids = [r["query_id"] for r in query_meta]
            text_map: dict[int, str] = {}
            for i in range(0, len(query_ids), batch_size):
                batch = query_ids[i : i + batch_size]
                text_map.update(fetch_query_texts_batch(cursor, batch))
                log.debug("Fetched texts for batch %d–%d", i, i + len(batch))
    finally:
        conn.close()

    entries: list[dict[str, Any]] = []
    for row in query_meta:
        qid = row["query_id"]
        query_text = text_map.get(qid, "")
        if not query_text.strip():
            continue  # ← SUBSTITUTE: decide whether to push rows with missing text

        entry = {
            "query_id": str(qid),
            "query_text": query_text,
            "start_time": _safe_isoformat(row.get("start_time")),
            "end_time": _safe_isoformat(row.get("end_time")),
            "user": str(row.get("user_id")) if row.get("user_id") is not None else None,
            "database_name": row.get("database_name"),
            "elapsed_time_us": row.get("elapsed_time"),
        }
        entries.append(entry)

    log.info("Collected %d query log entries", len(entries))

    manifest = {
        "log_type": LOG_TYPE,
        "collected_at": collected_at,
        "lookback_hours": lookback_hours,
        "lookback_lag_hours": lookback_lag_hours,
        "query_log_count": len(entries),
        "entries": entries,
    }
    with open(manifest_path, "w") as fh:
        json.dump(manifest, fh, indent=2)
    log.info("Manifest written to %s (%d entries)", manifest_path, len(entries))

    return entries


def main() -> None:
    parser = argparse.ArgumentParser(description="Collect Redshift query logs to a manifest file")
    parser.add_argument("--host", default=os.getenv("REDSHIFT_HOST"))         # ← SUBSTITUTE
    parser.add_argument("--db", default=os.getenv("REDSHIFT_DB"))             # ← SUBSTITUTE
    parser.add_argument("--user", default=os.getenv("REDSHIFT_USER"))         # ← SUBSTITUTE
    parser.add_argument("--password", default=os.getenv("REDSHIFT_PASSWORD")) # ← SUBSTITUTE
    parser.add_argument("--port", type=int, default=int(os.getenv("REDSHIFT_PORT", "5439")))
    parser.add_argument("--lookback-hours", type=int, default=LOOKBACK_HOURS)
    parser.add_argument("--lookback-lag-hours", type=int, default=LOOKBACK_LAG_HOURS)
    parser.add_argument("--batch-size", type=int, default=BATCH_SIZE)
    parser.add_argument("--max-queries", type=int, default=MAX_QUERIES)
    parser.add_argument("--manifest", default="manifest_query_logs.json")
    args = parser.parse_args()

    required = ["host", "db", "user", "password"]
    missing = [k for k in required if getattr(args, k) is None]
    if missing:
        parser.error(f"Missing required arguments/env vars: {missing}")

    collect(
        host=args.host,
        db=args.db,
        user=args.user,
        password=args.password,
        manifest_path=args.manifest,
        port=args.port,
        lookback_hours=args.lookback_hours,
        lookback_lag_hours=args.lookback_lag_hours,
        batch_size=args.batch_size,
        max_queries=args.max_queries,
    )


if __name__ == "__main__":
    main()

"""
Redshift — Lineage Push (push-only)
=====================================
Reads a JSON manifest file produced by collect_lineage.py and pushes the lineage
events to Monte Carlo via the push ingestion API, with configurable batching to
keep compressed payloads under 1 MB.

Substitution points (search for "← SUBSTITUTE"):
  - MCD_INGEST_ID / MCD_INGEST_TOKEN : Monte Carlo API credentials
  - MCD_RESOURCE_UUID      : UUID of the Redshift connection in Monte Carlo
  - PUSH_BATCH_SIZE       : number of events per API call (default 500)

Prerequisites:
  pip install pycarlo
"""

from __future__ import annotations

import argparse
import json
import logging
import os
from concurrent.futures import ThreadPoolExecutor, as_completed
from datetime import datetime, timezone
from typing import Any

from pycarlo.core import Client, Session
from pycarlo.features.ingestion import IngestionService
from pycarlo.features.ingestion.models import (
    LineageAssetRef,
    LineageEvent,
)

logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s")
log = logging.getLogger(__name__)

RESOURCE_TYPE = "redshift"
DEFAULT_BATCH_SIZE = 500  # ← SUBSTITUTE: conservative default to stay under 1 MB compressed


def _ref_from_dict(d: dict[str, Any]) -> LineageAssetRef:
    return LineageAssetRef(
        type="TABLE",
        name=d["asset_name"],
        database=d.get("database", ""),
        schema=d.get("schema", ""),
    )


def _event_from_dict(d: dict[str, Any]) -> LineageEvent:
    """Reconstruct a LineageEvent from a manifest dict."""
    sources = [_ref_from_dict(s) for s in d.get("sources", [])]
    destination = _ref_from_dict(d["destination"])
    return LineageEvent(
        sources=sources,
        destination=destination,
    )


def push(
    manifest_path: str,
    resource_uuid: str,
    key_id: str,
    key_token: str,
    batch_size: int = DEFAULT_BATCH_SIZE,
) -> dict[str, Any]:
    """Read a collect manifest and push lineage events to Monte Carlo in batches.

    Returns a summary dict with invocation IDs and counts.
    """
    with open(manifest_path) as fh:
        manifest = json.load(fh)

    event_dicts: list[dict[str, Any]] = manifest["events"]
    events = [_event_from_dict(d) for d in event_dicts]
    log.info("Loaded %d lineage events from %s", len(events), manifest_path)

    if not events:
        log.info("No lineage events to push.")
        summary = {
            "resource_uuid": resource_uuid,
            "resource_type": RESOURCE_TYPE,
            "invocation_ids": [],
            "pushed_at": datetime.now(timezone.utc).isoformat(),
            "event_count": 0,
            "batch_count": 0,
            "batch_size": batch_size,
        }
        push_manifest_path = manifest_path.replace(".json", "_push_result.json")
        with open(push_manifest_path, "w") as fh:
            json.dump(summary, fh, indent=2)
        return summary

    # Split into batches
    batches = []
    for i in range(0, len(events), batch_size):
        batches.append(events[i : i + batch_size])
    total_batches = len(batches)

    def _push_batch(batch: list, batch_num: int) -> str | None:
        """Push a single batch using a dedicated Session (thread-safe)."""
        log.info("Pushing batch %d/%d (%d events) ...", batch_num, total_batches, len(batch))
        client = Client(session=Session(mcd_id=key_id, mcd_token=key_token, scope="Ingestion"))
        service = IngestionService(mc_client=client)
        result = service.send_lineage(
            resource_uuid=resource_uuid,
            resource_type=RESOURCE_TYPE,
            events=batch,
        )
        invocation_id = service.extract_invocation_id(result)
        if invocation_id:
            log.info("Batch %d: invocation_id=%s", batch_num, invocation_id)
        return invocation_id

    # Push batches in parallel (each thread gets its own pycarlo Session)
    max_workers = min(4, total_batches)
    invocation_ids: list[str | None] = [None] * total_batches

    with ThreadPoolExecutor(max_workers=max_workers) as pool:
        futures = {
            pool.submit(_push_batch, batch, i + 1): i
            for i, batch in enumerate(batches)
        }
        for future in as_completed(futures):
            idx = futures[future]
            try:
                invocation_ids[idx] = future.result()
            except Exception as exc:
                log.error("ERROR pushing batch %d: %s", idx + 1, exc)
                raise

    log.info("All %d batches pushed (%d workers)", total_batches, max_workers)

    summary = {
        "resource_uuid": resource_uuid,
        "resource_type": RESOURCE_TYPE,
        "invocation_ids": invocation_ids,
        "pushed_at": datetime.now(timezone.utc).isoformat(),
        "event_count": len(events),
        "batch_count": total_batches,
        "batch_size": batch_size,
        "lookback_hours": manifest.get("lookback_hours"),
        "queries_scanned": manifest.get("queries_scanned"),
    }

    push_manifest_path = manifest_path.replace(".json", "_push_result.json")
    with open(push_manifest_path, "w") as fh:
        json.dump(summary, fh, indent=2)
    log.info("Push result written to %s", push_manifest_path)

    return summary


def main() -> None:
    parser = argparse.ArgumentParser(description="Push Redshift lineage to Monte Carlo from manifest")
    parser.add_argument("--manifest", default="manifest_lineage.json")
    parser.add_argument("--resource-uuid", default=os.getenv("MCD_RESOURCE_UUID"))
    parser.add_argument("--key-id", default=os.getenv("MCD_INGEST_ID"))
    parser.add_argument("--key-token", default=os.getenv("MCD_INGEST_TOKEN"))
    parser.add_argument("--batch-size", type=int, default=DEFAULT_BATCH_SIZE)
    args = parser.parse_args()

    required = ["resource_uuid", "key_id", "key_token"]
    missing = [k for k in required if getattr(args, k) is None]
    if missing:
        parser.error(f"Missing required arguments/env vars: {missing}")

    push(
        manifest_path=args.manifest,
        resource_uuid=args.resource_uuid,
        key_id=args.key_id,
        key_token=args.key_token,
        batch_size=args.batch_size,
    )


if __name__ == "__main__":
    main()

"""
Redshift — Metadata Push (push-only)
======================================
Reads a JSON manifest file produced by collect_metadata.py and pushes the assets
to Monte Carlo via the push ingestion API, with configurable batching to keep
compressed payloads under 1 MB.

Substitution points (search for "← SUBSTITUTE"):
  - MCD_INGEST_ID / MCD_INGEST_TOKEN : Monte Carlo API credentials
  - MCD_RESOURCE_UUID      : UUID of the Redshift connection in Monte Carlo
  - PUSH_BATCH_SIZE       : number of assets per API call (default 500)

Prerequisites:
  pip install pycarlo
"""

from __future__ import annotations

import argparse
import json
import logging
import os
from concurrent.futures import ThreadPoolExecutor, as_completed
from datetime import datetime, timezone
from typing import Any

from pycarlo.core import Client, Session
from pycarlo.features.ingestion import IngestionService
from pycarlo.features.ingestion.models import (
    AssetField,
    AssetFreshness,
    AssetMetadata,
    AssetVolume,
    RelationalAsset,
)

logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s")
log = logging.getLogger(__name__)

RESOURCE_TYPE = "redshift"
DEFAULT_BATCH_SIZE = 500  # ← SUBSTITUTE: conservative default to stay under 1 MB compressed


def _asset_from_dict(d: dict[str, Any]) -> RelationalAsset:
    """Reconstruct a RelationalAsset from a manifest dict."""
    fields = [
        AssetField(
            name=f["name"],
            type=f.get("type"),
            description=f.get("description"),
        )
        for f in d.get("fields", [])
    ]

    volume = None
    if d.get("row_count") is not None or d.get("byte_count") is not None:
        volume = AssetVolume(
            row_count=d.get("row_count"),
            byte_count=d.get("byte_count"),
        )

    freshness = None
    if d.get("last_updated") is not None:
        freshness = AssetFreshness(last_update_time=d.get("last_updated"))

    return RelationalAsset(
        type=d.get("asset_type", "TABLE"),
        metadata=AssetMetadata(
            name=d["asset_name"],
            database=d["database"],   # ← SUBSTITUTE: use database as top-level namespace
            schema=d["schema"],
            description=d.get("description"),
        ),
        fields=fields,
        volume=volume,
        freshness=freshness,
    )


def push(
    manifest_path: str,
    resource_uuid: str,
    key_id: str,
    key_token: str,
    batch_size: int = DEFAULT_BATCH_SIZE,
) -> dict[str, Any]:
    """Read a collect manifest and push assets to Monte Carlo in batches.

    Returns a summary dict with invocation IDs and counts.
    """
    with open(manifest_path) as fh:
        manifest = json.load(fh)

    asset_dicts: list[dict[str, Any]] = manifest["assets"]
    assets = [_asset_from_dict(d) for d in asset_dicts]
    log.info("Loaded %d assets from %s", len(assets), manifest_path)

    # Split into batches
    batches = []
    for i in range(0, max(len(assets), 1), batch_size):
        batches.append(assets[i : i + batch_size])
    total_batches = len(batches)

    def _push_batch(batch: list, batch_num: int) -> str | None:
        """Push a single batch using a dedicated Session (thread-safe)."""
        client = Client(session=Session(mcd_id=key_id, mcd_token=key_token, scope="Ingestion"))
        service = IngestionService(mc_client=client)
        result = service.send_metadata(
            resource_uuid=resource_uuid,
            resource_type=RESOURCE_TYPE,
            events=batch,
        )
        invocation_id = service.extract_invocation_id(result)
        log.info("Pushed batch %d/%d (%d assets) — invocation_id=%s", batch_num, total_batches, len(batch), invocation_id)
        return invocation_id

    # Push batches in parallel (each thread gets its own pycarlo Session)
    max_workers = min(4, total_batches)
    invocation_ids: list[str | None] = [None] * total_batches

    with ThreadPoolExecutor(max_workers=max_workers) as pool:
        futures = {
            pool.submit(_push_batch, batch, i + 1): i
            for i, batch in enumerate(batches)
        }
        for future in as_completed(futures):
            idx = futures[future]
            try:
                invocation_ids[idx] = future.result()
            except Exception as exc:
                log.error("ERROR pushing batch %d: %s", idx + 1, exc)
                raise

    log.info("All %d batches pushed (%d workers)", total_batches, max_workers)

    summary = {
        "resource_uuid": resource_uuid,
        "resource_type": RESOURCE_TYPE,
        "invocation_ids": invocation_ids,
        "pushed_at": datetime.now(timezone.utc).isoformat(),
        "asset_count": len(assets),
        "batch_count": total_batches,
        "batch_size": batch_size,
    }

    push_manifest_path = manifest_path.replace(".json", "_push_result.json")
    with open(push_manifest_path, "w") as fh:
        json.dump(summary, fh, indent=2)
    log.info("Push result written to %s", push_manifest_path)

    return summary


def main() -> None:
    parser = argparse.ArgumentParser(description="Push Redshift metadata to Monte Carlo from manifest")
    parser.add_argument("--manifest", default="manifest_metadata.json")
    parser.add_argument("--resource-uuid", default=os.getenv("MCD_RESOURCE_UUID"))
    parser.add_argument("--key-id", default=os.getenv("MCD_INGEST_ID"))
    parser.add_argument("--key-token", default=os.getenv("MCD_INGEST_TOKEN"))
    parser.add_argument("--batch-size", type=int, default=DEFAULT_BATCH_SIZE)
    args = parser.parse_args()

    required = ["resource_uuid", "key_id", "key_token"]
    missing = [k for k in required if getattr(args, k) is None]
    if missing:
        parser.error(f"Missing required arguments/env vars: {missing}")

    push(
        manifest_path=args.manifest,
        resource_uuid=args.resource_uuid,
        key_id=args.key_id,
        key_token=args.key_token,
        batch_size=args.batch_size,
    )


if __name__ == "__main__":
    main()

"""
Redshift — Query Log Push (push-only)
=======================================
Reads a JSON manifest file produced by collect_query_logs.py and pushes the query
log entries to Monte Carlo via the push ingestion API, with configurable batching
to keep compressed payloads under 1 MB.

Substitution points (search for "← SUBSTITUTE"):
  - MCD_INGEST_ID / MCD_INGEST_TOKEN : Monte Carlo API credentials
  - MCD_RESOURCE_UUID      : UUID of the Redshift connection in Monte Carlo
  - PUSH_BATCH_SIZE       : number of entries per API call (default 100)

Prerequisites:
  pip install pycarlo
"""

from __future__ import annotations

import argparse
import json
import logging
import os
from concurrent.futures import ThreadPoolExecutor, as_completed
from datetime import datetime, timezone
from typing import Any

from dateutil.parser import isoparse
from pycarlo.core import Client, Session
from pycarlo.features.ingestion import IngestionService
from pycarlo.features.ingestion.models import QueryLogEntry

logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s")
log = logging.getLogger(__name__)

LOG_TYPE = "redshift"
DEFAULT_BATCH_SIZE = 100  # ← SUBSTITUTE: conservative default to stay under 1 MB compressed

# Truncate query_text longer than this to prevent 413 errors.
# Some SQL statements (e.g., generated by BI tools) can be 100KB+ and blow up
# compressed payloads even at small batch sizes.
_MAX_QUERY_TEXT_LEN = 10_000


def _build_query_log_entries(entry_dicts: list[dict[str, Any]]) -> list[QueryLogEntry]:
    """Convert manifest query dicts into QueryLogEntry objects."""
    entries = []
    truncated = 0
    for d in entry_dicts:
        query_text = d.get("query_text") or ""

        # Truncate very long SQL to prevent 413 Request Too Large
        if len(query_text) > _MAX_QUERY_TEXT_LEN:
            query_text = query_text[:_MAX_QUERY_TEXT_LEN] + "... [TRUNCATED]"
            truncated += 1

        extra = {}
        if d.get("database_name") is not None:
            extra["database_name"] = d["database_name"]
        if d.get("elapsed_time_us") is not None:
            extra["elapsed_time_us"] = d["elapsed_time_us"]

        start_time = d.get("start_time")
        end_time = d.get("end_time")

        entries.append(
            QueryLogEntry(
                query_id=d.get("query_id"),
                query_text=query_text,
                start_time=isoparse(start_time) if start_time else None,
                end_time=isoparse(end_time) if end_time else None,
                user=d.get("user"),
                extra=extra or None,
            )
        )
    if truncated:
        log.info("Truncated %d query text(s) exceeding %d chars", truncated, _MAX_QUERY_TEXT_LEN)
    return entries


def push(
    manifest_path: str,
    resource_uuid: str,
    key_id: str,
    key_token: str,
    batch_size: int = DEFAULT_BATCH_SIZE,
) -> dict[str, Any]:
    """Read a collect manifest and push query log entries to Monte Carlo in batches.

    Returns a summary dict with invocation IDs and counts.
    """
    with open(manifest_path) as fh:
        manifest = json.load(fh)

    entry_dicts: list[dict[str, Any]] = manifest["entries"]
    entries = _build_query_log_entries(entry_dicts)
    log.info("Loaded %d query log entries from %s", len(entries), manifest_path)

    if not entries:
        log.info("No query log entries to push.")
        summary = {
            "resource_uuid": resource_uuid,
            "log_type": LOG_TYPE,
            "invocation_ids": [],
            "pushed_at": datetime.now(timezone.utc).isoformat(),
            "query_log_count": 0,
            "batch_count": 0,
            "batch_size": batch_size,
        }
        push_manifest_path = manifest_path.replace(".json", "_push_result.json")
        with open(push_manifest_path, "w") as fh:
            json.dump(summary, fh, indent=2)
        return summary

    # Split into batches
    batches = []
    for i in range(0, len(entries), batch_size):
        batches.append(entries[i : i + batch_size])
    total_batches = len(batches)

    def _push_batch(batch: list, batch_num: int) -> str | None:
        """Push a single batch using a dedicated Session (thread-safe)."""
        client = Client(session=Session(mcd_id=key_id, mcd_token=key_token, scope="Ingestion"))
        service = IngestionService(mc_client=client)
        result = service.send_query_logs(
            resource_uuid=resource_uuid,
            log_type=LOG_TYPE,
            events=batch,
        )
        invocation_id = service.extract_invocation_id(result)
        log.info("Pushed batch %d/%d (%d entries) — invocation_id=%s", batch_num, total_batches, len(batch), invocation_id)
        return invocation_id

    # Push batches in parallel (each thread gets its own pycarlo Session)
    max_workers = min(4, total_batches)
    invocation_ids: list[str | None] = [None] * total_batches

    with ThreadPoolExecutor(max_workers=max_workers) as pool:
        futures = {
            pool.submit(_push_batch, batch, i + 1): i
            for i, batch in enumerate(batches)
        }
        for future in as_completed(futures):
            idx = futures[future]
            try:
                invocation_ids[idx] = future.result()
            except Exception as exc:
                log.error("ERROR pushing batch %d: %s", idx + 1, exc)
                raise

    log.info("All %d batches pushed (%d workers)", total_batches, max_workers)

    summary = {
        "resource_uuid": resource_uuid,
        "log_type": LOG_TYPE,
        "invocation_ids": invocation_ids,
        "pushed_at": datetime.now(timezone.utc).isoformat(),
        "query_log_count": len(entries),
        "batch_count": total_batches,
        "batch_size": batch_size,
        "lookback_hours": manifest.get("lookback_hours"),
        "lookback_lag_hours": manifest.get("lookback_lag_hours"),
    }

    push_manifest_path = manifest_path.replace(".json", "_push_result.json")
    with open(push_manifest_path, "w") as fh:
        json.dump(summary, fh, indent=2)
    log.info("Push result written to %s", push_manifest_path)

    return summary


def main() -> None:
    parser = argparse.ArgumentParser(description="Push Redshift query logs to Monte Carlo from manifest")
    parser.add_argument("--manifest", default="manifest_query_logs.json")
    parser.add_argument("--resource-uuid", default=os.getenv("MCD_RESOURCE_UUID"))
    parser.add_argument("--key-id", default=os.getenv("MCD_INGEST_ID"))
    parser.add_argument("--key-token", default=os.getenv("MCD_INGEST_TOKEN"))
    parser.add_argument("--batch-size", type=int, default=DEFAULT_BATCH_SIZE)
    args = parser.parse_args()

    required = ["resource_uuid", "key_id", "key_token"]
    missing = [k for k in required if getattr(args, k) is None]
    if missing:
        parser.error(f"Missing required arguments/env vars: {missing}")

    push(
        manifest_path=args.manifest,
        resource_uuid=args.resource_uuid,
        key_id=args.key_id,
        key_token=args.key_token,
        batch_size=args.batch_size,
    )


if __name__ == "__main__":
    main()

#!/usr/bin/env python3
"""
Collect lineage from Snowflake and push it to Monte Carlo — combined.

Imports ``collect()`` from ``collect_lineage`` and ``push()`` from
``push_lineage``, runs both in sequence.

Substitution points
-------------------
- SNOWFLAKE_ACCOUNT    (env) / --account    (CLI) : Snowflake account identifier
- SNOWFLAKE_USER       (env) / --user       (CLI) : Snowflake username
- SNOWFLAKE_PASSWORD   (env) / --password   (CLI) : Snowflake password
- SNOWFLAKE_WAREHOUSE  (env) / --warehouse  (CLI) : Snowflake virtual warehouse
- MCD_INGEST_ID     (env) / --key-id     (CLI) : Monte Carlo ingestion key ID
- MCD_INGEST_TOKEN  (env) / --key-token  (CLI) : Monte Carlo ingestion key token
- MCD_RESOURCE_UUID     (env) / --resource-uuid (CLI) : MC resource UUID for this connection

Prerequisites
-------------
    pip install pycarlo snowflake-connector-python

Usage (table-level):
    python collect_and_push_lineage.py \\
        --account  <SNOWFLAKE_ACCOUNT> \\
        --user     <SNOWFLAKE_USER> \\
        --password <SNOWFLAKE_PASSWORD> \\
        --warehouse <SNOWFLAKE_WAREHOUSE> \\
        --key-id  <MCD_INGEST_ID> \\
        --key-token <MCD_INGEST_TOKEN> \\
        --resource-uuid <MCD_RESOURCE_UUID>

Usage (column-level):
    python collect_and_push_lineage.py ... --column-lineage
"""

from __future__ import annotations

import argparse
import os

from collect_lineage import collect, _LOOKBACK_HOURS
from push_lineage import push, _BATCH_SIZE


def main() -> None:
    parser = argparse.ArgumentParser(
        description="Collect Snowflake lineage from ACCOUNT_USAGE and push to Monte Carlo",
    )
    parser.add_argument(
        "--account",
        default=os.environ.get("SNOWFLAKE_ACCOUNT"),
        help="Snowflake account identifier (env: SNOWFLAKE_ACCOUNT)",
    )
    parser.add_argument(
        "--user",
        default=os.environ.get("SNOWFLAKE_USER"),
        help="Snowflake username (env: SNOWFLAKE_USER)",
    )
    parser.add_argument(
        "--password",
        default=os.environ.get("SNOWFLAKE_PASSWORD"),
        help="Snowflake password (env: SNOWFLAKE_PASSWORD)",
    )
    parser.add_argument(
        "--warehouse",
        default=os.environ.get("SNOWFLAKE_WAREHOUSE"),
        help="Snowflake virtual warehouse (env: SNOWFLAKE_WAREHOUSE)",
    )
    parser.add_argument(
        "--key-id",
        default=os.environ.get("MCD_INGEST_ID"),
        help="Monte Carlo ingestion key ID (env: MCD_INGEST_ID)",
    )
    parser.add_argument(
        "--key-token",
        default=os.environ.get("MCD_INGEST_TOKEN"),
        help="Monte Carlo ingestion key token (env: MCD_INGEST_TOKEN)",
    )
    parser.add_argument(
        "--resource-uuid",
        default=os.environ.get("MCD_RESOURCE_UUID"),
        help="Monte Carlo resource UUID for this Snowflake connection (env: MCD_RESOURCE_UUID)",
    )
    parser.add_argument(
        "--lookback-hours",
        type=int,
        default=_LOOKBACK_HOURS,
        help=f"Hours of QUERY_HISTORY to scan (default: {_LOOKBACK_HOURS})",
    )
    parser.add_argument(
        "--column-lineage",
        action="store_true",
        help="Push column-level lineage instead of table-level",
    )
    parser.add_argument(
        "--output-file",
        default="lineage_output.json",
        help="Path for the intermediate collect manifest (default: lineage_output.json)",
    )
    parser.add_argument(
        "--push-result-file",
        default="lineage_push_result.json",
        help="Path to write the push result (default: lineage_push_result.json)",
    )
    parser.add_argument(
        "--batch-size",
        type=int,
        default=_BATCH_SIZE,
        help=f"Max events per push batch (default: {_BATCH_SIZE})",
    )
    args = parser.parse_args()

    missing = [
        name
        for name, val in [
            ("--account", args.account),
            ("--user", args.user),
            ("--password", args.password),
            ("--warehouse", args.warehouse),
            ("--key-id", args.key_id),
            ("--key-token", args.key_token),
            ("--resource-uuid", args.resource_uuid),
        ]
        if not val
    ]
    if missing:
        parser.error(f"Missing required arguments: {', '.join(missing)}")

    # Step 1: Collect
    collect(
        account=args.account,
        user=args.user,
        password=args.password,
        warehouse=args.warehouse,
        lookback_hours=args.lookback_hours,
        column_lineage=args.column_lineage,
        output_file=args.output_file,
    )

    # Step 2: Push
    push(
        input_file=args.output_file,
        resource_uuid=args.resource_uuid,
        key_id=args.key_id,
        key_token=args.key_token,
        batch_size=args.batch_size,
        output_file=args.push_result_file,
    )

    print("Done.")


if __name__ == "__main__":
    main()

#!/usr/bin/env python3
"""
Collect table metadata from Snowflake and push it to Monte Carlo — combined.

Imports ``collect()`` from ``collect_metadata`` and ``push()`` from
``push_metadata``, runs both in sequence.

Substitution points
-------------------
- SNOWFLAKE_ACCOUNT    (env) / --account    (CLI) : Snowflake account identifier (e.g. xy12345.us-east-1)
- SNOWFLAKE_USER       (env) / --user       (CLI) : Snowflake username
- SNOWFLAKE_PASSWORD   (env) / --password   (CLI) : Snowflake password
- SNOWFLAKE_WAREHOUSE  (env) / --warehouse  (CLI) : Snowflake virtual warehouse
- MCD_INGEST_ID     (env) / --key-id     (CLI) : Monte Carlo ingestion key ID
- MCD_INGEST_TOKEN  (env) / --key-token  (CLI) : Monte Carlo ingestion key token
- MCD_RESOURCE_UUID     (env) / --resource-uuid (CLI) : MC resource UUID for this connection

Prerequisites
-------------
    pip install pycarlo snowflake-connector-python

Usage
-----
    python collect_and_push_metadata.py \\
        --account  <SNOWFLAKE_ACCOUNT> \\
        --user     <SNOWFLAKE_USER> \\
        --password <SNOWFLAKE_PASSWORD> \\
        --warehouse <SNOWFLAKE_WAREHOUSE> \\
        --key-id  <MCD_INGEST_ID> \\
        --key-token <MCD_INGEST_TOKEN> \\
        --resource-uuid <MCD_RESOURCE_UUID>
"""

import argparse
import os

from collect_metadata import collect
from push_metadata import push, _BATCH_SIZE


def main() -> None:
    parser = argparse.ArgumentParser(
        description="Collect Snowflake table metadata and push to Monte Carlo",
    )
    parser.add_argument(
        "--account",
        default=os.environ.get("SNOWFLAKE_ACCOUNT"),
        help="Snowflake account identifier, e.g. xy12345.us-east-1 (env: SNOWFLAKE_ACCOUNT)",  # ← SUBSTITUTE
    )
    parser.add_argument(
        "--user",
        default=os.environ.get("SNOWFLAKE_USER"),
        help="Snowflake username (env: SNOWFLAKE_USER)",  # ← SUBSTITUTE
    )
    parser.add_argument(
        "--password",
        default=os.environ.get("SNOWFLAKE_PASSWORD"),
        help="Snowflake password (env: SNOWFLAKE_PASSWORD)",  # ← SUBSTITUTE
    )
    parser.add_argument(
        "--warehouse",
        default=os.environ.get("SNOWFLAKE_WAREHOUSE"),
        help="Snowflake virtual warehouse (env: SNOWFLAKE_WAREHOUSE)",  # ← SUBSTITUTE
    )
    parser.add_argument(
        "--key-id",
        default=os.environ.get("MCD_INGEST_ID"),
        help="Monte Carlo ingestion key ID (env: MCD_INGEST_ID)",
    )
    parser.add_argument(
        "--key-token",
        default=os.environ.get("MCD_INGEST_TOKEN"),
        help="Monte Carlo ingestion key token (env: MCD_INGEST_TOKEN)",
    )
    parser.add_argument(
        "--resource-uuid",
        default=os.environ.get("MCD_RESOURCE_UUID"),
        help="Monte Carlo resource UUID for this Snowflake connection (env: MCD_RESOURCE_UUID)",
    )
    parser.add_argument(
        "--output-file",
        default="metadata_output.json",
        help="Path for the intermediate collect manifest (default: metadata_output.json)",
    )
    parser.add_argument(
        "--push-result-file",
        default="metadata_push_result.json",
        help="Path to write the push result (default: metadata_push_result.json)",
    )
    parser.add_argument(
        "--batch-size",
        type=int,
        default=_BATCH_SIZE,
        help=f"Max assets per push batch (default: {_BATCH_SIZE})",
    )
    args = parser.parse_args()

    missing = [
        name
        for name, val in [
            ("--account", args.account),
            ("--user", args.user),
            ("--password", args.password),
            ("--warehouse", args.warehouse),
            ("--key-id", args.key_id),
            ("--key-token", args.key_token),
            ("--resource-uuid", args.resource_uuid),
        ]
        if not val
    ]
    if missing:
        parser.error(f"Missing required arguments: {', '.join(missing)}")

    # Step 1: Collect
    collect(
        account=args.account,
        user=args.user,
        password=args.password,
        warehouse=args.warehouse,
        output_file=args.output_file,
    )

    # Step 2: Push
    push(
        input_file=args.output_file,
        resource_uuid=args.resource_uuid,
        key_id=args.key_id,
        key_token=args.key_token,
        batch_size=args.batch_size,
        output_file=args.push_result_file,
    )

    print("Done.")


if __name__ == "__main__":
    main()

#!/usr/bin/env python3
"""
Collect query logs from Snowflake and push them to Monte Carlo — combined.

Imports ``collect()`` from ``collect_query_logs`` and ``push()`` from
``push_query_logs``, runs both in sequence.

Substitution points
-------------------
- SNOWFLAKE_ACCOUNT    (env) / --account    (CLI) : Snowflake account identifier
- SNOWFLAKE_USER       (env) / --user       (CLI) : Snowflake username
- SNOWFLAKE_PASSWORD   (env) / --password   (CLI) : Snowflake password
- SNOWFLAKE_WAREHOUSE  (env) / --warehouse  (CLI) : Snowflake virtual warehouse
- MCD_INGEST_ID     (env) / --key-id     (CLI) : Monte Carlo ingestion key ID
- MCD_INGEST_TOKEN  (env) / --key-token  (CLI) : Monte Carlo ingestion key token
- MCD_RESOURCE_UUID     (env) / --resource-uuid (CLI) : MC resource UUID for this connection

Prerequisites
-------------
    pip install pycarlo snowflake-connector-python

Usage
-----
    python collect_and_push_query_logs.py \\
        --account  <SNOWFLAKE_ACCOUNT> \\
        --user     <SNOWFLAKE_USER> \\
        --password <SNOWFLAKE_PASSWORD> \\
        --warehouse <SNOWFLAKE_WAREHOUSE> \\
        --key-id  <MCD_INGEST_ID> \\
        --key-token <MCD_INGEST_TOKEN> \\
        --resource-uuid <MCD_RESOURCE_UUID>
"""

import argparse
import os

from collect_query_logs import collect
from push_query_logs import push, _BATCH_SIZE


def main() -> None:
    parser = argparse.ArgumentParser(
        description="Collect Snowflake query logs from ACCOUNT_USAGE and push to Monte Carlo",
    )
    parser.add_argument(
        "--account",
        default=os.environ.get("SNOWFLAKE_ACCOUNT"),
        help="Snowflake account identifier, e.g. xy12345.us-east-1 (env: SNOWFLAKE_ACCOUNT)",  # ← SUBSTITUTE
    )
    parser.add_argument(
        "--user",
        default=os.environ.get("SNOWFLAKE_USER"),
        help="Snowflake username (env: SNOWFLAKE_USER)",
    )
    parser.add_argument(
        "--password",
        default=os.environ.get("SNOWFLAKE_PASSWORD"),
        help="Snowflake password (env: SNOWFLAKE_PASSWORD)",
    )
    parser.add_argument(
        "--warehouse",
        default=os.environ.get("SNOWFLAKE_WAREHOUSE"),
        help="Snowflake virtual warehouse (env: SNOWFLAKE_WAREHOUSE)",  # ← SUBSTITUTE
    )
    parser.add_argument(
        "--key-id",
        default=os.environ.get("MCD_INGEST_ID"),
        help="Monte Carlo ingestion key ID (env: MCD_INGEST_ID)",
    )
    parser.add_argument(
        "--key-token",
        default=os.environ.get("MCD_INGEST_TOKEN"),
        help="Monte Carlo ingestion key token (env: MCD_INGEST_TOKEN)",
    )
    parser.add_argument(
        "--resource-uuid",
        default=os.environ.get("MCD_RESOURCE_UUID"),
        help="Monte Carlo resource UUID for this Snowflake connection (env: MCD_RESOURCE_UUID)",
    )
    parser.add_argument(
        "--output-file",
        default="query_logs_output.json",
        help="Path for the intermediate collect manifest (default: query_logs_output.json)",
    )
    parser.add_argument(
        "--push-result-file",
        default="query_logs_push_result.json",
        help="Path to write the push result (default: query_logs_push_result.json)",
    )
    parser.add_argument(
        "--batch-size",
        type=int,
        default=_BATCH_SIZE,
        help=f"Max entries per push batch (default: {_BATCH_SIZE})",
    )
    args = parser.parse_args()

    missing = [
        name
        for name, val in [
            ("--account", args.account),
            ("--user", args.user),
            ("--password", args.password),
            ("--warehouse", args.warehouse),
            ("--key-id", args.key_id),
            ("--key-token", args.key_token),
            ("--resource-uuid", args.resource_uuid),
        ]
        if not val
    ]
    if missing:
        parser.error(f"Missing required arguments: {', '.join(missing)}")

    # Step 1: Collect
    collect(
        account=args.account,
        user=args.user,
        password=args.password,
        warehouse=args.warehouse,
        output_file=args.output_file,
    )

    # Step 2: Push
    push(
        input_file=args.output_file,
        resource_uuid=args.resource_uuid,
        key_id=args.key_id,
        key_token=args.key_token,
        batch_size=args.batch_size,
        output_file=args.push_result_file,
    )

    print("Done.")


if __name__ == "__main__":
    main()

#!/usr/bin/env python3
"""
Collect table and column lineage from Snowflake — collection only.

Queries ACCOUNT_USAGE for DML/DDL statements in the last 24 hours, parses each
QUERY_TEXT with regex to extract source and destination tables, then writes the
resulting lineage edges to a JSON manifest file.

Can be run standalone via CLI or imported (use the ``collect()`` function).

Note: ACCOUNT_USAGE views have an approximate latency of 45 minutes, so very
recent queries may not yet appear.

Substitution points
-------------------
- SNOWFLAKE_ACCOUNT    (env) / --account    (CLI) : Snowflake account identifier
- SNOWFLAKE_USER       (env) / --user       (CLI) : Snowflake username
- SNOWFLAKE_PASSWORD   (env) / --password   (CLI) : Snowflake password
- SNOWFLAKE_WAREHOUSE  (env) / --warehouse  (CLI) : Snowflake virtual warehouse

Prerequisites
-------------
    pip install snowflake-connector-python

Usage (table-level):
    python collect_lineage.py \\
        --account  <SNOWFLAKE_ACCOUNT> \\
        --user     <SNOWFLAKE_USER> \\
        --password <SNOWFLAKE_PASSWORD> \\
        --warehouse <SNOWFLAKE_WAREHOUSE>

Usage (column-level):
    python collect_lineage.py ... --column-lineage
"""

from __future__ import annotations

import argparse
import json
import os
import re
from dataclasses import dataclass, field
from datetime import datetime, timezone

import snowflake.connector

# ← SUBSTITUTE: set RESOURCE_TYPE to match your Monte Carlo connection type
RESOURCE_TYPE = "snowflake"


def _check_available_memory(min_gb: float = 2.0) -> None:
    """Warn if available memory is below the threshold."""
    try:
        if hasattr(os, "sysconf"):  # Linux / macOS
            page_size = os.sysconf("SC_PAGE_SIZE")
            avail_pages = os.sysconf("SC_AVPHYS_PAGES")
            avail_gb = (page_size * avail_pages) / (1024 ** 3)
        else:
            return  # Windows — skip check
    except (ValueError, OSError):
        return
    if avail_gb < min_gb:
        print(
            f"WARNING: Only {avail_gb:.1f} GB of memory available "
            f"(minimum recommended: {min_gb:.1f} GB). "
            f"Consider reducing the lookback window or increasing available memory."
        )

# Hours to look back in ACCOUNT_USAGE.QUERY_HISTORY
# ← SUBSTITUTE: adjust the lookback window to match your collection cadence
_LOOKBACK_HOURS = 24

# Regex for CTAS: CREATE [OR REPLACE] [TRANSIENT] TABLE [IF NOT EXISTS] [db.][schema.]table AS SELECT
_CTAS_RE = re.compile(
    r"CREATE\s+(?:OR\s+REPLACE\s+)?(?:TRANSIENT\s+)?TABLE\s+(?:IF\s+NOT\s+EXISTS\s+)?"
    r"(?:(?P<dest_db>\w+)\.)?(?:(?P<dest_schema>\w+)\.)?(?P<dest_table>\w+)"
    r".*?AS\s+SELECT\s+(?P<select_cols>.+?)\s+FROM\s+"
    r"(?:(?P<src_db>\w+)\.)?(?:(?P<src_schema>\w+)\.)?(?P<src_table>\w+)",
    re.IGNORECASE | re.DOTALL,
)

# Regex for INSERT INTO [db.][schema.]table SELECT ... FROM [db.][schema.]table
_INSERT_RE = re.compile(
    r"INSERT\s+(?:INTO|OVERWRITE)\s+"
    r"(?:(?P<dest_db>\w+)\.)?(?:(?P<dest_schema>\w+)\.)?(?P<dest_table>\w+)"
    r".*?SELECT\s+(?P<select_cols>.+?)\s+FROM\s+"
    r"(?:(?P<src_db>\w+)\.)?(?:(?P<src_schema>\w+)\.)?(?P<src_table>\w+)",
    re.IGNORECASE | re.DOTALL,
)

# Regex for CREATE [OR REPLACE] VIEW [db.][schema.]view AS SELECT ... FROM ...
_CREATE_VIEW_RE = re.compile(
    r"CREATE\s+(?:OR\s+REPLACE\s+)?(?:SECURE\s+)?VIEW\s+"
    r"(?:(?P<dest_db>\w+)\.)?(?:(?P<dest_schema>\w+)\.)?(?P<dest_table>\w+)"
    r".*?AS\s+SELECT\s+(?P<select_cols>.+?)\s+FROM\s+"
    r"(?:(?P<src_db>\w+)\.)?(?:(?P<src_schema>\w+)\.)?(?P<src_table>\w+)",
    re.IGNORECASE | re.DOTALL,
)

# Additional JOIN sources
_JOIN_RE = re.compile(
    r"JOIN\s+(?:(?P<src_db>\w+)\.)?(?:(?P<src_schema>\w+)\.)?(?P<src_table>\w+)",
    re.IGNORECASE,
)

# Simple column alias extraction from SELECT clause
_COL_RE = re.compile(r"(?:(\w+)\.)?(\w+)(?:\s+AS\s+(\w+))?", re.IGNORECASE)
_SQL_KEYWORDS = {
    "FROM", "SELECT", "WHERE", "JOIN", "ON", "AS", "*", "AND", "OR",
    "GROUP", "ORDER", "BY", "HAVING", "LIMIT", "DISTINCT", "CASE", "WHEN",
    "THEN", "ELSE", "END", "NULL", "NOT", "IN", "IS", "BETWEEN",
}


@dataclass
class _LineageEdge:
    dest_db: str
    dest_schema: str
    dest_table: str
    sources: list[tuple[str, str, str]] = field(default_factory=list)
    # col_mappings: (dest_col, src_table, src_col)
    col_mappings: list[tuple[str, str, str]] = field(default_factory=list)


def _parse_select_cols(select_clause: str, src_table: str) -> list[tuple[str, str, str]]:
    mappings = []
    for m in _COL_RE.finditer(select_clause):
        src_col = m.group(2)
        dest_col = m.group(3) or src_col
        if src_col.upper() in _SQL_KEYWORDS:
            continue
        mappings.append((dest_col, src_table, src_col))
    return mappings


def _parse_edges(rows: list[dict]) -> list[_LineageEdge]:
    """Parse QUERY_HISTORY rows into _LineageEdge objects."""
    edges: dict[str, _LineageEdge] = {}

    for row in rows:
        query_text = row.get("QUERY_TEXT") or ""
        default_db = (row.get("DATABASE_NAME") or "").lower()
        sql_clean = re.sub(r"\s+", " ", query_text).strip()

        for pattern in (_CTAS_RE, _INSERT_RE, _CREATE_VIEW_RE):
            m = pattern.search(sql_clean)
            if not m:
                continue

            dest_db = (m.group("dest_db") or default_db).lower()
            dest_schema = (m.group("dest_schema") or "public").lower()
            dest_table = m.group("dest_table").lower()
            src_db = (m.group("src_db") or default_db).lower()
            src_schema = (m.group("src_schema") or "public").lower()
            src_table = m.group("src_table").lower()
            select_cols = m.group("select_cols")

            key = f"{dest_db}.{dest_schema}.{dest_table}"
            if key not in edges:
                edges[key] = _LineageEdge(
                    dest_db=dest_db, dest_schema=dest_schema, dest_table=dest_table
                )

            edge = edges[key]
            src_triple = (src_db, src_schema, src_table)
            if src_triple not in edge.sources:
                edge.sources.append(src_triple)

            for jm in _JOIN_RE.finditer(sql_clean):
                jt = jm.group("src_table").lower()
                jschema = (jm.group("src_schema") or src_schema).lower()
                jdb = (jm.group("src_db") or src_db).lower()
                jp = (jdb, jschema, jt)
                if jp not in edge.sources:
                    edge.sources.append(jp)

            edge.col_mappings.extend(_parse_select_cols(select_cols, src_table))
            break

    return list(edges.values())


def _fetch_query_history(conn, lookback_hours: int) -> list[dict]:
    cursor = conn.cursor()
    cursor.execute(
        f"""
        SELECT QUERY_ID, QUERY_TEXT, START_TIME, END_TIME, USER_NAME, DATABASE_NAME, EXECUTION_STATUS
        FROM SNOWFLAKE.ACCOUNT_USAGE.QUERY_HISTORY
        WHERE START_TIME >= DATEADD(hour, -{lookback_hours}, CURRENT_TIMESTAMP())
          AND EXECUTION_STATUS = 'SUCCESS'
          AND QUERY_TYPE IN ('CREATE_TABLE_AS_SELECT', 'INSERT', 'MERGE', 'CREATE_VIEW')
        ORDER BY START_TIME
        LIMIT 50000
        """
        # ← SUBSTITUTE: adjust QUERY_TYPE list, LIMIT, or add a WHERE clause to scope to specific databases
    )
    columns = [col[0] for col in cursor.description]
    rows = []
    while True:
        batch = cursor.fetchmany(1000)
        if not batch:
            break
        rows.extend(dict(zip(columns, row)) for row in batch)
    cursor.close()
    return rows


def collect(
    account: str,
    user: str,
    password: str,
    warehouse: str,
    lookback_hours: int = _LOOKBACK_HOURS,
    column_lineage: bool = False,
    output_file: str = "lineage_output.json",
) -> dict:
    """
    Connect to Snowflake, collect lineage edges, and write a JSON manifest.

    Returns the manifest dict.
    """
    _check_available_memory()
    print(f"Connecting to Snowflake account: {account} ...")
    conn = snowflake.connector.connect(
        account=account,
        user=user,
        password=password,
        warehouse=warehouse,
    )

    print(f"Fetching QUERY_HISTORY for the last {lookback_hours} hour(s) ...")
    rows = _fetch_query_history(conn, lookback_hours)
    conn.close()
    print(f"  Retrieved {len(rows)} qualifying query/queries.")

    if not rows:
        print("No lineage queries found in the specified window.")
        manifest = {
            "resource_type": RESOURCE_TYPE,
            "collected_at": datetime.now(tz=timezone.utc).isoformat(),
            "column_lineage": column_lineage,
            "edges": [],
        }
        with open(output_file, "w") as fh:
            json.dump(manifest, fh, indent=2)
        return manifest

    edges = _parse_edges(rows)
    print(f"  Parsed {len(edges)} lineage edge(s).")

    manifest = {
        "resource_type": RESOURCE_TYPE,
        "collected_at": datetime.now(tz=timezone.utc).isoformat(),
        "column_lineage": column_lineage,
        "edges": [
            {
                "destination": {
                    "database": e.dest_db,
                    "schema": e.dest_schema,
                    "table": e.dest_table,
                },
                "sources": [
                    {"database": sdb, "schema": sschema, "table": stbl}
                    for sdb, sschema, stbl in e.sources
                ],
                "col_mappings": [
                    {"dest_col": dc, "src_table": st, "src_col": sc}
                    for dc, st, sc in e.col_mappings
                ],
            }
            for e in edges
        ],
    }
    with open(output_file, "w") as fh:
        json.dump(manifest, fh, indent=2)
    print(f"Lineage manifest written to {output_file}")

    return manifest


def main() -> None:
    parser = argparse.ArgumentParser(
        description="Collect Snowflake lineage from ACCOUNT_USAGE and write to a manifest file",
    )
    parser.add_argument(
        "--account",
        default=os.environ.get("SNOWFLAKE_ACCOUNT"),
        help="Snowflake account identifier (env: SNOWFLAKE_ACCOUNT)",
    )
    parser.add_argument(
        "--user",
        default=os.environ.get("SNOWFLAKE_USER"),
        help="Snowflake username (env: SNOWFLAKE_USER)",
    )
    parser.add_argument(
        "--password",
        default=os.environ.get("SNOWFLAKE_PASSWORD"),
        help="Snowflake password (env: SNOWFLAKE_PASSWORD)",
    )
    parser.add_argument(
        "--warehouse",
        default=os.environ.get("SNOWFLAKE_WAREHOUSE"),
        help="Snowflake virtual warehouse (env: SNOWFLAKE_WAREHOUSE)",
    )
    parser.add_argument(
        "--lookback-hours",
        type=int,
        default=_LOOKBACK_HOURS,
        help=f"Hours of QUERY_HISTORY to scan (default: {_LOOKBACK_HOURS})",
    )
    parser.add_argument(
        "--column-lineage",
        action="store_true",
        help="Include column-level lineage mappings in the manifest",
    )
    parser.add_argument(
        "--output-file",
        default="lineage_output.json",
        help="Path to write the lineage manifest (default: lineage_output.json)",
    )
    args = parser.parse_args()

    missing = [
        name
        for name, val in [
            ("--account", args.account),
            ("--user", args.user),
            ("--password", args.password),
            ("--warehouse", args.warehouse),
        ]
        if not val
    ]
    if missing:
        parser.error(f"Missing required arguments: {', '.join(missing)}")

    collect(
        account=args.account,
        user=args.user,
        password=args.password,
        warehouse=args.warehouse,
        lookback_hours=args.lookback_hours,
        column_lineage=args.column_lineage,
        output_file=args.output_file,
    )
    print("Done.")


if __name__ == "__main__":
    main()

#!/usr/bin/env python3
"""
Collect table metadata from Snowflake — collection only.

Connects to Snowflake, discovers all accessible databases and schemas, then
queries INFORMATION_SCHEMA.TABLES for volume/freshness and
INFORMATION_SCHEMA.COLUMNS for field definitions.  The collected assets are
written to a JSON manifest file.

Can be run standalone via CLI or imported (use the ``collect()`` function).

Substitution points
-------------------
- SNOWFLAKE_ACCOUNT    (env) / --account    (CLI) : Snowflake account identifier (e.g. xy12345.us-east-1)
- SNOWFLAKE_USER       (env) / --user       (CLI) : Snowflake username
- SNOWFLAKE_PASSWORD   (env) / --password   (CLI) : Snowflake password
- SNOWFLAKE_WAREHOUSE  (env) / --warehouse  (CLI) : Snowflake virtual warehouse

Prerequisites
-------------
    pip install snowflake-connector-python

Usage
-----
    python collect_metadata.py \\
        --account  <SNOWFLAKE_ACCOUNT> \\
        --user     <SNOWFLAKE_USER> \\
        --password <SNOWFLAKE_PASSWORD> \\
        --warehouse <SNOWFLAKE_WAREHOUSE>
"""

import argparse
import json
import os
from datetime import datetime, timezone

import snowflake.connector

# ← SUBSTITUTE: set RESOURCE_TYPE to match your Monte Carlo connection type
RESOURCE_TYPE = "snowflake"


def _check_available_memory(min_gb: float = 2.0) -> None:
    """Warn if available memory is below the threshold."""
    try:
        if hasattr(os, "sysconf"):  # Linux / macOS
            page_size = os.sysconf("SC_PAGE_SIZE")
            avail_pages = os.sysconf("SC_AVPHYS_PAGES")
            avail_gb = (page_size * avail_pages) / (1024 ** 3)
        else:
            return  # Windows — skip check
    except (ValueError, OSError):
        return
    if avail_gb < min_gb:
        print(
            f"WARNING: Only {avail_gb:.1f} GB of memory available "
            f"(minimum recommended: {min_gb:.1f} GB). "
            f"Consider reducing the lookback window or increasing available memory."
        )

# Databases that are Snowflake system databases — skip them
_SKIP_DATABASES = {"SNOWFLAKE", "SNOWFLAKE_SAMPLE_DATA"}

# Schemas that are Snowflake system schemas — skip them
_SKIP_SCHEMAS = {"INFORMATION_SCHEMA"}


# Snowflake TABLE_TYPE → Monte Carlo RelationalAsset.type mapping.
# The MC API only accepts "TABLE" or "VIEW" (uppercase).
_TABLE_TYPE_MAP = {
    "BASE TABLE": "TABLE",
    "TABLE": "TABLE",
    "DYNAMIC TABLE": "TABLE",
    "EXTERNAL TABLE": "TABLE",
    "VIEW": "VIEW",
    "MATERIALIZED VIEW": "VIEW",
    "SECURE VIEW": "VIEW",
}


def _normalize_table_type(raw_type: str | None) -> str:
    """Map Snowflake's TABLE_TYPE value to MC-accepted 'TABLE' or 'VIEW'."""
    if not raw_type:
        return "TABLE"
    return _TABLE_TYPE_MAP.get(raw_type.upper(), "TABLE")


def _connect(account: str, user: str, password: str, warehouse: str):
    # ← SUBSTITUTE: add role= or authenticator= kwargs if your org requires them
    return snowflake.connector.connect(
        account=account,
        user=user,
        password=password,
        warehouse=warehouse,
    )


def _collect_assets(conn) -> list[dict]:
    """Collect table metadata from Snowflake and return as a list of dicts."""
    cursor = conn.cursor()
    assets: list[dict] = []

    # --- Discover databases ---
    cursor.execute("SHOW DATABASES")
    # SHOW DATABASES returns (created_on, name, …); column index 1 is the name
    all_db_rows = []
    while True:
        chunk = cursor.fetchmany(1000)
        if not chunk:
            break
        all_db_rows.extend(chunk)
    databases = [row[1] for row in all_db_rows if row[1] not in _SKIP_DATABASES]
    print(f"  Found {len(databases)} database(s): {databases}")

    for db in databases:
        # --- Discover schemas in each database ---
        try:
            cursor.execute(f'SHOW SCHEMAS IN DATABASE "{db}"')
        except Exception as exc:
            print(f"  WARNING: could not list schemas in {db}: {exc}")
            continue

        # Column index 1 is the schema name
        all_schema_rows = []
        while True:
            chunk = cursor.fetchmany(1000)
            if not chunk:
                break
            all_schema_rows.extend(chunk)
        schemas = [row[1] for row in all_schema_rows if row[1] not in _SKIP_SCHEMAS]

        # --- Collect tables, volume, and freshness via INFORMATION_SCHEMA ---
        try:
            cursor.execute(
                f"""
                SELECT
                    TABLE_CATALOG,
                    TABLE_SCHEMA,
                    TABLE_NAME,
                    TABLE_TYPE,
                    ROW_COUNT,
                    BYTES,
                    LAST_ALTERED,
                    COMMENT
                FROM "{db}".INFORMATION_SCHEMA.TABLES
                WHERE TABLE_SCHEMA != 'INFORMATION_SCHEMA'
                ORDER BY TABLE_SCHEMA, TABLE_NAME
                """
            )
        except Exception as exc:
            print(f"  WARNING: could not query INFORMATION_SCHEMA.TABLES in {db}: {exc}")
            continue

        table_rows = []
        while True:
            chunk = cursor.fetchmany(1000)
            if not chunk:
                break
            table_rows.extend(chunk)
        print(f"  {db}: {len(table_rows)} table(s)")

        # Build a set of schema names present in the table result to know which
        # INFORMATION_SCHEMA.COLUMNS queries to run
        schemas_with_tables: set[str] = {row[1] for row in table_rows}

        # Pre-fetch all columns for this database in one query per schema
        columns_by_table: dict[tuple[str, str], list[dict]] = {}
        for schema in schemas_with_tables:
            if schema not in schemas:
                continue  # respect the earlier schema skip list
            try:
                cursor.execute(
                    f"""
                    SELECT TABLE_NAME, COLUMN_NAME, DATA_TYPE, COMMENT
                    FROM "{db}".INFORMATION_SCHEMA.COLUMNS
                    WHERE TABLE_SCHEMA = %s
                    ORDER BY TABLE_NAME, ORDINAL_POSITION
                    """,
                    (schema,),
                )
            except Exception as exc:
                print(f"  WARNING: could not fetch columns for {db}.{schema}: {exc}")
                continue

            all_col_rows = []
            while True:
                chunk = cursor.fetchmany(1000)
                if not chunk:
                    break
                all_col_rows.extend(chunk)
            for col_row in all_col_rows:
                table_name, col_name, data_type, col_comment = col_row
                key = (schema, table_name)
                if key not in columns_by_table:
                    columns_by_table[key] = []
                columns_by_table[key].append(
                    {
                        "name": col_name,
                        "type": data_type,
                        "description": col_comment or None,
                    }
                )

        # Build asset dicts
        for row in table_rows:
            tbl_catalog, tbl_schema, tbl_name, tbl_type, row_count, byte_count, last_altered, tbl_comment = row

            volume = None
            if row_count is not None or byte_count is not None:
                volume = {
                    "row_count": int(row_count) if row_count is not None else None,
                    "byte_count": int(byte_count) if byte_count is not None else None,
                }

            freshness = None
            if last_altered is not None:
                freshness = {
                    "last_update_time": last_altered.isoformat() if hasattr(last_altered, "isoformat") else str(last_altered),
                }

            fields = columns_by_table.get((tbl_schema, tbl_name), [])

            assets.append(
                {
                    "type": _normalize_table_type(tbl_type),
                    "database": tbl_catalog,
                    "schema": tbl_schema,
                    "name": tbl_name,
                    "description": tbl_comment or None,
                    "fields": fields,
                    "volume": volume,
                    "freshness": freshness,
                }
            )
            print(f"    + {tbl_catalog}.{tbl_schema}.{tbl_name} ({len(fields)} columns)")

    cursor.close()
    return assets


def collect(
    account: str,
    user: str,
    password: str,
    warehouse: str,
    output_file: str = "metadata_output.json",
) -> dict:
    """
    Connect to Snowflake, collect table metadata, and write a JSON manifest.

    Returns the manifest dict.
    """
    _check_available_memory()
    print(f"Connecting to Snowflake account: {account} ...")
    conn = _connect(account, user, password, warehouse)

    print("Collecting table metadata ...")
    assets = _collect_assets(conn)
    conn.close()
    print(f"\nCollected {len(assets)} table(s).")

    manifest = {
        "resource_type": RESOURCE_TYPE,
        "collected_at": datetime.now(tz=timezone.utc).isoformat(),
        "assets": assets,
    }
    with open(output_file, "w") as fh:
        json.dump(manifest, fh, indent=2)
    print(f"Asset manifest written to {output_file}")

    return manifest


def main() -> None:
    parser = argparse.ArgumentParser(
        description="Collect Snowflake table metadata and write to a manifest file",
    )
    parser.add_argument(
        "--account",
        default=os.environ.get("SNOWFLAKE_ACCOUNT"),
        help="Snowflake account identifier, e.g. xy12345.us-east-1 (env: SNOWFLAKE_ACCOUNT)",  # ← SUBSTITUTE
    )
    parser.add_argument(
        "--user",
        default=os.environ.get("SNOWFLAKE_USER"),
        help="Snowflake username (env: SNOWFLAKE_USER)",  # ← SUBSTITUTE
    )
    parser.add_argument(
        "--password",
        default=os.environ.get("SNOWFLAKE_PASSWORD"),
        help="Snowflake password (env: SNOWFLAKE_PASSWORD)",  # ← SUBSTITUTE
    )
    parser.add_argument(
        "--warehouse",
        default=os.environ.get("SNOWFLAKE_WAREHOUSE"),
        help="Snowflake virtual warehouse (env: SNOWFLAKE_WAREHOUSE)",  # ← SUBSTITUTE
    )
    parser.add_argument(
        "--output-file",
        default="metadata_output.json",
        help="Path to write the output manifest (default: metadata_output.json)",
    )
    args = parser.parse_args()

    missing = [
        name
        for name, val in [
            ("--account", args.account),
            ("--user", args.user),
            ("--password", args.password),
            ("--warehouse", args.warehouse),
        ]
        if not val
    ]
    if missing:
        parser.error(f"Missing required arguments: {', '.join(missing)}")

    collect(
        account=args.account,
        user=args.user,
        password=args.password,
        warehouse=args.warehouse,
        output_file=args.output_file,
    )
    print("Done.")


if __name__ == "__main__":
    main()

#!/usr/bin/env python3
"""
Collect query logs from Snowflake ACCOUNT_USAGE.QUERY_HISTORY — collection only.

Queries a 24-hour window ending 1 hour ago (ACCOUNT_USAGE views have an
approximate 45-minute ingestion latency, so the last hour is intentionally
skipped to avoid incomplete data).  The collected query logs are written to a
JSON manifest file.

Can be run standalone via CLI or imported (use the ``collect()`` function).

Substitution points
-------------------
- SNOWFLAKE_ACCOUNT    (env) / --account    (CLI) : Snowflake account identifier
- SNOWFLAKE_USER       (env) / --user       (CLI) : Snowflake username
- SNOWFLAKE_PASSWORD   (env) / --password   (CLI) : Snowflake password
- SNOWFLAKE_WAREHOUSE  (env) / --warehouse  (CLI) : Snowflake virtual warehouse

Prerequisites
-------------
    pip install snowflake-connector-python

Usage
-----
    python collect_query_logs.py \\
        --account  <SNOWFLAKE_ACCOUNT> \\
        --user     <SNOWFLAKE_USER> \\
        --password <SNOWFLAKE_PASSWORD> \\
        --warehouse <SNOWFLAKE_WAREHOUSE>
"""

import argparse
import json
import os
from datetime import datetime, timezone

import snowflake.connector

# ← SUBSTITUTE: set LOG_TYPE to match your warehouse type (query logs use log_type, not resource_type)
LOG_TYPE = "snowflake"


def _check_available_memory(min_gb: float = 2.0) -> None:
    """Warn if available memory is below the threshold."""
    try:
        if hasattr(os, "sysconf"):  # Linux / macOS
            page_size = os.sysconf("SC_PAGE_SIZE")
            avail_pages = os.sysconf("SC_AVPHYS_PAGES")
            avail_gb = (page_size * avail_pages) / (1024 ** 3)
        else:
            return  # Windows — skip check
    except (ValueError, OSError):
        return
    if avail_gb < min_gb:
        print(
            f"WARNING: Only {avail_gb:.1f} GB of memory available "
            f"(minimum recommended: {min_gb:.1f} GB). "
            f"Consider reducing the lookback window or increasing available memory."
        )

# How many hours to look back from the trailing-edge cutoff
# ← SUBSTITUTE: adjust to match your collection cadence (e.g. 2 for every-2-hours runs)
_WINDOW_HOURS = 25

# Hours to skip at the trailing edge — ACCOUNT_USAGE has ~45-minute latency;
# skipping 1 hour provides a comfortable buffer.
# ← SUBSTITUTE: lower to 0 if you have confirmed real-time access to ACCOUNT_USAGE
_TRAILING_SKIP_HOURS = 1

# Maximum rows to collect per run — increase if your warehouse has higher query volume
# ← SUBSTITUTE: adjust based on your Snowflake query volume
_QUERY_LIMIT = 10000


def _fetch_query_history(conn) -> list[dict]:
    """
    Fetch recent query history from SNOWFLAKE.ACCOUNT_USAGE.QUERY_HISTORY.

    Collection window: [NOW - _WINDOW_HOURS, NOW - _TRAILING_SKIP_HOURS]
    This intentionally excludes the most recent hour to avoid the ACCOUNT_USAGE
    ingestion latency gap.
    """
    cursor = conn.cursor()
    cursor.execute(
        f"""
        SELECT
            QUERY_ID,
            QUERY_TEXT,
            START_TIME,
            END_TIME,
            USER_NAME,
            DATABASE_NAME,
            WAREHOUSE_NAME,
            BYTES_SCANNED,
            ROWS_PRODUCED,
            EXECUTION_STATUS,
            QUERY_TAG,
            ROLE_NAME
        FROM SNOWFLAKE.ACCOUNT_USAGE.QUERY_HISTORY
        WHERE START_TIME >= DATEADD(hour, -{_WINDOW_HOURS}, CURRENT_TIMESTAMP())
          AND START_TIME <  DATEADD(hour, -{_TRAILING_SKIP_HOURS}, CURRENT_TIMESTAMP())
          AND EXECUTION_STATUS = 'SUCCESS'
        ORDER BY START_TIME
        LIMIT {_QUERY_LIMIT}
        """
        # ← SUBSTITUTE: add AND DATABASE_NAME = '<db>' or AND WAREHOUSE_NAME = '<wh>'
        #   to restrict collection to a specific database or warehouse
    )
    columns = [col[0] for col in cursor.description]
    rows = []
    while True:
        chunk = cursor.fetchmany(1000)
        if not chunk:
            break
        rows.extend(dict(zip(columns, row)) for row in chunk)
    cursor.close()
    return rows


def _iso(dt: object) -> str | None:
    if dt is None:
        return None
    return dt.isoformat() if hasattr(dt, "isoformat") else str(dt)


def collect(
    account: str,
    user: str,
    password: str,
    warehouse: str,
    output_file: str = "query_logs_output.json",
) -> dict:
    """
    Connect to Snowflake, collect query logs, and write a JSON manifest.

    Returns the manifest dict.
    """
    _check_available_memory()
    print(f"Connecting to Snowflake account: {account} ...")
    conn = snowflake.connector.connect(
        account=account,
        user=user,
        password=password,
        warehouse=warehouse,
    )

    print(
        f"Fetching QUERY_HISTORY (last {_WINDOW_HOURS}h, excluding final {_TRAILING_SKIP_HOURS}h, "
        f"limit {_QUERY_LIMIT}) ..."
    )
    rows = _fetch_query_history(conn)
    conn.close()
    print(f"  Retrieved {len(rows)} query log row(s).")

    if not rows:
        print("No query log rows found in the specified window.")
        manifest = {
            "log_type": LOG_TYPE,
            "collected_at": datetime.now(tz=timezone.utc).isoformat(),
            "entry_count": 0,
            "window_start": None,
            "window_end": None,
            "queries": [],
        }
        with open(output_file, "w") as fh:
            json.dump(manifest, fh, indent=2, default=str)
        return manifest

    start_times = [r["START_TIME"] for r in rows if r.get("START_TIME") is not None]
    end_times = [r["END_TIME"] for r in rows if r.get("END_TIME") is not None]

    manifest = {
        "log_type": LOG_TYPE,
        "collected_at": datetime.now(tz=timezone.utc).isoformat(),
        "entry_count": len(rows),
        "window_start": _iso(min(start_times)) if start_times else None,
        "window_end": _iso(max(end_times)) if end_times else None,
        "queries": [
            {
                "query_id": r.get("QUERY_ID"),
                "query_text": r.get("QUERY_TEXT") or "",
                "start_time": _iso(r.get("START_TIME")),
                "end_time": _iso(r.get("END_TIME")),
                "user": r.get("USER_NAME"),
                "warehouse": r.get("WAREHOUSE_NAME"),
                "bytes_scanned": r.get("BYTES_SCANNED"),
                "rows_produced": r.get("ROWS_PRODUCED"),
            }
            for r in rows
        ],
    }
    with open(output_file, "w") as fh:
        json.dump(manifest, fh, indent=2, default=str)
    print(f"Query log manifest written to {output_file}")

    return manifest


def main() -> None:
    parser = argparse.ArgumentParser(
        description="Collect Snowflake query logs from ACCOUNT_USAGE and write to a manifest file",
    )
    parser.add_argument(
        "--account",
        default=os.environ.get("SNOWFLAKE_ACCOUNT"),
        help="Snowflake account identifier, e.g. xy12345.us-east-1 (env: SNOWFLAKE_ACCOUNT)",  # ← SUBSTITUTE
    )
    parser.add_argument(
        "--user",
        default=os.environ.get("SNOWFLAKE_USER"),
        help="Snowflake username (env: SNOWFLAKE_USER)",
    )
    parser.add_argument(
        "--password",
        default=os.environ.get("SNOWFLAKE_PASSWORD"),
        help="Snowflake password (env: SNOWFLAKE_PASSWORD)",
    )
    parser.add_argument(
        "--warehouse",
        default=os.environ.get("SNOWFLAKE_WAREHOUSE"),
        help="Snowflake virtual warehouse (env: SNOWFLAKE_WAREHOUSE)",  # ← SUBSTITUTE
    )
    parser.add_argument(
        "--output-file",
        default="query_logs_output.json",
        help="Path to write the output manifest (default: query_logs_output.json)",
    )
    args = parser.parse_args()

    missing = [
        name
        for name, val in [
            ("--account", args.account),
            ("--user", args.user),
            ("--password", args.password),
            ("--warehouse", args.warehouse),
        ]
        if not val
    ]
    if missing:
        parser.error(f"Missing required arguments: {', '.join(missing)}")

    collect(
        account=args.account,
        user=args.user,
        password=args.password,
        warehouse=args.warehouse,
        output_file=args.output_file,
    )
    print("Done.")


if __name__ == "__main__":
    main()

#!/usr/bin/env python3
"""
Push lineage events to Monte Carlo from a JSON manifest — push only.

Reads a manifest file produced by ``collect_lineage.py`` and sends the lineage
events to Monte Carlo using the pycarlo push ingestion API.  Large payloads are
split into batches to stay under the 1 MB compressed limit.

Can be run standalone via CLI or imported (use the ``push()`` function).

Substitution points
-------------------
- MCD_INGEST_ID     (env) / --key-id     (CLI) : Monte Carlo ingestion key ID
- MCD_INGEST_TOKEN  (env) / --key-token  (CLI) : Monte Carlo ingestion key token
- MCD_RESOURCE_UUID     (env) / --resource-uuid (CLI) : MC resource UUID for this connection

Prerequisites
-------------
    pip install pycarlo

Usage
-----
    python push_lineage.py \\
        --key-id  <MCD_INGEST_ID> \\
        --key-token <MCD_INGEST_TOKEN> \\
        --resource-uuid <MCD_RESOURCE_UUID> \\
        --input-file lineage_output.json
"""

from __future__ import annotations

import argparse
import json
import os
from concurrent.futures import ThreadPoolExecutor, as_completed
from datetime import datetime, timezone

from pycarlo.core import Client, Session
from pycarlo.features.ingestion import IngestionService
from pycarlo.features.ingestion.models import (
    ColumnLineageField,
    ColumnLineageSourceField,
    LineageAssetRef,
    LineageEvent,
)

# ← SUBSTITUTE: set RESOURCE_TYPE to match your Monte Carlo connection type
RESOURCE_TYPE = "snowflake"

# Maximum events per batch — conservative default to keep compressed payload under 1 MB
# ← SUBSTITUTE: tune based on average edge complexity (number of sources, column mappings)
_BATCH_SIZE = 500


def _build_table_lineage_events(edges: list[dict]) -> list[LineageEvent]:
    """Build table-level LineageEvent objects from manifest edge dicts."""
    events = []
    for edge in edges:
        dest = edge["destination"]
        sources = edge.get("sources", [])
        if not sources:
            continue
        events.append(
            LineageEvent(
                destination=LineageAssetRef(
                    type="TABLE",
                    name=dest["table"],
                    database=dest["database"],
                    schema=dest["schema"],
                ),
                sources=[
                    LineageAssetRef(
                        type="TABLE",
                        name=s["table"],
                        database=s["database"],
                        schema=s["schema"],
                    )
                    for s in sources
                ],
            )
        )
    return events


def _build_column_lineage_events(edges: list[dict]) -> list[LineageEvent]:
    """Build column-level LineageEvent objects from manifest edge dicts."""
    events = []
    for edge in edges:
        dest = edge["destination"]
        sources = edge.get("sources", [])
        col_mappings = edge.get("col_mappings", [])
        if not sources:
            continue

        dest_asset_id = f"{dest['database']}__{dest['schema']}__{dest['table']}"
        source_asset_ids = {
            (s["database"], s["schema"], s["table"]): f"{s['database']}__{s['schema']}__{s['table']}"
            for s in sources
        }

        col_fields: dict[str, ColumnLineageField] = {}
        for mapping in col_mappings:
            dest_col = mapping["dest_col"]
            src_table = mapping["src_table"]
            src_col = mapping["src_col"]
            # Match src_table to the first source with that table name
            match = next(
                (s for s in sources if s["table"] == src_table),
                sources[0] if sources else None,
            )
            if not match:
                continue
            src_aid = source_asset_ids[(match["database"], match["schema"], match["table"])]
            if dest_col not in col_fields:
                col_fields[dest_col] = ColumnLineageField(name=dest_col, source_fields=[])
            col_fields[dest_col].source_fields.append(
                ColumnLineageSourceField(asset_id=src_aid, field_name=src_col)
            )

        events.append(
            LineageEvent(
                destination=LineageAssetRef(
                    type="TABLE",
                    name=dest["table"],
                    database=dest["database"],
                    schema=dest["schema"],
                    asset_id=dest_asset_id,
                ),
                sources=[
                    LineageAssetRef(
                        type="TABLE",
                        name=s["table"],
                        database=s["database"],
                        schema=s["schema"],
                        asset_id=source_asset_ids[(s["database"], s["schema"], s["table"])],
                    )
                    for s in sources
                ],
                fields=list(col_fields.values()) if col_fields else None,
            )
        )
    return events


def push(
    input_file: str,
    resource_uuid: str,
    key_id: str,
    key_token: str,
    batch_size: int = _BATCH_SIZE,
    output_file: str = "lineage_push_result.json",
) -> dict:
    """
    Read a lineage manifest and push events to Monte Carlo in batches.

    Returns a result dict with invocation IDs for each batch.
    """
    with open(input_file) as fh:
        manifest = json.load(fh)

    edges = manifest.get("edges", [])
    resource_type = manifest.get("resource_type", RESOURCE_TYPE)
    column_lineage = manifest.get("column_lineage", False)

    if column_lineage:
        events = _build_column_lineage_events(edges)
        label = "column-level"
    else:
        events = _build_table_lineage_events(edges)
        label = "table-level"

    print(f"Loaded {len(events)} {label} lineage event(s) from {input_file}")

    if not events:
        print("No lineage events to push.")
        push_result = {
            "resource_uuid": resource_uuid,
            "resource_type": resource_type,
            "invocation_ids": [],
            "pushed_at": datetime.now(tz=timezone.utc).isoformat(),
            "total_events": 0,
            "batch_count": 0,
            "batch_size": batch_size,
        }
        with open(output_file, "w") as fh:
            json.dump(push_result, fh, indent=2)
        return push_result

    # Split into batches
    batches = []
    for i in range(0, len(events), batch_size):
        batches.append(events[i : i + batch_size])
    total_batches = len(batches)

    def _push_batch(batch: list, batch_num: int) -> str | None:
        """Push a single batch using a dedicated Session (thread-safe)."""
        print(f"  Pushing batch {batch_num}/{total_batches} ({len(batch)} events) ...")
        client = Client(session=Session(mcd_id=key_id, mcd_token=key_token, scope="Ingestion"))
        service = IngestionService(mc_client=client)
        result = service.send_lineage(
            resource_uuid=resource_uuid,
            resource_type=resource_type,
            events=batch,
        )
        invocation_id = service.extract_invocation_id(result)
        if invocation_id:
            print(f"    Batch {batch_num}: invocation_id={invocation_id}")
        return invocation_id

    # Push batches in parallel (each thread gets its own pycarlo Session)
    max_workers = min(4, total_batches)
    invocation_ids: list[str | None] = [None] * total_batches

    with ThreadPoolExecutor(max_workers=max_workers) as pool:
        futures = {
            pool.submit(_push_batch, batch, i + 1): i
            for i, batch in enumerate(batches)
        }
        for future in as_completed(futures):
            idx = futures[future]
            try:
                invocation_ids[idx] = future.result()
            except Exception as exc:
                print(f"    ERROR pushing batch {idx + 1}: {exc}")
                raise

    print(f"  All {total_batches} batches pushed ({max_workers} workers)")

    push_result = {
        "resource_uuid": resource_uuid,
        "resource_type": resource_type,
        "invocation_ids": invocation_ids,
        "pushed_at": datetime.now(tz=timezone.utc).isoformat(),
        "total_events": len(events),
        "batch_count": total_batches,
        "batch_size": batch_size,
        "edges": edges,  # preserve for downstream validation
    }
    with open(output_file, "w") as fh:
        json.dump(push_result, fh, indent=2)
    print(f"Push result written to {output_file}")

    return push_result


def main() -> None:
    parser = argparse.ArgumentParser(
        description="Push Snowflake lineage from a manifest to Monte Carlo",
    )
    parser.add_argument(
        "--key-id",
        default=os.environ.get("MCD_INGEST_ID"),
        help="Monte Carlo ingestion key ID (env: MCD_INGEST_ID)",
    )
    parser.add_argument(
        "--key-token",
        default=os.environ.get("MCD_INGEST_TOKEN"),
        help="Monte Carlo ingestion key token (env: MCD_INGEST_TOKEN)",
    )
    parser.add_argument(
        "--resource-uuid",
        default=os.environ.get("MCD_RESOURCE_UUID"),
        help="Monte Carlo resource UUID for this Snowflake connection (env: MCD_RESOURCE_UUID)",
    )
    parser.add_argument(
        "--input-file",
        default="lineage_output.json",
        help="Path to the collect manifest to read (default: lineage_output.json)",
    )
    parser.add_argument(
        "--output-file",
        default="lineage_push_result.json",
        help="Path to write the push result (default: lineage_push_result.json)",
    )
    parser.add_argument(
        "--batch-size",
        type=int,
        default=_BATCH_SIZE,
        help=f"Max events per push batch (default: {_BATCH_SIZE})",
    )
    args = parser.parse_args()

    missing = [
        name
        for name, val in [
            ("--key-id", args.key_id),
            ("--key-token", args.key_token),
            ("--resource-uuid", args.resource_uuid),
        ]
        if not val
    ]
    if missing:
        parser.error(f"Missing required arguments: {', '.join(missing)}")

    push(
        input_file=args.input_file,
        resource_uuid=args.resource_uuid,
        key_id=args.key_id,
        key_token=args.key_token,
        batch_size=args.batch_size,
        output_file=args.output_file,
    )
    print("Done.")


if __name__ == "__main__":
    main()

#!/usr/bin/env python3
"""
Push table metadata to Monte Carlo from a JSON manifest — push only.

Reads a manifest file produced by ``collect_metadata.py`` and sends the assets
to Monte Carlo as RelationalAsset events using the pycarlo push ingestion API.
Large payloads are split into batches to stay under the 1 MB compressed limit.

Can be run standalone via CLI or imported (use the ``push()`` function).

Substitution points
-------------------
- MCD_INGEST_ID     (env) / --key-id     (CLI) : Monte Carlo ingestion key ID
- MCD_INGEST_TOKEN  (env) / --key-token  (CLI) : Monte Carlo ingestion key token
- MCD_RESOURCE_UUID     (env) / --resource-uuid (CLI) : MC resource UUID for this connection

Prerequisites
-------------
    pip install pycarlo

Usage
-----
    python push_metadata.py \\
        --key-id  <MCD_INGEST_ID> \\
        --key-token <MCD_INGEST_TOKEN> \\
        --resource-uuid <MCD_RESOURCE_UUID> \\
        --input-file metadata_output.json
"""

import argparse
import json
import os
from concurrent.futures import ThreadPoolExecutor, as_completed
from datetime import datetime, timezone

from pycarlo.core import Client, Session
from pycarlo.features.ingestion import IngestionService
from pycarlo.features.ingestion.models import (
    AssetField,
    AssetFreshness,
    AssetMetadata,
    AssetVolume,
    RelationalAsset,
)

# ← SUBSTITUTE: set RESOURCE_TYPE to match your Monte Carlo connection type
RESOURCE_TYPE = "snowflake"

# Maximum assets per batch — conservative default to keep compressed payload under 1 MB
# ← SUBSTITUTE: tune based on average asset size (fields per table, description length, etc.)
_BATCH_SIZE = 500


def _asset_from_dict(d: dict) -> RelationalAsset:
    """Reconstruct a RelationalAsset from a manifest dict entry."""
    fields = [
        AssetField(
            name=f["name"],
            type=f.get("type"),
            description=f.get("description"),
        )
        for f in d.get("fields", [])
    ]

    volume = None
    if d.get("volume"):
        volume = AssetVolume(
            row_count=d["volume"].get("row_count"),
            byte_count=d["volume"].get("byte_count"),
        )

    freshness = None
    if d.get("freshness"):
        freshness = AssetFreshness(
            last_update_time=d["freshness"].get("last_update_time"),
        )

    return RelationalAsset(
        type=d.get("type", "TABLE"),
        metadata=AssetMetadata(
            name=d["name"],
            database=d["database"],
            schema=d["schema"],
            description=d.get("description"),
        ),
        fields=fields,
        volume=volume,
        freshness=freshness,
    )


def push(
    input_file: str,
    resource_uuid: str,
    key_id: str,
    key_token: str,
    batch_size: int = _BATCH_SIZE,
    output_file: str = "metadata_push_result.json",
) -> dict:
    """
    Read a metadata manifest and push assets to Monte Carlo in batches.

    Returns a result dict with invocation IDs for each batch.
    """
    with open(input_file) as fh:
        manifest = json.load(fh)

    asset_dicts = manifest.get("assets", [])
    resource_type = manifest.get("resource_type", RESOURCE_TYPE)
    assets = [_asset_from_dict(d) for d in asset_dicts]
    print(f"Loaded {len(assets)} asset(s) from {input_file}")

    # Split into batches
    batches = []
    for i in range(0, max(len(assets), 1), batch_size):
        batches.append(assets[i : i + batch_size])
    total_batches = len(batches)

    def _push_batch(batch: list, batch_num: int) -> str | None:
        """Push a single batch using a dedicated Session (thread-safe)."""
        client = Client(session=Session(mcd_id=key_id, mcd_token=key_token, scope="Ingestion"))
        service = IngestionService(mc_client=client)
        result = service.send_metadata(
            resource_uuid=resource_uuid,
            resource_type=resource_type,
            events=batch,
        )
        invocation_id = service.extract_invocation_id(result)
        print(f"  Pushed batch {batch_num}/{total_batches} ({len(batch)} assets) — invocation_id={invocation_id}")
        return invocation_id

    # Push batches in parallel (each thread gets its own pycarlo Session)
    max_workers = min(4, total_batches)
    invocation_ids: list[str | None] = [None] * total_batches

    with ThreadPoolExecutor(max_workers=max_workers) as pool:
        futures = {
            pool.submit(_push_batch, batch, i + 1): i
            for i, batch in enumerate(batches)
        }
        for future in as_completed(futures):
            idx = futures[future]
            try:
                invocation_ids[idx] = future.result()
            except Exception as exc:
                print(f"    ERROR pushing batch {idx + 1}: {exc}")
                raise

    print(f"  All {total_batches} batches pushed ({max_workers} workers)")

    push_result = {
        "resource_uuid": resource_uuid,
        "resource_type": resource_type,
        "invocation_ids": invocation_ids,
        "pushed_at": datetime.now(tz=timezone.utc).isoformat(),
        "total_assets": len(assets),
        "batch_count": total_batches,
        "batch_size": batch_size,
    }
    with open(output_file, "w") as fh:
        json.dump(push_result, fh, indent=2)
    print(f"Push result written to {output_file}")

    return push_result


def main() -> None:
    parser = argparse.ArgumentParser(
        description="Push Snowflake table metadata from a manifest to Monte Carlo",
    )
    parser.add_argument(
        "--key-id",
        default=os.environ.get("MCD_INGEST_ID"),
        help="Monte Carlo ingestion key ID (env: MCD_INGEST_ID)",
    )
    parser.add_argument(
        "--key-token",
        default=os.environ.get("MCD_INGEST_TOKEN"),
        help="Monte Carlo ingestion key token (env: MCD_INGEST_TOKEN)",
    )
    parser.add_argument(
        "--resource-uuid",
        default=os.environ.get("MCD_RESOURCE_UUID"),
        help="Monte Carlo resource UUID for this Snowflake connection (env: MCD_RESOURCE_UUID)",
    )
    parser.add_argument(
        "--input-file",
        default="metadata_output.json",
        help="Path to the collect manifest to read (default: metadata_output.json)",
    )
    parser.add_argument(
        "--output-file",
        default="metadata_push_result.json",
        help="Path to write the push result (default: metadata_push_result.json)",
    )
    parser.add_argument(
        "--batch-size",
        type=int,
        default=_BATCH_SIZE,
        help=f"Max assets per push batch (default: {_BATCH_SIZE})",
    )
    args = parser.parse_args()

    missing = [
        name
        for name, val in [
            ("--key-id", args.key_id),
            ("--key-token", args.key_token),
            ("--resource-uuid", args.resource_uuid),
        ]
        if not val
    ]
    if missing:
        parser.error(f"Missing required arguments: {', '.join(missing)}")

    push(
        input_file=args.input_file,
        resource_uuid=args.resource_uuid,
        key_id=args.key_id,
        key_token=args.key_token,
        batch_size=args.batch_size,
        output_file=args.output_file,
    )
    print("Done.")


if __name__ == "__main__":
    main()

#!/usr/bin/env python3
"""
Push query logs to Monte Carlo from a JSON manifest — push only.

Reads a manifest file produced by ``collect_query_logs.py`` and sends the query
log entries to Monte Carlo using the pycarlo push ingestion API.  Large payloads
are split into batches to stay under the 1 MB compressed limit.

Can be run standalone via CLI or imported (use the ``push()`` function).

Substitution points
-------------------
- MCD_INGEST_ID     (env) / --key-id     (CLI) : Monte Carlo ingestion key ID
- MCD_INGEST_TOKEN  (env) / --key-token  (CLI) : Monte Carlo ingestion key token
- MCD_RESOURCE_UUID     (env) / --resource-uuid (CLI) : MC resource UUID for this connection

Prerequisites
-------------
    pip install pycarlo

Usage
-----
    python push_query_logs.py \\
        --key-id  <MCD_INGEST_ID> \\
        --key-token <MCD_INGEST_TOKEN> \\
        --resource-uuid <MCD_RESOURCE_UUID> \\
        --input-file query_logs_output.json
"""

import argparse
import json
import os
from concurrent.futures import ThreadPoolExecutor, as_completed
from datetime import datetime, timezone

from dateutil.parser import isoparse
from pycarlo.core import Client, Session
from pycarlo.features.ingestion import IngestionService
from pycarlo.features.ingestion.models import QueryLogEntry

# ← SUBSTITUTE: set LOG_TYPE to match your warehouse type (query logs use log_type, not resource_type)
LOG_TYPE = "snowflake"

# Maximum entries per batch — conservative default to keep compressed payload under 1 MB.
# Query logs include full SQL text — keep batches small to stay under the 1 MB
# compressed payload limit.  50 entries can trigger 413 on active warehouses.
# ← SUBSTITUTE: tune based on average query length
_BATCH_SIZE = 100

# Truncate query_text longer than this to prevent 413 errors.
# Some SQL statements (e.g., generated by BI tools) can be 100KB+ and blow up
# compressed payloads even at small batch sizes.
_MAX_QUERY_TEXT_LEN = 10_000


def _build_query_log_entries(queries: list[dict]) -> list[QueryLogEntry]:
    """Convert manifest query dicts into QueryLogEntry objects."""
    entries = []
    truncated = 0
    for q in queries:
        start_time = q.get("start_time")
        end_time = q.get("end_time")
        query_text = q.get("query_text") or ""
        query_id = q.get("query_id")
        user_name = q.get("user")
        warehouse_name = q.get("warehouse")
        bytes_scanned = q.get("bytes_scanned")
        rows_produced = q.get("rows_produced")

        # Truncate very long SQL to prevent 413 Request Too Large
        if len(query_text) > _MAX_QUERY_TEXT_LEN:
            query_text = query_text[:_MAX_QUERY_TEXT_LEN] + "... [TRUNCATED]"
            truncated += 1

        extra = {}
        if warehouse_name is not None:
            extra["warehouse_name"] = warehouse_name
        if bytes_scanned is not None:
            extra["bytes_scanned"] = int(bytes_scanned)

        entries.append(
            QueryLogEntry(
                start_time=isoparse(start_time) if start_time else None,
                end_time=isoparse(end_time) if end_time else None,
                query_text=query_text,
                query_id=query_id,
                user=user_name,
                returned_rows=int(rows_produced) if rows_produced is not None else None,
                extra=extra or None,
            )
        )
    if truncated:
        print(f"  Truncated {truncated} query text(s) exceeding {_MAX_QUERY_TEXT_LEN} chars")
    return entries


def push(
    input_file: str,
    resource_uuid: str,
    key_id: str,
    key_token: str,
    batch_size: int = _BATCH_SIZE,
    output_file: str = "query_logs_push_result.json",
) -> dict:
    """
    Read a query log manifest and push entries to Monte Carlo in batches.

    Returns a result dict with invocation IDs for each batch.
    """
    with open(input_file) as fh:
        manifest = json.load(fh)

    queries = manifest.get("queries", [])
    log_type = manifest.get("log_type", LOG_TYPE)
    entries = _build_query_log_entries(queries)
    print(f"Loaded {len(entries)} query log entry/entries from {input_file}")

    if not entries:
        print("No query log entries to push.")
        push_result = {
            "resource_uuid": resource_uuid,
            "log_type": log_type,
            "invocation_ids": [],
            "pushed_at": datetime.now(tz=timezone.utc).isoformat(),
            "total_entries": 0,
            "batch_count": 0,
            "batch_size": batch_size,
        }
        with open(output_file, "w") as fh:
            json.dump(push_result, fh, indent=2)
        return push_result

    # Split into batches
    batches = []
    for i in range(0, len(entries), batch_size):
        batches.append(entries[i : i + batch_size])
    total_batches = len(batches)

    def _push_batch(batch: list, batch_num: int) -> str | None:
        """Push a single batch using a dedicated Session (thread-safe)."""
        client = Client(session=Session(mcd_id=key_id, mcd_token=key_token, scope="Ingestion"))
        service = IngestionService(mc_client=client)
        result = service.send_query_logs(
            resource_uuid=resource_uuid,
            log_type=log_type,
            events=batch,
        )
        invocation_id = service.extract_invocation_id(result)
        print(f"  Pushed batch {batch_num}/{total_batches} ({len(batch)} entries) — invocation_id={invocation_id}")
        return invocation_id

    # Push batches in parallel (each thread gets its own pycarlo Session)
    max_workers = min(4, total_batches)
    invocation_ids: list[str | None] = [None] * total_batches

    with ThreadPoolExecutor(max_workers=max_workers) as pool:
        futures = {
            pool.submit(_push_batch, batch, i + 1): i
            for i, batch in enumerate(batches)
        }
        for future in as_completed(futures):
            idx = futures[future]
            try:
                invocation_ids[idx] = future.result()
            except Exception as exc:
                print(f"    ERROR pushing batch {idx + 1}: {exc}")
                raise

    print(f"  All {total_batches} batches pushed ({max_workers} workers)")

    push_result = {
        "resource_uuid": resource_uuid,
        "log_type": log_type,
        "invocation_ids": invocation_ids,
        "pushed_at": datetime.now(tz=timezone.utc).isoformat(),
        "total_entries": len(entries),
        "batch_count": total_batches,
        "batch_size": batch_size,
    }
    with open(output_file, "w") as fh:
        json.dump(push_result, fh, indent=2)
    print(f"Push result written to {output_file}")

    return push_result


def main() -> None:
    parser = argparse.ArgumentParser(
        description="Push Snowflake query logs from a manifest to Monte Carlo",
    )
    parser.add_argument(
        "--key-id",
        default=os.environ.get("MCD_INGEST_ID"),
        help="Monte Carlo ingestion key ID (env: MCD_INGEST_ID)",
    )
    parser.add_argument(
        "--key-token",
        default=os.environ.get("MCD_INGEST_TOKEN"),
        help="Monte Carlo ingestion key token (env: MCD_INGEST_TOKEN)",
    )
    parser.add_argument(
        "--resource-uuid",
        default=os.environ.get("MCD_RESOURCE_UUID"),
        help="Monte Carlo resource UUID for this Snowflake connection (env: MCD_RESOURCE_UUID)",
    )
    parser.add_argument(
        "--input-file",
        default="query_logs_output.json",
        help="Path to the collect manifest to read (default: query_logs_output.json)",
    )
    parser.add_argument(
        "--output-file",
        default="query_logs_push_result.json",
        help="Path to write the push result (default: query_logs_push_result.json)",
    )
    parser.add_argument(
        "--batch-size",
        type=int,
        default=_BATCH_SIZE,
        help=f"Max entries per push batch (default: {_BATCH_SIZE})",
    )
    args = parser.parse_args()

    missing = [
        name
        for name, val in [
            ("--key-id", args.key_id),
            ("--key-token", args.key_token),
            ("--resource-uuid", args.resource_uuid),
        ]
        if not val
    ]
    if missing:
        parser.error(f"Missing required arguments: {', '.join(missing)}")

    push(
        input_file=args.input_file,
        resource_uuid=args.resource_uuid,
        key_id=args.key_id,
        key_token=args.key_token,
        batch_size=args.batch_size,
        output_file=args.output_file,
    )
    print("Done.")


if __name__ == "__main__":
    main()