Databricks.txt

========================================================================================
스토리지 연결
========================================================================================

storage_key = ""

spark.conf.set("fs.azure.account.key.wjhjsbs.blob.core.windows.net",storage_key)
spark.conf.set("fs.azure.account.key.wjhjsbs.dfs.core.windows.net",storage_key)
spark.conf.set("spark.databricks.delta.schema.autoMerge.enabled", "true")
spark.conf.set("spark.sql.session.timeZone", "Asia/Seoul")

========================================================================================
 스트리밍
========================================================================================

from pyspark.sql.functions import (
    input_file_name,
    collect_list,
    concat_ws,
    current_timestamp,
    col
)

from pyspark.sql.types import *

schema = StructType([
    StructField("id", StringType(), True),
    StructField("input", StringType(), True),
    StructField("output", StringType(), True),
    StructField("model_name", StringType(), True),
    StructField("api_version", StringType(), True),
    StructField("prompt_tokens", IntegerType(), True),
    StructField("completion_tokens", IntegerType(), True),
    StructField("total_tokens", IntegerType(), True),
    StructField("estimated_cost", DoubleType(), True),
    StructField("timestamp", StringType(), True),
])

df_raw = (
    spark.readStream
    .format("cloudFiles")
    .option("cloudFiles.format", "json")
    .option("cloudFiles.schemaLocation", "abfss://stream@wjhjsbs.dfs.core.windows.net/bronze/axtft/_schema")
    .option("multiLine", "true")
    .schema(schema)
    .load("abfss://axtft@wjhjsbs.dfs.core.windows.net")
    .withColumn("file_path", input_file_name())
)

def write_batch(batch_df, batch_id):

    # 빈 batch 방어
    if batch_df.isEmpty():
        print(f"[BATCH {batch_id}] Empty batch - skipped")
        return

    # 컬럼 존재 확인 (디버깅용)
    print(f"[BATCH {batch_id}] Columns: {batch_df.columns}")

    df_final = (
        batch_df
            .withColumn("ingest_ts", current_timestamp())
    )

    (
        df_final
            .write
            .format("delta")
            .mode("append")
            .option("mergeSchema", "true")
            .saveAsTable("bronze.axtft")
    )
    
query = (
    df_raw.writeStream
        .foreachBatch(write_batch)
        .option(
            "checkpointLocation",
            "abfss://stream@wjhjsbs.dfs.core.windows.net/bronze/axtft/_checkpoint"
        )
        .trigger(once=True)
        .start()
)

========================================================================================
 REST API
========================================================================================

curl -sS -X POST "https://adb-7405615198282020.0.azuredatabricks.net/api/2.0/sql/statements/" \
  -H "Authorization: Bearer abc" \
  -H "Content-Type: application/json" \
  -d '{
    "warehouse_id": "56fa0508658093ef",
    "statement": "SELECT * FROM samples.accuweather.forecast_daily_calendar_imperial LIMIT 50 OFFSET 100",
    "disposition": "INLINE",
    "format": "JSON_ARRAY",
    "wait_timeout": "5s"
  }'


curl -sS -X GET "https://adb-7405615198282020.0.azuredatabricks.net/api/2.0/sql/statements/01f113b6-086c-1bad-9b30-6477910b0af7"   -H "Authorization: Bearer abc"

========================================================================================
AWS Databricks <-> Entra ID
========================================================================================


from fastapi import FastAPI, Header, HTTPException
import os
import requests

app = FastAPI()

DATABRICKS_HOST = os.environ["DATABRICKS_HOST"]
WAREHOUSE_ID = os.environ["WAREHOUSE_ID"]
DATABRICKS_CLIENT_ID = os.getenv("DATABRICKS_CLIENT_ID")  # 필요하면 설정

def token_exchange(entra_jwt: str) -> str:
    url = f"{DATABRICKS_HOST}/oidc/v1/token"
    data = {
        "subject_token": entra_jwt,
        "subject_token_type": "urn:ietf:params:oauth:token-type:jwt",
        "grant_type": "urn:ietf:params:oauth:grant-type:token-exchange",
        "scope": "all-apis",
    }
    if DATABRICKS_CLIENT_ID:
        data["client_id"] = DATABRICKS_CLIENT_ID

    r = requests.post(url, data=data, timeout=15)
    if r.status_code != 200:
        raise HTTPException(status_code=401, detail={"token_exchange_failed": r.text})
    return r.json()["access_token"]

def sql_statements(dbx_token: str, sql: str) -> dict:
    url = f"{DATABRICKS_HOST}/api/2.0/sql/statements/"
    headers = {"Authorization": f"Bearer {dbx_token}"}
    payload = {"warehouse_id": WAREHOUSE_ID, "statement": sql, "wait_timeout": "10s"}
    r = requests.post(url, json=payload, headers=headers, timeout=30)
    r.raise_for_status()
    return r.json()

@app.post("/query")
def query(statement: str, authorization: str = Header(...)):
    if not authorization.lower().startswith("bearer "):
        raise HTTPException(status_code=401, detail="Missing Bearer token")
    entra_token = authorization.split(" ", 1)[1].strip()

    dbx_token = token_exchange(entra_token)
    return sql_statements(dbx_token, statement)


========================================================================================
컬럼 마스킹
========================================================================================
-- 1) 함수 생성
CREATE OR REPLACE FUNCTION main.security.mask_unit_price(price DOUBLE)
RETURN CASE
  WHEN is_account_group_member('price_access') THEN price
  ELSE NULL
END;

-- 2) 마스크 적용
ALTER TABLE main.demo.sales_sample
ALTER COLUMN unit_price
SET MASK main.security.mask_unit_price;

-- 3) 확인
DESCRIBE EXTENDED main.demo.sales_sample;

-- 4) 테스트
SELECT order_id, product, unit_price
FROM main.demo.sales_sample;


========================================================================================
행 필터
========================================================================================

CREATE OR REPLACE FUNCTION region_filter(region STRING)
RETURN
  CASE
    WHEN is_account_group_member('korea_team') THEN region = 'Korea'
    WHEN is_account_group_member('usa_team') THEN region = 'USA'
    ELSE FALSE
  END;

ALTER TABLE sales_sample
SET ROW FILTER region_filter ON (region);