pip로 설치 (가장 일반적)

# pip로 설치 (가장 일반적)
pip install databricks-connect==16.2.*

# 버전은 사용할 Databricks Runtime과 일치시켜야 합니다
# 예: DBR 16.2를 사용한다면 databricks-connect==16.2.*

# 설치 확인
python -c "from databricks.connect import DatabricksSession; print('OK')"

# Databricks CLI로 OAuth 인증 설정
databricks configure --profile CONNECT

# 입력:
# Databricks Host: https://dbc-abc123.cloud.databricks.com
# (브라우저에서 OAuth 인증 완료)

# Python 코드에서 프로필 사용
from databricks.connect import DatabricksSession

spark = DatabricksSession.builder.profile("CONNECT").getOrCreate()

from databricks.connect import DatabricksSession

spark = DatabricksSession.builder \
    .host("https://dbc-abc123.cloud.databricks.com") \
    .token("dapi...") \
    .clusterId("0123-456789-abcdefgh") \
    .getOrCreate()

export DATABRICKS_HOST="https://dbc-abc123.cloud.databricks.com"
export DATABRICKS_TOKEN="dapi..."
export DATABRICKS_CLUSTER_ID="0123-456789-abcdefgh"

# 환경 변수가 설정되면 별도 설정 없이 세션 생성 가능
from databricks.connect import DatabricksSession

spark = DatabricksSession.builder.getOrCreate()

VS Code 설정 순서:
1. "Databricks" Extension 설치 (Marketplace)
2. Extension에서 Workspace 연결 설정
3. 클러스터 선택
4. Python 인터프리터를 databricks-connect가 설치된 환경으로 설정
5. 코드 작성 후 F5로 실행/디버깅

PyCharm 설정 순서:
1. Project Interpreter에서 databricks-connect 설치된 venv 선택
2. Run/Debug Configuration 생성
3. 환경 변수 또는 코드 내에서 인증 설정
4. 코드 실행 시 자동으로 원격 클러스터에서 처리

from databricks.connect import DatabricksSession

# 세션 생성 (원격 클러스터에 연결)
spark = DatabricksSession.builder.profile("CONNECT").getOrCreate()

# Unity Catalog 테이블 읽기
df = spark.read.table("catalog.schema.my_table")

# 로컬에서 작성한 변환 로직이 원격에서 실행됩니다
result = (
    df
    .filter(df.status == "active")
    .groupBy("region")
    .agg({"amount": "sum", "customer_id": "count"})
    .orderBy("sum(amount)", ascending=False)
)

# 결과를 로컬로 가져오기
result.show()

# 로컬 데이터로 DataFrame 생성 후 원격에서 처리
from pyspark.sql import Row

local_data = [
    Row(name="김철수", department="엔지니어링", salary=80000),
    Row(name="이영희", department="마케팅", salary=75000),
    Row(name="박민수", department="엔지니어링", salary=85000),
]

df = spark.createDataFrame(local_data)

# 원격 클러스터에서 집계 실행
dept_stats = (
    df.groupBy("department")
    .agg(
        {"salary": "avg", "*": "count"}
    )
)
dept_stats.show()

# Unity Catalog 테이블에 SQL 쿼리 실행
result = spark.sql("""
    SELECT
        region,
        COUNT(*) AS order_count,
        SUM(amount) AS total_revenue
    FROM catalog.gold.orders
    WHERE order_date >= '2025-01-01'
    GROUP BY region
    ORDER BY total_revenue DESC
""")

# pandas DataFrame으로 변환 (로컬로 가져오기)
pandas_df = result.toPandas()
print(pandas_df)

# pytest와 함께 사용
import pytest
from databricks.connect import DatabricksSession

@pytest.fixture(scope="session")
def spark():
    """테스트 세션용 Spark 세션"""
    return DatabricksSession.builder.profile("CONNECT").getOrCreate()

def test_data_quality(spark):
    """데이터 품질 검증 테스트"""
    df = spark.read.table("catalog.gold.customers")

    # Null 체크
    null_count = df.filter(df.customer_id.isNull()).count()
    assert null_count == 0, f"customer_id에 {null_count}개의 Null 값이 있습니다"

    # 유니크 체크
    total = df.count()
    distinct = df.select("customer_id").distinct().count()
    assert total == distinct, "customer_id에 중복이 있습니다"

def test_transformation_logic(spark):
    """변환 로직 단위 테스트"""
    from my_project.transforms import calculate_revenue

    input_df = spark.createDataFrame([
        {"product_id": "A", "quantity": 10, "price": 100.0},
        {"product_id": "B", "quantity": 5, "price": 200.0},
    ])

    result = calculate_revenue(input_df)
    rows = result.collect()

    assert rows[0]["revenue"] == 1000.0
    assert rows[1]["revenue"] == 1000.0

from databricks.connect import DatabricksSession

# Serverless에 연결 (cluster_id 대신 serverless 사용)
spark = DatabricksSession.builder \
    .host("https://dbc-abc123.cloud.databricks.com") \
    .token("dapi...") \
    .serverless(True) \
    .getOrCreate()

# 이후 사용법은 동일합니다
df = spark.read.table("catalog.schema.my_table")
df.show()

# RDD 대신 DataFrame API 사용
# ❌ 지원하지 않음
# rdd = spark.sparkContext.parallelize([1, 2, 3])

# ✅ DataFrame API 사용
from pyspark.sql import Row
df = spark.createDataFrame([Row(value=1), Row(value=2), Row(value=3)])

# DML(MERGE 등)은 SQL로 실행
# ❌ DataFrame Writer의 merge는 제한적
# ✅ SQL 사용
spark.sql("""
    MERGE INTO catalog.schema.target t
    USING catalog.schema.source s
    ON t.id = s.id
    WHEN MATCHED THEN UPDATE SET *
    WHEN NOT MATCHED THEN INSERT *
""")

# 연결 상태 확인
from databricks.connect import DatabricksSession

spark = DatabricksSession.builder.profile("CONNECT").getOrCreate()

# 간단한 연산으로 연결 테스트
spark.range(10).count()
print("연결 성공!")