diff --git a/.env.example b/.env.example new file mode 100644 index 0000000..cdfd8ca --- /dev/null +++ b/.env.example @@ -0,0 +1,34 @@ +# DR Sync Configuration — Environment Variables +# Copy this file to .env and fill in your values. +# These override common.py when DR_SYNC_SOURCE_HOST is set. + +# Cloud provider: aws, azure, or gcp +DR_SYNC_CLOUD_TYPE=azure + +# Source (primary) workspace +DR_SYNC_SOURCE_HOST=https:// +DR_SYNC_SOURCE_TOKEN= + +# Target (secondary) workspace +DR_SYNC_TARGET_HOST=https:// +DR_SYNC_TARGET_TOKEN= + +# Catalogs to replicate (comma-separated) +DR_SYNC_CATALOGS_TO_COPY=my-catalog1,my-catalog2 + +# Mapping file paths +DR_SYNC_CRED_MAPPING_FILE=data/azure_cred_mapping.csv +DR_SYNC_LOC_MAPPING_FILE=data/ext_location_mapping.csv +DR_SYNC_CATALOG_MAPPING_FILE=data/catalog_mapping.csv +DR_SYNC_SCHEMA_MAPPING_FILE=data/schema_mapping.csv + +# Execution settings +DR_SYNC_LANDING_ZONE_URL=path/to/storage/ +DR_SYNC_NUM_EXEC=4 +DR_SYNC_WAREHOUSE_SIZE=Small +DR_SYNC_RESPONSE_BACKOFF=0.5 +DR_SYNC_METASTORE_ID= +DR_SYNC_MANIFEST_NAME=manifest + +# Runtime flags +DR_SYNC_DRY_RUN=false diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..9603e23 --- /dev/null +++ b/.gitignore @@ -0,0 +1,61 @@ +# Agents +.claude/ +.cursor/ +.copilot/ +.github/copilot-instructions.md +.aider* +.codeium/ +.tabnine/ +.sourcegraph/ +.cody/ +.victor/ +CLAUDE.md + +# Python +__pycache__/ +*.py[cod] +*$py.class +*.egg-info/ +*.egg +dist/ +build/ +*.whl +.eggs/ +*.so + +# Virtual environments +.venv/ +venv/ +ENV/ +env/ + +# IDE +.idea/ +.vscode/ +*.swp +*.swo +*~ +.project +.settings/ +*.sublime-project +*.sublime-workspace + +# OS +.DS_Store +Thumbs.db + +# Environment / secrets +.env +.env.* +*.pem +*.key + +# Pre-commit +.pre-commit-config.yaml +scripts/ + +# Logs and reports +*.log +*.out +*.report +*.summary diff --git a/common.py b/common.py index 6250a76..6dc8f88 100644 --- a/common.py +++ b/common.py @@ -2,19 +2,21 @@ # # contains all variables/settings to be used in other scripts -cloud_type = "azure" # cloud where primary/secondary metastores exist -cred_mapping_file = "data/azure_cred_mapping.csv" # path to credential mapping file +cloud_type = "azure" # cloud where primary/secondary metastores exist +cred_mapping_file = "data/azure_cred_mapping.csv" # path to credential mapping file loc_mapping_file = "data/ext_location_mapping.csv" # path to locations mapping file -catalog_mapping_file = "data/catalog_mapping.csv" # path to catalog mapping file -schema_mapping_file = "data/schema_mapping.csv" # path to schema mapping file -source_host = "" # source hostname, including https:// -source_pat = "" # source PAT -target_host = "" # target hostname, including https:// -target_pat = "" # targe PAT -catalogs_to_copy = ["my-catalog1", "my-catalog2"] # list of catalogs to replicate -landing_zone_url = "path/to/storage/" # if using sync_tables, intermediate storage location -num_exec = 4 # number of parallel threads to execute -warehouse_size = "Small" # serverless warehouse size in target WS -response_backoff = 0.5 # polling backoff for checking query status -metastore_id = "" # global metastore ID for secondary metastore -manifest_name = "manifest" # name of the manifest file, if written +catalog_mapping_file = "data/catalog_mapping.csv" # path to catalog mapping file +schema_mapping_file = "data/schema_mapping.csv" # path to schema mapping file +source_host = "" # source hostname, including https:// +source_pat = "" # source PAT +target_host = "" # target hostname, including https:// +target_pat = "" # targe PAT +catalogs_to_copy = ["my-catalog1", "my-catalog2"] # list of catalogs to replicate +landing_zone_url = ( + "path/to/storage/" # if using sync_tables, intermediate storage location +) +num_exec = 4 # number of parallel threads to execute +warehouse_size = "Small" # serverless warehouse size in target WS +response_backoff = 0.5 # polling backoff for checking query status +metastore_id = "" # global metastore ID for secondary metastore +manifest_name = "manifest" # name of the manifest file, if written diff --git a/dr_sync/__init__.py b/dr_sync/__init__.py new file mode 100644 index 0000000..25a0af7 --- /dev/null +++ b/dr_sync/__init__.py @@ -0,0 +1,39 @@ +"""DR Sync — shared utilities for Databricks Disaster Recovery scripts.""" + +from dr_sync.exceptions import ( + DRSyncError, + ConfigurationError, + MappingError, + StatementError, + WarehouseError, + SyncError, +) +from dr_sync.sql_utils import ( + execute_statement_sync, + managed_warehouse, + drop_table_if_exists, +) +from dr_sync.workspace import create_client +from dr_sync.csv_mapping import load_mapping, lookup_value +from dr_sync.thread_utils import parallel_map, ProgressCounter +from dr_sync.config import DRSyncConfig +from dr_sync.log import setup_logging + +__all__ = [ + "DRSyncError", + "ConfigurationError", + "MappingError", + "StatementError", + "WarehouseError", + "SyncError", + "execute_statement_sync", + "managed_warehouse", + "drop_table_if_exists", + "create_client", + "load_mapping", + "lookup_value", + "parallel_map", + "ProgressCounter", + "DRSyncConfig", + "setup_logging", +] diff --git a/dr_sync/config.py b/dr_sync/config.py new file mode 100644 index 0000000..1547d99 --- /dev/null +++ b/dr_sync/config.py @@ -0,0 +1,131 @@ +"""Configuration management for DR sync scripts.""" + +import os +from dataclasses import dataclass, field +from typing import List + +from dr_sync.exceptions import ConfigurationError + + +@dataclass +class DRSyncConfig: + """Central configuration for all DR sync scripts. + + Can be populated from common.py (backward compat) or environment variables. + """ + + # Cloud and workspace settings + cloud_type: str = "azure" + source_host: str = "" + source_token: str = "" + target_host: str = "" + target_token: str = "" + + # Catalogs + catalogs_to_copy: List[str] = field(default_factory=list) + + # Mapping file paths + cred_mapping_file: str = "data/azure_cred_mapping.csv" + loc_mapping_file: str = "data/ext_location_mapping.csv" + catalog_mapping_file: str = "data/catalog_mapping.csv" + schema_mapping_file: str = "data/schema_mapping.csv" + + # Execution settings + landing_zone_url: str = "" + num_exec: int = 4 + warehouse_size: str = "Small" + response_backoff: float = 0.5 + metastore_id: str = "" + manifest_name: str = "manifest" + + # Runtime flags + dry_run: bool = False + + @classmethod + def from_common_module(cls): + """Create config by importing from common.py (backward compatible).""" + try: + import common + except ImportError: + raise ConfigurationError( + "common.py not found. Please create it or use environment variables." + ) + + kwargs = {} + field_map = { + "cloud_type": "cloud_type", + "source_host": "source_host", + "source_token": "source_pat", + "target_host": "target_host", + "target_token": "target_pat", + "catalogs_to_copy": "catalogs_to_copy", + "cred_mapping_file": "cred_mapping_file", + "loc_mapping_file": "loc_mapping_file", + "catalog_mapping_file": "catalog_mapping_file", + "schema_mapping_file": "schema_mapping_file", + "landing_zone_url": "landing_zone_url", + "num_exec": "num_exec", + "warehouse_size": "warehouse_size", + "response_backoff": "response_backoff", + "metastore_id": "metastore_id", + "manifest_name": "manifest_name", + } + + for config_key, common_key in field_map.items(): + if hasattr(common, common_key): + kwargs[config_key] = getattr(common, common_key) + + return cls(**kwargs) + + @classmethod + def from_env(cls): + """Create config from DR_SYNC_* environment variables.""" + + def get(name, default=""): + return os.environ.get(f"DR_SYNC_{name}", default) + + catalogs = get("CATALOGS_TO_COPY", "") + catalog_list = ( + [c.strip() for c in catalogs.split(",") if c.strip()] if catalogs else [] + ) + + return cls( + cloud_type=get("CLOUD_TYPE", "azure"), + source_host=get("SOURCE_HOST"), + source_token=get("SOURCE_TOKEN"), + target_host=get("TARGET_HOST"), + target_token=get("TARGET_TOKEN"), + catalogs_to_copy=catalog_list, + cred_mapping_file=get("CRED_MAPPING_FILE", "data/azure_cred_mapping.csv"), + loc_mapping_file=get("LOC_MAPPING_FILE", "data/ext_location_mapping.csv"), + catalog_mapping_file=get( + "CATALOG_MAPPING_FILE", "data/catalog_mapping.csv" + ), + schema_mapping_file=get("SCHEMA_MAPPING_FILE", "data/schema_mapping.csv"), + landing_zone_url=get("LANDING_ZONE_URL"), + num_exec=int(get("NUM_EXEC", "4")), + warehouse_size=get("WAREHOUSE_SIZE", "Small"), + response_backoff=float(get("RESPONSE_BACKOFF", "0.5")), + metastore_id=get("METASTORE_ID"), + manifest_name=get("MANIFEST_NAME", "manifest"), + dry_run=get("DRY_RUN", "false").lower() in ("true", "1", "yes"), + ) + + def validate(self) -> List[str]: + """Validate configuration and return list of errors (empty = valid).""" + errors = [] + + if not self.target_host: + errors.append("target_host is required") + if not self.target_token: + errors.append("target_token is required") + if not self.catalogs_to_copy: + errors.append("catalogs_to_copy must not be empty") + if self.cloud_type not in ("aws", "azure", "gcp"): + errors.append( + f"cloud_type must be one of aws, azure, gcp (got {self.cloud_type!r})" + ) + if self.num_exec < 1: + errors.append(f"num_exec must be >= 1 (got {self.num_exec})") + + return errors diff --git a/dr_sync/csv_mapping.py b/dr_sync/csv_mapping.py new file mode 100644 index 0000000..60144f0 --- /dev/null +++ b/dr_sync/csv_mapping.py @@ -0,0 +1,108 @@ +"""CSV mapping file loading and lookup utilities.""" + +import logging +import os + +import pandas as pd + +from dr_sync.exceptions import MappingError + +logger = logging.getLogger("dr_sync") + + +def load_mapping(filepath, required_columns=None): + """Load a CSV mapping file with validation. + + Args: + filepath: Path to the CSV file. + required_columns: Optional list of column names that must be present. + + Returns: + pandas DataFrame with the mapping data. + + Raises: + MappingError: If the file doesn't exist or required columns are missing. + """ + if not os.path.exists(filepath): + raise MappingError(filepath, "", "", f"Mapping file not found: {filepath}") + + df = pd.read_csv(filepath, keep_default_na=False) + + if required_columns: + missing = set(required_columns) - set(df.columns) + if missing: + raise MappingError( + filepath, "", "", f"Missing required columns in {filepath}: {missing}" + ) + + return df + + +def lookup_value(df, key_col, key_val, value_col): + """Safe lookup of a single value from a mapping DataFrame. + + Args: + df: pandas DataFrame to search. + key_col: Column name to match against. + key_val: Value to search for. + value_col: Column name to return. + + Returns: + The matched value, or None if no match found. + """ + matches = df[value_col].loc[df[key_col] == key_val] + if matches.empty: + return None + return matches.iloc[0] + + +def validate_catalog_mapping(filepath): + """Validate catalog mapping CSV for duplicates and valid storage roots. + + Returns: + List of error strings (empty = valid). + """ + errors = [] + df = load_mapping( + filepath, required_columns=["source_catalog", "target_storage_root"] + ) + dupes = df[df.duplicated(subset=["source_catalog"], keep=False)] + if not dupes.empty: + dupe_names = dupes["source_catalog"].unique().tolist() + errors.append(f"Duplicate source_catalog entries: {dupe_names}") + return errors + + +def validate_cred_mapping(filepath, cloud_type): + """Validate credential mapping CSV has cloud-specific required columns. + + Returns: + List of error strings (empty = valid). + """ + errors = [] + if cloud_type == "aws": + df = load_mapping( + filepath, required_columns=["source_cred_name", "target_iam_role"] + ) + empty_roles = df[df["target_iam_role"] == ""] + if not empty_roles.empty: + names = empty_roles["source_cred_name"].tolist() + errors.append(f"Empty target_iam_role for credentials: {names}") + elif cloud_type == "azure": + df = load_mapping(filepath, required_columns=["source_cred_name"]) + return errors + + +def validate_ext_location_mapping(filepath): + """Validate external location mapping CSV for non-empty URLs. + + Returns: + List of error strings (empty = valid). + """ + errors = [] + df = load_mapping(filepath, required_columns=["source_loc_name", "target_url"]) + empty_urls = df[df["target_url"] == ""] + if not empty_urls.empty: + names = empty_urls["source_loc_name"].tolist() + errors.append(f"Empty target_url for locations: {names}") + return errors diff --git a/dr_sync/exceptions.py b/dr_sync/exceptions.py new file mode 100644 index 0000000..c35882f --- /dev/null +++ b/dr_sync/exceptions.py @@ -0,0 +1,41 @@ +"""Custom exception hierarchy for DR sync operations.""" + + +class DRSyncError(Exception): + """Base exception for all DR sync errors.""" + + +class ConfigurationError(DRSyncError): + """Raised when configuration is invalid or missing.""" + + +class MappingError(DRSyncError): + """Raised when CSV mapping lookup fails.""" + + def __init__(self, mapping_file, key_col, key_val, message=None): + self.mapping_file = mapping_file + self.key_col = key_col + self.key_val = key_val + msg = message or f"No mapping found for {key_col}={key_val!r} in {mapping_file}" + super().__init__(msg) + + +class StatementError(DRSyncError): + """Raised when a SQL statement execution fails.""" + + def __init__(self, statement, message): + self.statement = statement + super().__init__(f"Statement failed: {message}\nSQL: {statement[:200]}") + + +class WarehouseError(DRSyncError): + """Raised when warehouse operations fail.""" + + +class SyncError(DRSyncError): + """Raised when syncing a specific resource fails.""" + + def __init__(self, resource_type, resource_name, message): + self.resource_type = resource_type + self.resource_name = resource_name + super().__init__(f"Failed to sync {resource_type} {resource_name!r}: {message}") diff --git a/dr_sync/log.py b/dr_sync/log.py new file mode 100644 index 0000000..ad22381 --- /dev/null +++ b/dr_sync/log.py @@ -0,0 +1,40 @@ +"""Logging setup for DR sync scripts.""" + +import logging +import sys + + +def setup_logging(level="INFO", log_file=None): + """Configure logging with console and optional file handler. + + Args: + level: Log level string (DEBUG, INFO, WARNING, ERROR). + log_file: Optional path to write logs to a file. + + Returns: + The root logger configured for DR sync. + """ + logger = logging.getLogger("dr_sync") + logger.setLevel(getattr(logging, level.upper(), logging.INFO)) + + # avoid adding duplicate handlers on re-invocation + if logger.handlers: + return logger + + formatter = logging.Formatter( + "%(asctime)s [%(levelname)s] %(message)s", + datefmt="%Y-%m-%d %H:%M:%S", + ) + + # console handler (stdout for notebook compatibility) + console = logging.StreamHandler(sys.stdout) + console.setFormatter(formatter) + logger.addHandler(console) + + # optional file handler + if log_file: + file_handler = logging.FileHandler(log_file) + file_handler.setFormatter(formatter) + logger.addHandler(file_handler) + + return logger diff --git a/dr_sync/sql_utils.py b/dr_sync/sql_utils.py new file mode 100644 index 0000000..14004eb --- /dev/null +++ b/dr_sync/sql_utils.py @@ -0,0 +1,136 @@ +"""SQL statement execution utilities and warehouse lifecycle management.""" + +import logging +import time +from contextlib import contextmanager + +from databricks.sdk.service.sql import ( + Disposition, + StatementState, + CreateWarehouseRequestWarehouseType, + ExecuteStatementRequestOnWaitTimeout, +) +from databricks.sdk.service import sql as dbsql + +from dr_sync.exceptions import StatementError, WarehouseError + +logger = logging.getLogger("dr_sync") + + +def execute_statement_sync( + client, warehouse_id, statement, backoff=0.5, timeout_seconds=3600 +): + """Execute a SQL statement and poll until completion. + + Args: + client: WorkspaceClient instance. + warehouse_id: ID of the warehouse to execute on. + statement: SQL statement string. + backoff: Seconds between polling attempts. + timeout_seconds: Maximum seconds to wait before cancelling. + + Returns: + The final statement response object. + + Raises: + StatementError: If the statement fails or times out. + """ + resp = client.statement_execution.execute_statement( + warehouse_id=warehouse_id, + wait_timeout="0s", + on_wait_timeout=ExecuteStatementRequestOnWaitTimeout("CONTINUE"), + disposition=Disposition("EXTERNAL_LINKS"), + statement=statement, + ) + + start = time.monotonic() + while resp.status.state in {StatementState.PENDING, StatementState.RUNNING}: + if time.monotonic() - start > timeout_seconds: + try: + client.statement_execution.cancel_execution(resp.statement_id) + except Exception: + pass + raise StatementError(statement, f"Timed out after {timeout_seconds}s") + time.sleep(backoff) + resp = client.statement_execution.get_statement(resp.statement_id) + + if resp.status.state != StatementState.SUCCEEDED: + error_msg = resp.status.error.message if resp.status.error else "Unknown error" + raise StatementError(statement, error_msg) + + return resp + + +@contextmanager +def managed_warehouse(client, size="Small", name_prefix="sdk"): + """Context manager that creates a serverless warehouse and deletes it on exit. + + Args: + client: WorkspaceClient instance. + size: Warehouse cluster size. + name_prefix: Prefix for the warehouse name. + + Yields: + The warehouse ID string. + + Raises: + WarehouseError: If warehouse creation fails. + """ + wh_type = CreateWarehouseRequestWarehouseType("PRO") + + try: + wh = client.warehouses.create( + name=f"{name_prefix}-{time.time_ns()}", + cluster_size=size, + max_num_clusters=1, + auto_stop_mins=10, + warehouse_type=wh_type, + enable_serverless_compute=True, + tags=dbsql.EndpointTags( + custom_tags=[dbsql.EndpointTagPair(key="Owner", value="dr-sync-tool")] + ), + ).result() + except Exception as e: + raise WarehouseError(f"Failed to create warehouse: {e}") from e + + try: + yield wh.id + finally: + try: + client.warehouses.delete(wh.id) + logger.info("Cleaned up warehouse %s", wh.id) + except Exception as e: + logger.warning("Could not delete warehouse %s: %s", wh.id, e) + + +def drop_table_if_exists( + client, warehouse_id, catalog, schema, table_name, backoff=0.5 +): + """Drop a table if it exists via SQL statement execution. + + Returns: + dict with status (1=success, 0=failure) and table identifiers. + """ + fqn = f"{catalog}.{schema}.{table_name}" + logger.info("Dropping table %s...", fqn) + + try: + execute_statement_sync( + client, + warehouse_id, + f"DROP TABLE IF EXISTS {fqn}", + backoff=backoff, + ) + return { + "status": 1, + "catalog": catalog, + "schema": schema, + "table_name": table_name, + } + except Exception: + return { + "status": 0, + "catalog": catalog, + "schema": schema, + "table_name": table_name, + } diff --git a/dr_sync/thread_utils.py b/dr_sync/thread_utils.py new file mode 100644 index 0000000..77b2ac5 --- /dev/null +++ b/dr_sync/thread_utils.py @@ -0,0 +1,53 @@ +"""Thread pool utilities for parallel execution with error isolation.""" + +import threading +from concurrent.futures import ThreadPoolExecutor, as_completed + + +def parallel_map(func, items, max_workers=4): + """Execute func on each item in parallel with per-item error isolation. + + Unlike executor.map(), uses as_completed() so one failure doesn't block others. + + Args: + func: Callable that takes a single item and returns a result. + items: Iterable of items to process. + max_workers: Maximum concurrent workers. + + Returns: + List of results (or exception objects for failed items). + """ + results = [] + with ThreadPoolExecutor(max_workers=max_workers) as executor: + future_to_item = {executor.submit(func, item): item for item in items} + for future in as_completed(future_to_item): + try: + results.append(future.result()) + except Exception as e: + results.append(e) + return results + + +class ProgressCounter: + """Thread-safe counter for progress reporting.""" + + def __init__(self, total, label="items"): + self._count = 0 + self._total = total + self._label = label + self._lock = threading.Lock() + + def increment(self, item_name=""): + """Increment counter and print progress.""" + with self._lock: + self._count += 1 + msg = f"[{self._count}/{self._total}]" + if item_name: + msg += f" Processed {self._label}: {item_name}" + print(msg) + return self._count + + @property + def count(self): + with self._lock: + return self._count diff --git a/dr_sync/workspace.py b/dr_sync/workspace.py new file mode 100644 index 0000000..d16946f --- /dev/null +++ b/dr_sync/workspace.py @@ -0,0 +1,32 @@ +"""WorkspaceClient factory with flexible authentication.""" + +import os + +from databricks.sdk import WorkspaceClient + + +def create_client(host=None, token=None, profile=None): + """Create a WorkspaceClient with auth resolution. + + Resolution order: explicit args -> environment variables -> SDK default chain. + + Args: + host: Workspace URL (e.g. https://adb-xxx.azuredatabricks.net). + token: Personal access token. + profile: Databricks CLI profile name. + + Returns: + Configured WorkspaceClient instance. + """ + resolved_host = host or os.environ.get("DATABRICKS_HOST") + resolved_token = token or os.environ.get("DATABRICKS_TOKEN") + + kwargs = {} + if resolved_host: + kwargs["host"] = resolved_host + if resolved_token: + kwargs["token"] = resolved_token + if profile: + kwargs["profile"] = profile + + return WorkspaceClient(**kwargs) diff --git a/examples/clone_to_secondary.py b/examples/clone_to_secondary.py index 181da32..9a594ea 100644 --- a/examples/clone_to_secondary.py +++ b/examples/clone_to_secondary.py @@ -1,7 +1,7 @@ # clone_to_secondary.py # # This script clones tables from the catalogs listed in catalogs_to_copy into the bucket specified by dest_bucket, which -# can be an S3 bucket, ADLS storage account, or GCS bucket. Tables will be tracked in a manifest file, also written to +# can be an S3 bucket, ADLS storage account, or GCS bucket. Tables will be tracked in a manifest file, also written to # dest_bucket, containing the catalog, schema and table name, table type, and write location. # # We assume this script is run on a Databricks cluster; if it is run locally, you may need to add additional configuration @@ -10,7 +10,7 @@ import pandas as pd # script inputs -catalogs_to_copy = ["my_catalog1, my_catalog2"] +catalogs_to_copy = ["my_catalog1", "my_catalog2"] dest_bucket = "s3://path/to/intermediate/location" manifest_name = "manifest" @@ -24,37 +24,45 @@ # loop through all catalogs to copy, then copy all tables excluding system tables. for catalog in catalogs_to_copy: - filtered_tables = system_info.filter((system_info.table_catalog == catalog) - & (system_info.table_schema != "information_schema") - & (system_info.table_type != "VIEW")) - - for table in filtered_tables.collect(): - schema = table['table_schema'] - table_name = table['table_name'] - table_type = table['table_type'] - - # skip views - if table_type == "VIEW": - continue - - print(f"Copying table {schema}.{table_name}...") - sqlstring = f"CREATE TABLE delta.`{dest_bucket}/{catalog}_{schema}_{table_name}` DEEP CLONE {catalog}.{schema}.{table_name}" - sql(sqlstring) - - # the below will be used to create the manifest table in the secondary region - copied_table_names.append(table_name) - copied_table_types.append(table_type) - copied_table_schemas.append(schema) - copied_table_catalogs.append(catalog) - copied_table_locations.append(f"{dest_bucket}/{catalog}_{schema}_{table_name}") + filtered_tables = system_info.filter( + (system_info.table_catalog == catalog) + & (system_info.table_schema != "information_schema") + & (system_info.table_type != "VIEW") + ) + + for table in filtered_tables.collect(): + schema = table["table_schema"] + table_name = table["table_name"] + table_type = table["table_type"] + + # skip views + if table_type == "VIEW": + continue + + print(f"Copying table {schema}.{table_name}...") + sqlstring = f"CREATE TABLE delta.`{dest_bucket}/{catalog}_{schema}_{table_name}` DEEP CLONE {catalog}.{schema}.{table_name}" + sql(sqlstring) + + # the below will be used to create the manifest table in the secondary region + copied_table_names.append(table_name) + copied_table_types.append(table_type) + copied_table_schemas.append(schema) + copied_table_catalogs.append(catalog) + copied_table_locations.append(f"{dest_bucket}/{catalog}_{schema}_{table_name}") # create the manifest as a df and write to a table in dr target # this contains catalog, schema, table and location -manifest_df = pd.DataFrame({"catalog": copied_table_catalogs, - "schema": copied_table_schemas, - "table": copied_table_names, - "location": copied_table_locations, - "type": copied_table_types}) +manifest_df = pd.DataFrame( + { + "catalog": copied_table_catalogs, + "schema": copied_table_schemas, + "table": copied_table_names, + "location": copied_table_locations, + "type": copied_table_types, + } +) -spark.createDataFrame(manifest_df).write.mode("overwrite").format("delta").save(f"{dest_bucket}/{manifest_name}") +spark.createDataFrame(manifest_df).write.mode("overwrite").format("delta").save( + f"{dest_bucket}/{manifest_name}" +) display(manifest_df) diff --git a/examples/clone_to_secondary_par.py b/examples/clone_to_secondary_par.py index b155a81..96cfe3f 100644 --- a/examples/clone_to_secondary_par.py +++ b/examples/clone_to_secondary_par.py @@ -3,7 +3,7 @@ # Parallelized version of clone_to_secondary.py. # # This script clones tables from the catalogs listed in catalogs_to_copy into the bucket specified by dest_bucket, which -# can be an S3 bucket, ADLS storage account, or GCS bucket. Tables will be tracked in a manifest file, also written to +# can be an S3 bucket, ADLS storage account, or GCS bucket. Tables will be tracked in a manifest file, also written to # dest_bucket, containing the catalog, schema and table name, table type, and write location. # # We assume this script is run on a Databricks cluster; if it is run locally, you may need to add additional configuration @@ -13,24 +13,30 @@ from itertools import repeat from concurrent.futures import ThreadPoolExecutor + # helper function to copy tables def copy_table(catalog, schema, table_name, table_type, dest_bucket): try: - sqlstring = f"CREATE TABLE delta.`{dest_bucket}/{catalog}_{schema}_{table_name}` DEEP CLONE {catalog}.{schema}.{table_name}" - sql(sqlstring) + sqlstring = f"CREATE TABLE delta.`{dest_bucket}/{catalog}_{schema}_{table_name}` DEEP CLONE {catalog}.{schema}.{table_name}" + sql(sqlstring) - # return the table params in dict; used to build manifest - return {"catalog":catalog, - "schema":schema, - "table_name":table_name, - "table_type":table_type, - "dest_bucket":dest_bucket} + # return the table params in dict; used to build manifest + return { + "catalog": catalog, + "schema": schema, + "table_name": table_name, + "table_type": table_type, + "dest_bucket": dest_bucket, + } except Exception: - return {"catalog":catalog, - "schema":schema, - "table_name":table_name, - "table_type":"ERROR", - "dest_bucket":"N/A"} + return { + "catalog": catalog, + "schema": schema, + "table_name": table_name, + "table_type": "ERROR", + "dest_bucket": "N/A", + } + # script inputs catalogs_to_copy = ["my_catalog1", "my_catalog2"] @@ -48,39 +54,60 @@ def copy_table(catalog, schema, table_name, table_type, dest_bucket): # loop through all catalogs to copy, then copy all tables excluding system tables. for catalog in catalogs_to_copy: - filtered_tables = system_info.filter((system_info.table_catalog == catalog) - & (system_info.table_schema != "information_schema") - & (system_info.table_type != "VIEW")).collect() + filtered_tables = system_info.filter( + (system_info.table_catalog == catalog) + & (system_info.table_schema != "information_schema") + & (system_info.table_type != "VIEW") + ).collect() + + # get schemas, tables and types in list form + schemas = [row["table_schema"] for row in filtered_tables] + table_names = [row["table_name"] for row in filtered_tables] + table_types = [row["table_type"] for row in filtered_tables] - # get schemas, tables and types in list form - schemas = [row['table_schema'] for row in filtered_tables] - table_names = [row['table_name'] for row in filtered_tables] - table_types = [row['table_type'] for row in filtered_tables] + # use ThreadPoolExecutor to copy tables in parallel + with ThreadPoolExecutor(max_workers=num_exec) as executor: + threads = executor.map( + copy_table, + repeat(catalog), + schemas, + table_names, + table_types, + repeat(dest_bucket), + ) - # use ThreadPoolExecutor to copy tables in parallel - with ThreadPoolExecutor(max_workers = num_exec) as executor: - threads = executor.map(copy_table, - repeat(catalog), - schemas, - table_names, - table_types, - repeat(dest_bucket)) - - # wait for threads to execute and build lists for manifest - for thread in threads: - copied_table_names.append(thread["table_name"]) - copied_table_types.append(thread["table_type"]) - copied_table_schemas.append(thread["schema"]) - copied_table_catalogs.append(thread["catalog"]) - copied_table_locations.append("{}/{}_{}_{}".format(thread["dest_bucket"],thread["catalog"],thread["schema"],thread["table_name"])) - print("Copied table {}.{}.{}.".format(thread["catalog"],thread["schema"],thread["table_name"])) + # wait for threads to execute and build lists for manifest + for thread in threads: + copied_table_names.append(thread["table_name"]) + copied_table_types.append(thread["table_type"]) + copied_table_schemas.append(thread["schema"]) + copied_table_catalogs.append(thread["catalog"]) + copied_table_locations.append( + "{}/{}_{}_{}".format( + thread["dest_bucket"], + thread["catalog"], + thread["schema"], + thread["table_name"], + ) + ) + print( + "Copied table {}.{}.{}.".format( + thread["catalog"], thread["schema"], thread["table_name"] + ) + ) # create the manifest as a df and write to a table in dr target # this contains catalog, schema, table and location -manifest_df = pd.DataFrame({"catalog": copied_table_catalogs, - "schema": copied_table_schemas, - "table": copied_table_names, - "location": copied_table_locations, - "type": copied_table_types}) +manifest_df = pd.DataFrame( + { + "catalog": copied_table_catalogs, + "schema": copied_table_schemas, + "table": copied_table_names, + "location": copied_table_locations, + "type": copied_table_types, + } +) -spark.createDataFrame(manifest_df).write.mode("overwrite").format("delta").save(f"{dest_bucket}/{manifest_name}") +spark.createDataFrame(manifest_df).write.mode("overwrite").format("delta").save( + f"{dest_bucket}/{manifest_name}" +) diff --git a/examples/create_tables_simple.py b/examples/create_tables_simple.py index 0675124..2708fce 100644 --- a/examples/create_tables_simple.py +++ b/examples/create_tables_simple.py @@ -13,9 +13,9 @@ # loop through manifest and create each table for row in manifest_df.collect(): - catalog = row['catalog'] - schema = row['schema'] - table_name = row['table'] + catalog = row["catalog"] + schema = row["schema"] + table_name = row["table"] location = row["location"] tbl_type = row["type"] @@ -28,4 +28,6 @@ sqlstring = f"CREATE TABLE {catalog}.{schema}.{table_name} USING delta LOCATION '{location}'" sql(sqlstring) else: - print(f"Skipping table {catalog}.{schema}.{table_name}; please check manifest file.") + print( + f"Skipping table {catalog}.{schema}.{table_name}; please check manifest file." + ) diff --git a/ruff.toml b/ruff.toml new file mode 100644 index 0000000..a4c0867 --- /dev/null +++ b/ruff.toml @@ -0,0 +1,2 @@ +# Databricks notebook builtins (spark, sql, display) are injected by the runtime. +builtins = ["spark", "sql", "display"] diff --git a/sync_catalogs_and_schemas.py b/sync_catalogs_and_schemas.py index 5e0507a..363e127 100644 --- a/sync_catalogs_and_schemas.py +++ b/sync_catalogs_and_schemas.py @@ -15,12 +15,27 @@ # Currently, we use PAT-based auth for the WorkspaceClient objects, so you must provide the host and token manually for # each workspace. You can update this to use other auth methods if desired. +import argparse +import os +from concurrent.futures import ThreadPoolExecutor from databricks.sdk import WorkspaceClient -import pandas as pd -from common import (target_pat, target_host, - source_pat, source_host, - catalogs_to_copy, catalog_mapping_file, - schema_mapping_file) +from dr_sync.csv_mapping import load_mapping +from dr_sync.config import DRSyncConfig +from dr_sync.log import setup_logging + +config = ( + DRSyncConfig.from_env() + if os.environ.get("DR_SYNC_SOURCE_HOST") + else DRSyncConfig.from_common_module() +) +logger = setup_logging() +target_host = config.target_host +target_pat = config.target_token +source_host = config.source_host +source_pat = config.source_token +catalogs_to_copy = config.catalogs_to_copy +catalog_mapping_file = config.catalog_mapping_file +schema_mapping_file = config.schema_mapping_file # create WorkspaceClient objects @@ -37,16 +52,20 @@ target_catalog_names = [x.name for x in target_catalogs] catalog_diff = list(set(source_catalog_names) - set(target_catalog_names)) catalogs_to_create = [x for x in source_catalogs if x.name in catalog_diff] -catalog_df = pd.read_csv(catalog_mapping_file, keep_default_na=False) +catalog_df = load_mapping(catalog_mapping_file) +catalog_lookup = catalog_df.set_index("source_catalog").to_dict("index") if not catalogs_to_create: - print("All source catalogs exist in target metastore.") + logger.info("All source catalogs exist in target metastore.") for catalog in catalogs_to_create: # skip shared or external catalogs if catalog.connection_name or catalog.share_name: - print(f"External Catalogs and Shared Catalogs are not currently supported by this script. \ - Skipping {catalog.name}...") + logger.warning( + "External Catalogs and Shared Catalogs are not currently supported by this script. " + "Skipping %s...", + catalog.name, + ) continue # get parameters that map directly between catalogs @@ -55,63 +74,126 @@ catalog_options = catalog.options catalog_properties = catalog.properties - print(f"Creating catalog {catalog_name}...") + logger.info("Creating catalog %s...", catalog_name) # get target storage root based off of catalog name - try: - storage_root = catalog_df['target_storage_root'].loc[catalog_df['source_catalog'] == catalog_name].iloc[0] - except (KeyError, IndexError): - print(f"Could not create catalog {catalog_name}. Please check mapping file.") + row = catalog_lookup.get(catalog_name) + if row is None: + logger.error( + "Could not create catalog %s. Please check mapping file.", catalog_name + ) continue + storage_root = row["target_storage_root"] # create catalog in target metastore + if config.dry_run: + logger.info("[DRY RUN] Would create catalog %s", catalog_name) + continue + if storage_root: - w_target.catalogs.create(name=catalog_name, - comment=catalog_comment, - options=catalog_options, - properties=catalog_properties, - storage_root=storage_root) + w_target.catalogs.create( + name=catalog_name, + comment=catalog_comment, + options=catalog_options, + properties=catalog_properties, + storage_root=storage_root, + ) else: - w_target.catalogs.create(name=catalog_name, - comment=catalog_comment, - options=catalog_options, - properties=catalog_properties) + w_target.catalogs.create( + name=catalog_name, + comment=catalog_comment, + options=catalog_options, + properties=catalog_properties, + ) - print(f"Created catalog {catalog_name}.") + logger.info("Created catalog %s.", catalog_name) -schema_df = pd.read_csv(schema_mapping_file, keep_default_na=False) +schema_df = load_mapping(schema_mapping_file) +# Build a lookup keyed by (source_catalog, source_schema) for O(1) access +schema_lookup = {} +for _, srow in schema_df.iterrows(): + key = (srow["source_catalog"], srow["source_schema"]) + schema_lookup[key] = srow.to_dict() -for catalog in source_catalogs: - source_schemas = [x for x in w_source.schemas.list(catalog.name)] - target_schemas = [x for x in w_target.schemas.list(catalog.name)] + +def create_schema(catalog_name, schema_obj, storage_root): + """Create a single schema in the target workspace. Returns a status dict.""" + schema_name = schema_obj.name + schema_comment = schema_obj.comment + schema_properties = schema_obj.properties + + try: + if storage_root: + w_target.schemas.create( + name=schema_name, + comment=schema_comment, + properties=schema_properties, + catalog_name=catalog_name, + storage_root=storage_root, + ) + else: + w_target.schemas.create( + name=schema_name, + comment=schema_comment, + properties=schema_properties, + catalog_name=catalog_name, + ) + logger.info("Created schema %s.%s.", catalog_name, schema_name) + except Exception as e: + logger.error("Error creating schema %s.%s: %s", catalog_name, schema_name, e) + + +# Collect all schema creation tasks across catalogs +schema_tasks = [] +for cat in source_catalogs: + source_schemas = [x for x in w_source.schemas.list(cat.name)] + target_schemas = [x for x in w_target.schemas.list(cat.name)] source_schema_names = [x.name for x in source_schemas] target_schema_names = [x.name for x in target_schemas] schema_diff = list(set(source_schema_names) - set(target_schema_names)) schemas_to_create = [x for x in source_schemas if x.name in schema_diff] for schema in schemas_to_create: - schema_name = schema.name - schema_comment = schema.comment - schema_properties = schema.properties - - try: - storage_root = (schema_df['target_storage_root'].loc[ - (schema_df['source_schema'] == schema_name) & - (schema_df['source_catalog'] == catalog.name)].iloc[0]) - except (KeyError, IndexError): - print(f"Could not create schema {catalog.name}.{schema_name}. Please check mapping file.") + row = schema_lookup.get((cat.name, schema.name)) + if row is None: + logger.error( + "Could not create schema %s.%s. Please check mapping file.", + cat.name, + schema.name, + ) continue - if storage_root: - w_target.schemas.create(name=schema_name, - comment=schema_comment, - properties=schema_properties, - catalog_name=catalog.name, - storage_root=storage_root) - else: - w_target.schemas.create(name=schema_name, - comment=schema_comment, - properties=schema_properties, - catalog_name=catalog.name) + storage_root = row["target_storage_root"] + + if config.dry_run: + logger.info("[DRY RUN] Would create schema %s.%s", cat.name, schema.name) + continue - print(f"Created schema {catalog.name}.{schema_name}.") + schema_tasks.append((cat.name, schema, storage_root)) + +# Execute schema creation in parallel +if schema_tasks: + catalog_names = [t[0] for t in schema_tasks] + schema_objs = [t[1] for t in schema_tasks] + storage_roots = [t[2] for t in schema_tasks] + with ThreadPoolExecutor(max_workers=config.num_exec) as executor: + list(executor.map(create_schema, catalog_names, schema_objs, storage_roots)) + +if __name__ == "__main__": + parser = argparse.ArgumentParser( + description="Sync catalogs and schemas between workspaces" + ) + parser.add_argument( + "--dry-run", + action="store_true", + help="Show planned operations without executing", + ) + parser.add_argument( + "--log-level", + default="INFO", + choices=["DEBUG", "INFO", "WARNING", "ERROR"], + help="Set logging level", + ) + args = parser.parse_args() + config.dry_run = args.dry_run + logger = setup_logging(level=args.log_level) diff --git a/sync_creds_and_locs.py b/sync_creds_and_locs.py index 33ac4fd..6725795 100644 --- a/sync_creds_and_locs.py +++ b/sync_creds_and_locs.py @@ -23,13 +23,27 @@ # cloud object information in the provided CSVs, especially for Azure; this could be done by directly interfacing with # the cloud provider CLI/APIs within this script (or as part of an external workflow). +import argparse +import os from databricks.sdk import WorkspaceClient from databricks.sdk.service import catalog -import pandas as pd -from common import (target_pat, target_host, - source_pat, source_host, - cred_mapping_file, loc_mapping_file, - cloud_type) +from dr_sync.csv_mapping import load_mapping +from dr_sync.config import DRSyncConfig +from dr_sync.log import setup_logging + +config = ( + DRSyncConfig.from_env() + if os.environ.get("DR_SYNC_SOURCE_HOST") + else DRSyncConfig.from_common_module() +) +logger = setup_logging() +target_host = config.target_host +target_pat = config.target_token +source_host = config.source_host +source_pat = config.source_token +cred_mapping_file = config.cred_mapping_file +loc_mapping_file = config.loc_mapping_file +cloud_type = config.cloud_type # create WorkspaceClient objects w_source = WorkspaceClient(host=source_host, token=source_pat) @@ -47,80 +61,113 @@ target_cred_names = [x.name for x in target_creds] cred_diff = list(set(source_cred_names) - set(target_cred_names)) creds_to_create = [x for x in source_creds if x.name in cred_diff] -cred_df = pd.read_csv(cred_mapping_file, keep_default_na=False) +cred_df = load_mapping(cred_mapping_file) +cred_lookup = cred_df.set_index("source_cred_name").to_dict("index") if not creds_to_create: - print("All source credentials exist in target metastore.") + logger.info("All source credentials exist in target metastore.") for cred in creds_to_create: # get parameters that map directly between creds cred_name = cred.name cred_read_only = cred.read_only cred_comment = cred.comment - print(f"Creating storage credential {cred_name}...") + logger.info("Creating storage credential %s...", cred_name) if cloud_type == "aws": # get cred IAM role based off of name - try: - iam_role_arn = cred_df['target_iam_role'].loc[cred_df['source_cred_name'] == cred_name].iloc[0] - except (KeyError, IndexError): - print(f"Could not create credential {cred_name}. Please check mapping file.") + row = cred_lookup.get(cred_name) + if row is None: + logger.error( + "Could not create credential %s. Please check mapping file.", cred_name + ) continue + iam_role_arn = row["target_iam_role"] # create storage credential in target WS cred_iam_role = catalog.AwsIamRole(role_arn=iam_role_arn) - w_target.storage_credentials.create(name=cred_name, - read_only=cred_read_only, - comment=cred_comment, - aws_iam_role=cred_iam_role) + if config.dry_run: + logger.info("[DRY RUN] Would create credential %s", cred_name) + continue + w_target.storage_credentials.create( + name=cred_name, + read_only=cred_read_only, + comment=cred_comment, + aws_iam_role=cred_iam_role, + ) elif cloud_type == "azure": # get SP and Mgd ID info based off of name - try: - managed_id_connector = \ - cred_df['target_mgd_id_connector'].loc[cred_df['source_cred_name'] == cred_name].iloc[0] - managed_id_identity = \ - cred_df['target_mgd_id_identity'].loc[cred_df['source_cred_name'] == cred_name].iloc[0] - sp_directory = cred_df['target_sp_directory'].loc[cred_df['source_cred_name'] == cred_name].iloc[0] - sp_appid = cred_df['target_sp_appid'].loc[cred_df['source_cred_name'] == cred_name].iloc[0] - sp_secret = cred_df['target_sp_secret'].loc[cred_df['source_cred_name'] == cred_name].iloc[0] - except (KeyError, IndexError): - print(f"Could not create credential {cred_name}. Please check mapping file.") + row = cred_lookup.get(cred_name) + if row is None: + logger.error( + "Could not create credential %s. Please check mapping file.", cred_name + ) + continue + managed_id_connector = row.get("target_mgd_id_connector", "") + managed_id_identity = row.get("target_mgd_id_identity", "") + sp_directory = row.get("target_sp_directory", "") + sp_appid = row.get("target_sp_appid", "") + sp_secret = row.get("target_sp_secret", "") + + if not managed_id_connector and not managed_id_identity and not sp_directory: + logger.error( + "Could not create credential %s. Please check mapping file.", cred_name + ) continue # create storage credential in target WS + if config.dry_run: + logger.info("[DRY RUN] Would create credential %s", cred_name) + continue if managed_id_connector: - cred_mgd_id = catalog.AzureManagedIdentityRequest(access_connector_id=managed_id_connector) - w_target.storage_credentials.create(name=cred_name, - read_only=cred_read_only, - comment=cred_comment, - azure_managed_identity=cred_mgd_id) + cred_mgd_id = catalog.AzureManagedIdentityRequest( + access_connector_id=managed_id_connector + ) + w_target.storage_credentials.create( + name=cred_name, + read_only=cred_read_only, + comment=cred_comment, + azure_managed_identity=cred_mgd_id, + ) elif managed_id_identity: - cred_mgd_id = catalog.AzureManagedIdentityRequest(access_connector_id=managed_id_identity) - w_target.storage_credentials.create(name=cred_name, - read_only=cred_read_only, - comment=cred_comment, - azure_managed_identity=cred_mgd_id) + cred_mgd_id = catalog.AzureManagedIdentityRequest( + access_connector_id=managed_id_identity + ) + w_target.storage_credentials.create( + name=cred_name, + read_only=cred_read_only, + comment=cred_comment, + azure_managed_identity=cred_mgd_id, + ) else: try: - cred_sp = catalog.AzureServicePrincipal(directory_id=sp_directory, - application_id=sp_appid, - client_secret=sp_secret) - w_target.storage_credentials.create(name=cred_name, - read_only=cred_read_only, - comment=cred_comment, - azure_service_principal=cred_sp) + cred_sp = catalog.AzureServicePrincipal( + directory_id=sp_directory, + application_id=sp_appid, + client_secret=sp_secret, + ) + w_target.storage_credentials.create( + name=cred_name, + read_only=cred_read_only, + comment=cred_comment, + azure_service_principal=cred_sp, + ) except Exception: - print(f"Could not create credential {cred_name}. Please make sure that only one of \ - managed_id_connector, managed_id_identity or service_principal info is provided in the mapping.") - + logger.error( + "Could not create credential %s. Please make sure that only one of " + "managed_id_connector, managed_id_identity or service_principal info " + "is provided in the mapping.", + cred_name, + ) + elif cloud_type == "gcp": - print("GCP not yet implemented.") + logger.warning("GCP not yet implemented.") continue else: - print("Cloud type must be one of AWS, GCP, or Azure.") + logger.error("Cloud type must be one of AWS, GCP, or Azure.") continue - print(f"Created storage credential {cred_name}.") + logger.info("Created storage credential %s.", cred_name) # compare source and target external locations # we can only do this by name since the URL and IDs will change between workspaces @@ -128,10 +175,11 @@ target_extloc_names = [x.name for x in target_extloc] loc_diff = list(set(source_extloc_names) - set(target_extloc_names)) locs_to_create = [x for x in source_extloc if x.name in loc_diff] -loc_df = pd.read_csv(loc_mapping_file, keep_default_na=False) +loc_df = load_mapping(loc_mapping_file) +loc_lookup = loc_df.set_index("source_loc_name").to_dict("index") if not locs_to_create: - print("All source external locations exist in target metastore.") + logger.info("All source external locations exist in target metastore.") for loc in locs_to_create: # get parameters that map directly between creds @@ -140,49 +188,98 @@ loc_comment = loc.comment loc_fallback = loc.fallback loc_read_only = loc.read_only - print(f"Creating external location {loc_name}...") + logger.info("Creating external location %s...", loc_name) if cloud_type == "aws": - try: - url = loc_df['target_url'].loc[loc_df['source_loc_name'] == loc_name].iloc[0] - access_pt = loc_df['target_access_pt'].loc[loc_df['source_loc_name'] == loc_name].iloc[0] - except (KeyError, IndexError): - print(f"Could not create location {loc_name}. Please check mapping file.") + row = loc_lookup.get(loc_name) + if row is None: + logger.error( + "Could not create location %s. Please check mapping file.", loc_name + ) + continue + url = row["target_url"] + access_pt = row.get("target_access_pt", "") + + if url is None: + logger.error( + "Could not create location %s. Please check mapping file.", loc_name + ) + continue + + if config.dry_run: + logger.info("[DRY RUN] Would create external location %s", loc_name) continue if access_pt: - w_target.external_locations.create(name=loc_name, - credential_name=loc_cred_name, - comment=loc_comment, - fallback=loc_fallback, - read_only=loc_read_only, - url=url, - access_point=access_pt) + w_target.external_locations.create( + name=loc_name, + credential_name=loc_cred_name, + comment=loc_comment, + fallback=loc_fallback, + read_only=loc_read_only, + url=url, + access_point=access_pt, + ) else: - w_target.external_locations.create(name=loc_name, - credential_name=loc_cred_name, - comment=loc_comment, - fallback=loc_fallback, - read_only=loc_read_only, - url=url) + w_target.external_locations.create( + name=loc_name, + credential_name=loc_cred_name, + comment=loc_comment, + fallback=loc_fallback, + read_only=loc_read_only, + url=url, + ) elif cloud_type == "azure": - try: - url = loc_df['target_url'].loc[loc_df['source_loc_name'] == loc_name].iloc[0] - except (KeyError, IndexError): - print(f"Could not create location {loc_name}. Please check mapping file.") + row = loc_lookup.get(loc_name) + if row is None: + logger.error( + "Could not create location %s. Please check mapping file.", loc_name + ) + continue + url = row["target_url"] + + if url is None: + logger.error( + "Could not create location %s. Please check mapping file.", loc_name + ) continue - w_target.external_locations.create(name=loc_name, - credential_name=loc_cred_name, - comment=loc_comment, - fallback=loc_fallback, - read_only=loc_read_only, - url=url) + if config.dry_run: + logger.info("[DRY RUN] Would create external location %s", loc_name) + continue + + w_target.external_locations.create( + name=loc_name, + credential_name=loc_cred_name, + comment=loc_comment, + fallback=loc_fallback, + read_only=loc_read_only, + url=url, + ) elif cloud_type == "gcp": - print("GCP not yet implemented.") + logger.warning("GCP not yet implemented.") continue else: - print("Cloud type must be one of AWS, GCP, or Azure.") + logger.error("Cloud type must be one of AWS, GCP, or Azure.") continue - print(f"External location {loc_name} created.") + logger.info("External location %s created.", loc_name) + +if __name__ == "__main__": + parser = argparse.ArgumentParser( + description="Sync storage credentials and external locations between workspaces" + ) + parser.add_argument( + "--dry-run", + action="store_true", + help="Show planned operations without executing", + ) + parser.add_argument( + "--log-level", + default="INFO", + choices=["DEBUG", "INFO", "WARNING", "ERROR"], + help="Set logging level", + ) + args = parser.parse_args() + config.dry_run = args.dry_run + logger = setup_logging(level=args.log_level) diff --git a/sync_ext_volumes.py b/sync_ext_volumes.py index eb3bc37..076b218 100644 --- a/sync_ext_volumes.py +++ b/sync_ext_volumes.py @@ -16,39 +16,78 @@ # -num_exec: the number of threads to spawn in the ThreadPoolExecutor. +import argparse +import os from itertools import repeat from databricks.sdk import WorkspaceClient from databricks.sdk.service import catalog from concurrent.futures import ThreadPoolExecutor from databricks.sdk.errors.platform import ResourceAlreadyExists -from common import (target_pat, target_host, - source_pat, source_host, - catalogs_to_copy, num_exec) +from dr_sync.config import DRSyncConfig +from dr_sync.log import setup_logging + +config = ( + DRSyncConfig.from_env() + if os.environ.get("DR_SYNC_SOURCE_HOST") + else DRSyncConfig.from_common_module() +) +logger = setup_logging() +target_host = config.target_host +target_pat = config.target_token +source_host = config.source_host +source_pat = config.source_token +catalogs_to_copy = config.catalogs_to_copy +num_exec = config.num_exec # helper function to create volumes and set appropriate owner def create_volume(w, catalog_name, schema_name, volume_name, location, owner): - print(f"Creating volume {volume_name} in {catalog_name}.{schema_name}...") + logger.info( + "Creating volume %s in %s.%s...", volume_name, catalog_name, schema_name + ) + + # dry-run guard: log what would be created without executing + if config.dry_run: + logger.info( + "[DRY RUN] Would create volume %s in %s.%s", + volume_name, + catalog_name, + schema_name, + ) + return { + "volume": f"{catalog_name}.{schema_name}.{volume_name}", + "status": "dry_run", + } # try creating new volume try: - volume = w.volumes.create(catalog_name=catalog_name, - schema_name=schema_name, - name=volume_name, - storage_location=location, - volume_type=catalog.VolumeType.EXTERNAL) + volume = w.volumes.create( + catalog_name=catalog_name, + schema_name=schema_name, + name=volume_name, + storage_location=location, + volume_type=catalog.VolumeType.EXTERNAL, + ) _ = w.volumes.update(name=volume.full_name, owner=owner) return {"volume": volume.full_name, "status": "success"} # if volume already exists, just update the owner (in case it has changed) except ResourceAlreadyExists: - _ = w.volumes.update(name=f"{catalog_name}.{schema_name}.{volume_name}", owner=owner) - return {"volume": f"{catalog_name}.{schema_name}.{volume_name}", "status": "already_exists"} + _ = w.volumes.update( + name=f"{catalog_name}.{schema_name}.{volume_name}", owner=owner + ) + return { + "volume": f"{catalog_name}.{schema_name}.{volume_name}", + "status": "already_exists", + } # for any other exception, return the error except Exception as e: - return {"volume": f"{catalog_name}.{schema_name}.{volume_name}", "status": f"ERROR: {e}"} + return { + "volume": f"{catalog_name}.{schema_name}.{volume_name}", + "status": f"ERROR: {e}", + } # create the WorkspaceClient pointed at the target WS @@ -64,29 +103,57 @@ def create_volume(w, catalog_name, schema_name, volume_name, location, owner): # but the owner has changed. for cat in catalogs_to_copy: filtered_volumes = system_info.filter( - (system_info.volume_catalog == cat) & - (system_info.volume_schema != "information_schema") & - (system_info.volume_type == "EXTERNAL")).collect() + (system_info.volume_catalog == cat) + & (system_info.volume_schema != "information_schema") + & (system_info.volume_type == "EXTERNAL") + ).collect() # get schemas, tables and locations in list form - schema_names = [row['volume_schema'] for row in filtered_volumes] - volume_names = [row['volume_name'] for row in filtered_volumes] - volume_locs = [row['storage_location'] for row in filtered_volumes] - volume_owners = [row['volume_owner'] for row in filtered_volumes] + schema_names = [row["volume_schema"] for row in filtered_volumes] + volume_names = [row["volume_name"] for row in filtered_volumes] + volume_locs = [row["storage_location"] for row in filtered_volumes] + volume_owners = [row["volume_owner"] for row in filtered_volumes] with ThreadPoolExecutor(max_workers=num_exec) as executor: - threads = executor.map(create_volume, - repeat(w_target), - repeat(cat), - schema_names, - volume_names, - volume_locs, - volume_owners) + threads = executor.map( + create_volume, + repeat(w_target), + repeat(cat), + schema_names, + volume_names, + volume_locs, + volume_owners, + ) for thread in threads: if thread["status"] == "success": - print("Created volume {}.".format(thread["volume"])) + logger.info("Created volume %s.", thread["volume"]) elif thread["status"] == "already_exists": - print("Skipped volume {} because it already exists.".format(thread["volume"])) + logger.warning( + "Skipped volume %s because it already exists.", thread["volume"] + ) else: - print("Could not create volume {}; error: {}".format(thread["volume"], thread["status"])) + logger.error( + "Could not create volume %s; error: %s", + thread["volume"], + thread["status"], + ) + +if __name__ == "__main__": + parser = argparse.ArgumentParser( + description="Sync external volumes between workspaces" + ) + parser.add_argument( + "--dry-run", + action="store_true", + help="Show planned operations without executing", + ) + parser.add_argument( + "--log-level", + default="INFO", + choices=["DEBUG", "INFO", "WARNING", "ERROR"], + help="Set logging level", + ) + args = parser.parse_args() + config.dry_run = args.dry_run + logger = setup_logging(level=args.log_level) diff --git a/sync_grs_ext.py b/sync_grs_ext.py index d43e205..6198937 100644 --- a/sync_grs_ext.py +++ b/sync_grs_ext.py @@ -23,103 +23,78 @@ # warehouse. All table load statuses will be written to the delta table at {target_bucket}/sync_status_{time.time_ns()}. +import argparse +import os import time import pandas as pd from itertools import repeat from databricks.sdk import WorkspaceClient -from databricks.sdk.service import sql as dbsql from concurrent.futures import ThreadPoolExecutor -from databricks.sdk.service.sql import Disposition -from databricks.sdk.service.sql import StatementState -from databricks.sdk.service.sql import CreateWarehouseRequestWarehouseType -from databricks.sdk.service.sql import ExecuteStatementRequestOnWaitTimeout -from common import (target_pat, target_host, - catalogs_to_copy, num_exec, - landing_zone_url, warehouse_size, - response_backoff) - - -# helper function to drop external tables -def drop_table(w, catalog, schema, table_name, warehouse): - print(f"Dropping table {catalog}.{schema}.{table_name}...") - - try: - sqlstring = f"DROP TABLE IF EXISTS {catalog}.{schema}.{table_name}" - resp = w.statement_execution.execute_statement(warehouse_id=warehouse, - wait_timeout="0s", - on_wait_timeout=ExecuteStatementRequestOnWaitTimeout("CONTINUE"), - disposition=Disposition("EXTERNAL_LINKS"), - statement=sqlstring) - - while resp.status.state in {StatementState.PENDING, StatementState.RUNNING}: - resp = w.statement_execution.get_statement(resp.statement_id) - time.sleep(response_backoff) - - if resp.status.state != StatementState.SUCCEEDED: - return {"status": 0, - "catalog": catalog, - "schema": schema, - "table_name": table_name} - - return {"status": 1, - "catalog": catalog, - "schema": schema, - "table_name": table_name} - - except Exception: - return {"status": 0, - "catalog": catalog, - "schema": schema, - "table_name": table_name} +from dr_sync.sql_utils import ( + execute_statement_sync, + managed_warehouse, + drop_table_if_exists, +) +from dr_sync.exceptions import StatementError +from dr_sync.config import DRSyncConfig +from dr_sync.log import setup_logging + +logger = setup_logging() +config = ( + DRSyncConfig.from_env() + if os.environ.get("DR_SYNC_SOURCE_HOST") + else DRSyncConfig.from_common_module() +) +target_host = config.target_host +target_pat = config.target_token +catalogs_to_copy = config.catalogs_to_copy +num_exec = config.num_exec +landing_zone_url = config.landing_zone_url +warehouse_size = config.warehouse_size +response_backoff = config.response_backoff # helper function to load tables from a specified location def load_table(w, catalog, schema, table_name, location, warehouse): - print(f"Creating EXTERNAL table {catalog}.{schema}.{table_name}...") + logger.info("Creating EXTERNAL table %s.%s.%s...", catalog, schema, table_name) try: sqlstring = f"CREATE TABLE {catalog}.{schema}.{table_name} USING delta LOCATION '{location}'" - resp = w.statement_execution.execute_statement(warehouse_id=warehouse, - wait_timeout="0s", - on_wait_timeout=ExecuteStatementRequestOnWaitTimeout("CONTINUE"), - disposition=Disposition("EXTERNAL_LINKS"), - statement=sqlstring) - - while resp.status.state in {StatementState.PENDING, StatementState.RUNNING}: - resp = w.statement_execution.get_statement(resp.statement_id) - time.sleep(response_backoff) - - if resp.status.state != StatementState.SUCCEEDED: - return {"catalog": catalog, - "schema": schema, - "table_name": table_name, - "location": location, - "status": f"FAIL: {resp.status.error.message}", - "creation_time": time.time_ns()} - - return {"catalog": catalog, - "schema": schema, - "table_name": table_name, - "location": location, - "status": "SUCCESS", - "creation_time": time.time_ns()} + execute_statement_sync(w, warehouse, sqlstring, backoff=response_backoff) + + return { + "catalog": catalog, + "schema": schema, + "table_name": table_name, + "location": location, + "status": "SUCCESS", + "creation_time": time.time_ns(), + } + + except StatementError as e: + return { + "catalog": catalog, + "schema": schema, + "table_name": table_name, + "location": location, + "status": f"FAIL: {e}", + "creation_time": time.time_ns(), + } except Exception as e: - return {"catalog": catalog, - "schema": schema, - "table_name": table_name, - "location": location, - "status": "FAIL: {e}", - "creation_time": time.time_ns()} - + return { + "catalog": catalog, + "schema": schema, + "table_name": table_name, + "location": location, + "status": f"FAIL: {e}", + "creation_time": time.time_ns(), + } -# other parameters -wh_type = CreateWarehouseRequestWarehouseType("PRO") # required for serverless warehouse # initialize lists for status tracking loaded_table_names = [] -loaded_table_types = [] loaded_table_schemas = [] loaded_table_catalogs = [] loaded_table_locations = [] @@ -129,78 +104,145 @@ def load_table(w, catalog, schema, table_name, location, warehouse): # create the WorkspaceClient pointed at the target WS w_target = WorkspaceClient(host=target_host, token=target_pat) -# create warehouse to run table creation statements -wh_target = w_target.warehouses.create(name=f'sdk-{time.time_ns()}', - cluster_size=warehouse_size, - max_num_clusters=1, - auto_stop_mins=10, - warehouse_type=wh_type, - enable_serverless_compute=True, - tags=dbsql.EndpointTags( - custom_tags=[ - dbsql.EndpointTagPair(key="Owner", value="dr-sync-tool")])).result() - -system_info = spark.sql("SELECT * FROM system.information_schema.tables") - -# loop through all catalogs to copy, then copy all tables excluding system tables. -# we also skip views; these need to be created separately since they cannot be cloned. -for cat in catalogs_to_copy: - filtered_tables = system_info.filter( - (system_info.table_catalog == cat) & - (system_info.table_schema != "information_schema") & - (system_info.table_type == "EXTERNAL")).collect() - - # get schemas, tables and types in list form - schemas = [row['table_schema'] for row in filtered_tables] - table_names = [row['table_name'] for row in filtered_tables] - table_locs = [row['storage_path'] for row in filtered_tables] - - with ThreadPoolExecutor(max_workers=num_exec) as executor: - threads = executor.map(drop_table, - repeat(w_target), - repeat(cat), - schemas, - table_names, - repeat(wh_target.id)) - - for thread in threads: - if thread["status"]: - print("Dropped table {}.{}.{}.".format(thread["catalog"], thread["schema"], thread["table_name"])) - else: - print( - "Error dropping table {}.{}.{}.".format(thread["catalog"], thread["schema"], thread["table_name"])) - - # use ThreadPool to copy tables in parallel - with ThreadPoolExecutor(max_workers=num_exec) as executor: - threads = executor.map(load_table, - repeat(w_target), - repeat(cat), - schemas, - table_names, - table_locs, - repeat(wh_target.id)) - - # wait for threads to execute and build lists for status table - for thread in threads: - loaded_table_names.append(thread["table_name"]) - loaded_table_schemas.append(thread["schema"]) - loaded_table_catalogs.append(thread["catalog"]) - loaded_table_locations.append(thread["location"]) - loaded_table_status.append(thread["status"]) - loaded_table_times.append(thread["creation_time"]) - print("Loaded table {}.{}.{}.".format(thread["catalog"], thread["schema"], thread["table_name"])) - -# create the table statuses as a df and write to a table in dr target -status_df = pd.DataFrame({"catalog": loaded_table_catalogs, - "schema": loaded_table_schemas, - "table": loaded_table_names, - "location": loaded_table_locations, - "type": loaded_table_types, - "status": loaded_table_status, - "create_time": loaded_table_times}) - -# table will get a specific timestamp-based location per run -(spark.createDataFrame(status_df) - .write.mode("overwrite") - .format("delta") - .save(f"{landing_zone_url}/sync_status_{time.time_ns()}")) +if config.dry_run: + # In dry-run mode, log what would happen without creating warehouses or executing SQL + for cat in catalogs_to_copy: + filtered_tables = spark.sql( + f""" + SELECT table_schema, table_name, storage_path + FROM system.information_schema.tables + WHERE table_catalog = '{cat}' + AND table_schema != 'information_schema' + AND table_type = 'EXTERNAL' + """ + ).collect() + + schemas = [row["table_schema"] for row in filtered_tables] + table_names = [row["table_name"] for row in filtered_tables] + table_locs = [row["storage_path"] for row in filtered_tables] + + logger.info( + "[DRY RUN] Would process %d external tables in catalog %s", + len(table_names), + cat, + ) + for schema, table_name, location in zip(schemas, table_names, table_locs): + logger.info( + "[DRY RUN] Would drop and recreate external table %s.%s.%s at %s", + cat, + schema, + table_name, + location, + ) +else: + # create warehouse to run table creation statements, guaranteed cleanup + with managed_warehouse(w_target, size=warehouse_size) as wh_id: + # loop through all catalogs to copy, then copy all tables excluding system tables. + # we also skip views; these need to be created separately since they cannot be cloned. + for cat in catalogs_to_copy: + filtered_tables = spark.sql( + f""" + SELECT table_schema, table_name, storage_path + FROM system.information_schema.tables + WHERE table_catalog = '{cat}' + AND table_schema != 'information_schema' + AND table_type = 'EXTERNAL' + """ + ).collect() + + # get schemas, tables and types in list form + schemas = [row["table_schema"] for row in filtered_tables] + table_names = [row["table_name"] for row in filtered_tables] + table_locs = [row["storage_path"] for row in filtered_tables] + + with ThreadPoolExecutor(max_workers=num_exec) as executor: + threads = executor.map( + drop_table_if_exists, + repeat(w_target), + repeat(wh_id), + repeat(cat), + schemas, + table_names, + ) + + for thread in threads: + if thread["status"]: + logger.info( + "Dropped table %s.%s.%s.", + thread["catalog"], + thread["schema"], + thread["table_name"], + ) + else: + logger.error( + "Error dropping table %s.%s.%s.", + thread["catalog"], + thread["schema"], + thread["table_name"], + ) + + # use ThreadPool to copy tables in parallel + with ThreadPoolExecutor(max_workers=num_exec) as executor: + threads = executor.map( + load_table, + repeat(w_target), + repeat(cat), + schemas, + table_names, + table_locs, + repeat(wh_id), + ) + + # wait for threads to execute and build lists for status table + for thread in threads: + loaded_table_names.append(thread["table_name"]) + loaded_table_schemas.append(thread["schema"]) + loaded_table_catalogs.append(thread["catalog"]) + loaded_table_locations.append(thread["location"]) + loaded_table_status.append(thread["status"]) + loaded_table_times.append(thread["creation_time"]) + logger.info( + "Loaded table %s.%s.%s.", + thread["catalog"], + thread["schema"], + thread["table_name"], + ) + + # create the table statuses as a df and write to a table in dr target + status_df = pd.DataFrame( + { + "catalog": loaded_table_catalogs, + "schema": loaded_table_schemas, + "table": loaded_table_names, + "location": loaded_table_locations, + "status": loaded_table_status, + "create_time": loaded_table_times, + } + ) + + # table will get a specific timestamp-based location per run + ( + spark.createDataFrame(status_df) + .write.mode("overwrite") + .format("delta") + .save(f"{landing_zone_url}/sync_status_{time.time_ns()}") + ) + +if __name__ == "__main__": + parser = argparse.ArgumentParser( + description="Sync GRS-replicated external tables to secondary Databricks workspace" + ) + parser.add_argument( + "--dry-run", + action="store_true", + help="Show planned operations without executing", + ) + parser.add_argument( + "--log-level", + default="INFO", + choices=["DEBUG", "INFO", "WARNING", "ERROR"], + help="Set logging level", + ) + args = parser.parse_args() + config.dry_run = args.dry_run + logger = setup_logging(level=args.log_level) diff --git a/sync_perms.py b/sync_perms.py index 1960146..7b76331 100644 --- a/sync_perms.py +++ b/sync_perms.py @@ -17,14 +17,28 @@ # -catalogs_to_copy: a list of the catalogs to be replicated between workspaces. # -num_exec: the number of threads to spawn in the ThreadPoolExecutor. +import argparse +import os from itertools import repeat from databricks.sdk.service import catalog from databricks.sdk import WorkspaceClient from concurrent.futures import ThreadPoolExecutor from databricks.sdk.errors.platform import NotFound -from common import (target_pat, target_host, - source_pat, source_host, - catalogs_to_copy, num_exec) +from dr_sync.config import DRSyncConfig +from dr_sync.log import setup_logging + +config = ( + DRSyncConfig.from_env() + if os.environ.get("DR_SYNC_SOURCE_HOST") + else DRSyncConfig.from_common_module() +) +logger = setup_logging() +target_host = config.target_host +target_pat = config.target_token +source_host = config.source_host +source_pat = config.source_token +catalogs_to_copy = config.catalogs_to_copy +num_exec = config.num_exec # helper function to update object grants between source and target WS @@ -40,23 +54,36 @@ def sync_grants(w_src, w_tgt, obj_name, obj_type): # get list of all distinct users with grants on the object user_list = {u.principal for u in source_grants.privilege_assignments}.union( - {u.principal for u in target_grants.privilege_assignments}) + {u.principal for u in target_grants.privilege_assignments} + ) # create PermissionsChange object for each user where a change exists change_list = [] for u in user_list: # get the source/target privileges; these may not exist in one or the other environment try: - source_privs = [x.privilege for x in - [p.privileges for p in source_grants.privilege_assignments if p.principal == u][0] - if x.privilege is not None] + source_privs = [ + x.privilege + for x in [ + p.privileges + for p in source_grants.privilege_assignments + if p.principal == u + ][0] + if x.privilege is not None + ] except IndexError: source_privs = [] try: - target_privs = [x.privilege for x in - [p.privileges for p in target_grants.privilege_assignments if p.principal == u][0] - if x.privilege is not None] + target_privs = [ + x.privilege + for x in [ + p.privileges + for p in target_grants.privilege_assignments + if p.principal == u + ][0] + if x.privilege is not None + ] except IndexError: target_privs = [] @@ -65,24 +92,22 @@ def sync_grants(w_src, w_tgt, obj_name, obj_type): # for the change list based on which types of changes exist if add_perms and rem_perms: - change_list.append(catalog.PermissionsChange( - add=add_perms, - remove=rem_perms, - principal=u)) + change_list.append( + catalog.PermissionsChange(add=add_perms, remove=rem_perms, principal=u) + ) elif add_perms: - change_list.append(catalog.PermissionsChange( - add=add_perms, - principal=u)) + change_list.append(catalog.PermissionsChange(add=add_perms, principal=u)) elif rem_perms: - change_list.append(catalog.PermissionsChange( - remove=rem_perms, - principal=u)) + change_list.append(catalog.PermissionsChange(remove=rem_perms, principal=u)) # if any grants changed, update the object in target if change_list: - w_tgt.grants.update(full_name=obj_name, - securable_type=obj_type, - changes=change_list) + if config.dry_run: + logger.info("[DRY RUN] Would update grants for %s (%s)", obj_name, obj_type) + return {"name": obj_name, "status": "DRY_RUN"} + w_tgt.grants.update( + full_name=obj_name, securable_type=obj_type, changes=change_list + ) return {"name": obj_name, "status": "SUCCESS"} else: return {"name": obj_name, "status": None} @@ -99,8 +124,9 @@ def sync_grants(w_src, w_tgt, obj_name, obj_type): # iterate through catalogs for cat in catalogs_to_copy: filtered_tables = table_info.filter( - (table_info.table_catalog == cat) & - (table_info.table_schema != "information_schema")).collect() + (table_info.table_catalog == cat) + & (table_info.table_schema != "information_schema") + ).collect() filtered_volumes = volume_info.filter(volume_info.volume_catalog == cat).collect() @@ -108,68 +134,115 @@ def sync_grants(w_src, w_tgt, obj_name, obj_type): res = sync_grants(w_source, w_target, cat, catalog.SecurableType.CATALOG) if res["status"] == "SUCCESS": - print(f"Synced grants for catalog {cat}.") + logger.info("Synced grants for catalog %s.", cat) elif res["status"] == "NotFound": - print(f"ERROR: catalog {cat} does not exist in target workspace. Sync metadata and re-run.") + logger.error( + "Catalog %s does not exist in target workspace. Sync metadata and re-run.", + cat, + ) else: - print(f"No changes to sync for catalog {cat}.") + logger.info("No changes to sync for catalog %s.", cat) # get list of fully qualified schemas and tables - schemas = {f"{cat}.{schema}" for schema in [row['table_schema'] for row in filtered_tables]} - table_names = [f"{cat}.{schema}.{table}" for schema, table in - zip([row['table_schema'] for row in filtered_tables], - [row['table_name'] for row in filtered_tables])] - volume_names = [f"{cat}.{schema}.{table}" for schema, table in - zip([row['volume_schema'] for row in filtered_volumes], - [row['volume_name'] for row in filtered_volumes])] + schemas = { + f"{cat}.{schema}" for schema in [row["table_schema"] for row in filtered_tables] + } + table_names = [ + f"{cat}.{schema}.{table}" + for schema, table in zip( + [row["table_schema"] for row in filtered_tables], + [row["table_name"] for row in filtered_tables], + ) + ] + volume_names = [ + f"{cat}.{schema}.{table}" + for schema, table in zip( + [row["volume_schema"] for row in filtered_volumes], + [row["volume_name"] for row in filtered_volumes], + ) + ] # update schema grants in parallel with ThreadPoolExecutor(max_workers=num_exec) as executor: - threads = executor.map(sync_grants, - repeat(w_source), - repeat(w_target), - schemas, - repeat(catalog.SecurableType.SCHEMA)) + threads = executor.map( + sync_grants, + repeat(w_source), + repeat(w_target), + schemas, + repeat(catalog.SecurableType.SCHEMA), + ) for thread in threads: name = thread["name"] if thread["status"] == "SUCCESS": - print(f"Synced grants for schema {name}.") + logger.info("Synced grants for schema %s.", name) elif thread["status"] == "NotFound": - print(f"ERROR: schema {name} does not exist in target workspace. Sync metadata and re-run.") + logger.error( + "Schema %s does not exist in target workspace. Sync metadata and re-run.", + name, + ) else: - print(f"No changes to sync for schema {name}.") + logger.info("No changes to sync for schema %s.", name) # update table grants in parallel with ThreadPoolExecutor(max_workers=num_exec) as executor: - threads = executor.map(sync_grants, - repeat(w_source), - repeat(w_target), - table_names, - repeat(catalog.SecurableType.TABLE)) + threads = executor.map( + sync_grants, + repeat(w_source), + repeat(w_target), + table_names, + repeat(catalog.SecurableType.TABLE), + ) for thread in threads: name = thread["name"] if thread["status"] == "SUCCESS": - print(f"Synced grants for table {name}.") + logger.info("Synced grants for table %s.", name) elif thread["status"] == "NotFound": - print(f"ERROR: table {name} does not exist in target workspace. Sync metadata and re-run.") + logger.error( + "Table %s does not exist in target workspace. Sync metadata and re-run.", + name, + ) else: - print(f"No changes to sync for table {name}.") + logger.info("No changes to sync for table %s.", name) # update volume grants in parallel with ThreadPoolExecutor(max_workers=num_exec) as executor: - threads = executor.map(sync_grants, - repeat(w_source), - repeat(w_target), - volume_names, - repeat(catalog.SecurableType.VOLUME)) + threads = executor.map( + sync_grants, + repeat(w_source), + repeat(w_target), + volume_names, + repeat(catalog.SecurableType.VOLUME), + ) for thread in threads: name = thread["name"] if thread["status"] == "SUCCESS": - print(f"Synced grants for volume {name}.") + logger.info("Synced grants for volume %s.", name) elif thread["status"] == "NotFound": - print(f"ERROR: volume {name} does not exist in target workspace. Sync volumes and re-run.") + logger.error( + "Volume %s does not exist in target workspace. Sync volumes and re-run.", + name, + ) else: - print(f"No changes to sync for volume {name}.") + logger.info("No changes to sync for volume %s.", name) + +if __name__ == "__main__": + parser = argparse.ArgumentParser( + description="Sync permissions (grants) between workspaces" + ) + parser.add_argument( + "--dry-run", + action="store_true", + help="Show planned operations without executing", + ) + parser.add_argument( + "--log-level", + default="INFO", + choices=["DEBUG", "INFO", "WARNING", "ERROR"], + help="Set logging level", + ) + args = parser.parse_args() + config.dry_run = args.dry_run + logger = setup_logging(level=args.log_level) diff --git a/sync_shared_tables.py b/sync_shared_tables.py index 9ed516b..25d7327 100644 --- a/sync_shared_tables.py +++ b/sync_shared_tables.py @@ -20,115 +20,135 @@ # -num_exec: the number of threads to spawn in the ThreadPoolExecutor. # -target_share_id: the sharing identifier of the secondary metastore. -import sys +import argparse +import os import time import pandas as pd from itertools import repeat from databricks.sdk import WorkspaceClient -from databricks.sdk.service import sql as dbsql from concurrent.futures import ThreadPoolExecutor -from databricks.sdk.service.sql import (Disposition, StatementState, - CreateWarehouseRequestWarehouseType, ExecuteStatementRequestOnWaitTimeout) from databricks.sdk.errors.platform import BadRequest from databricks.sdk.service.catalog import Privilege, PermissionsChange -from databricks.sdk.service.sharing import (AuthenticationType, SharedDataObjectUpdate, - SharedDataObjectUpdateAction, SharedDataObject, - SharedDataObjectDataObjectType, SharedDataObjectStatus) -from common import (target_pat, target_host, - source_pat, source_host, - catalogs_to_copy, num_exec, - landing_zone_url, warehouse_size, - response_backoff, metastore_id) +from databricks.sdk.service.sharing import ( + AuthenticationType, + SharedDataObjectUpdate, + SharedDataObjectUpdateAction, + SharedDataObject, + SharedDataObjectDataObjectType, + SharedDataObjectStatus, +) +from dr_sync.sql_utils import execute_statement_sync, managed_warehouse +from dr_sync.exceptions import StatementError +from dr_sync.config import DRSyncConfig +from dr_sync.log import setup_logging + +config = ( + DRSyncConfig.from_env() + if os.environ.get("DR_SYNC_SOURCE_HOST") + else DRSyncConfig.from_common_module() +) +logger = setup_logging() +target_host = config.target_host +target_pat = config.target_token +source_host = config.source_host +source_pat = config.source_token +catalogs_to_copy = config.catalogs_to_copy +num_exec = config.num_exec +landing_zone_url = config.landing_zone_url +warehouse_size = config.warehouse_size +response_backoff = config.response_backoff +metastore_id = config.metastore_id # helper function to clone a table from one catalog to another def clone_table(w, source_catalog, target_catalog, schema, table_name, warehouse): - print(f"Cloning table {source_catalog}.{schema}.{table_name}...") + logger.info("Cloning table %s.%s.%s...", source_catalog, schema, table_name) try: - sqlstring = (f"CREATE OR REPLACE TABLE {target_catalog}.{schema}.{table_name} " - f"DEEP CLONE {source_catalog}.{schema}.{table_name}") - - resp = w.statement_execution.execute_statement(warehouse_id=warehouse, - wait_timeout="0s", - on_wait_timeout=ExecuteStatementRequestOnWaitTimeout( - "CONTINUE"), - disposition=Disposition("EXTERNAL_LINKS"), - statement=sqlstring) - - while resp.status.state in {StatementState.PENDING, StatementState.RUNNING}: - resp = w.statement_execution.get_statement(resp.statement_id) - time.sleep(response_backoff) - - if resp.status.state != StatementState.SUCCEEDED: - return {"catalog": target_catalog, - "schema": schema, - "table_name": table_name, - "status": f"FAIL: {resp.status.error.message}", - "creation_time": time.time_ns()} - - return {"catalog": target_catalog, - "schema": schema, - "table_name": table_name, - "status": "SUCCESS", - "creation_time": time.time_ns()} + sqlstring = ( + f"CREATE OR REPLACE TABLE {target_catalog}.{schema}.{table_name} " + f"DEEP CLONE {source_catalog}.{schema}.{table_name}" + ) + + execute_statement_sync(w, warehouse, sqlstring, backoff=response_backoff) + + return { + "catalog": target_catalog, + "schema": schema, + "table_name": table_name, + "status": "SUCCESS", + "creation_time": time.time_ns(), + } + + except StatementError as e: + return { + "catalog": target_catalog, + "schema": schema, + "table_name": table_name, + "status": f"FAIL: {e}", + "creation_time": time.time_ns(), + } except Exception as e: - return {"catalog": target_catalog, - "schema": schema, - "table_name": table_name, - "status": f"FAIL: {e}", - "creation_time": time.time_ns()} + return { + "catalog": target_catalog, + "schema": schema, + "table_name": table_name, + "status": f"FAIL: {e}", + "creation_time": time.time_ns(), + } # other parameters -wh_type = CreateWarehouseRequestWarehouseType("PRO") # required for serverless warehouse write_results = False # set to true to write status df to disk # create the WorkspaceClients for source and target workspaces w_source = WorkspaceClient(host=source_host, token=source_pat) w_target = WorkspaceClient(host=target_host, token=target_pat) -# create warehouse in secondary to run table creation statements -print("Creating warehouse in secondary workspace...") -wh_target = w_target.warehouses.create(name=f'sdk-{time.time_ns()}', - cluster_size=warehouse_size, - max_num_clusters=1, - auto_stop_mins=10, - warehouse_type=wh_type, - enable_serverless_compute=True, - tags=dbsql.EndpointTags( - custom_tags=[ - dbsql.EndpointTagPair(key="Owner", value="dr-sync-tool")])).result() - # create the secondary metastore as a recipient try: - print(f"Creating recipient with id {metastore_id}...") - recipient = w_source.recipients.create(name="dr_automation_recipient", - authentication_type=AuthenticationType.DATABRICKS, - data_recipient_global_metastore_id=metastore_id) + logger.info("Creating recipient with id %s...", metastore_id) + recipient = w_source.recipients.create( + name="dr_automation_recipient", + authentication_type=AuthenticationType.DATABRICKS, + data_recipient_global_metastore_id=metastore_id, + ) except BadRequest: try: - recipient = [r for r in w_source.recipients.list() if r.data_recipient_global_metastore_id == metastore_id][0] - print(f"Recipient with id {metastore_id} already exists. Skipping creation...") + recipient = [ + r + for r in w_source.recipients.list() + if r.data_recipient_global_metastore_id == metastore_id + ][0] + logger.info( + "Recipient with id %s already exists. Skipping creation...", metastore_id + ) except IndexError: - print(f"Recipient with id {metastore_id} does not exist in source workspace. Please validate the id and create it manually.") - sys.exit() + raise RuntimeError( + f"Recipient with id {metastore_id} does not exist in source workspace. Please validate the id and create it manually." + ) # get all tables in the primary metastore system_info = spark.sql("SELECT * FROM system.information_schema.tables") # get local metastore id -local_metastore_id = [r["current_metastore()"] for r in spark.sql("SELECT current_metastore()").collect()][0] +local_metastore_id = [ + r["current_metastore()"] for r in spark.sql("SELECT current_metastore()").collect() +][0] # get remote provider name; it may or may not be the same as local_metastore_id try: - remote_provider_name = [p.name for p in w_target.providers.list() if - p.data_provider_global_metastore_id == local_metastore_id][0] + remote_provider_name = [ + p.name + for p in w_target.providers.list() + if p.data_provider_global_metastore_id == local_metastore_id + ][0] except IndexError: - print("Provider could not be found in target workspace; please check that it was created.") - sys.exit() + raise RuntimeError( + "Provider could not be found in target workspace; please check that it was created." + ) # initalize df lists cloned_table_names = [] @@ -137,83 +157,183 @@ def clone_table(w, source_catalog, target_catalog, schema, table_name, warehouse cloned_table_status = [] cloned_table_times = [] -# iterate through all catalogs to share -for cat in catalogs_to_copy: - filtered_tables = system_info.filter( - (system_info.table_catalog == cat) & - (system_info.table_schema != "information_schema") & - (system_info.table_type != "VIEW")).distinct().collect() +if config.dry_run: + # In dry-run mode, log what would happen without creating warehouses or executing SQL + for cat in catalogs_to_copy: + filtered_tables = ( + system_info.filter( + (system_info.table_catalog == cat) + & (system_info.table_schema != "information_schema") + & (system_info.table_type != "VIEW") + ) + .distinct() + .collect() + ) - unique_schemas = {row['table_schema'] for row in filtered_tables} - all_tables = [row["table_name"] for row in filtered_tables] - all_schemas = [row["table_schema"] for row in filtered_tables] + unique_schemas = {row["table_schema"] for row in filtered_tables} + all_tables = [row["table_name"] for row in filtered_tables] + all_schemas = [row["table_schema"] for row in filtered_tables] - # create the share for the current catalog and update permissions - print(f"Creating share for catalog {cat}...") - try: - share = w_source.shares.create(name=f"{cat}_share") - share_name = share.name - except BadRequest: - print(f"Share {cat}_share already exists. Skipping creation...") - share_name = f"{cat}_share" + logger.info( + "[DRY RUN] Would create/update share %s_share with %d schemas", + cat, + len(unique_schemas), + ) + for schema in unique_schemas: + logger.info("[DRY RUN] Would add schema %s.%s to share", cat, schema) + logger.info( + "[DRY RUN] Would create shared catalog %s_share in target workspace", cat + ) + logger.info( + "[DRY RUN] Would clone %d tables from %s_share to %s", + len(all_tables), + cat, + cat, + ) + for schema, table_name in zip(all_schemas, all_tables): + logger.info( + "[DRY RUN] Would clone table %s_share.%s.%s to %s.%s.%s", + cat, + schema, + table_name, + cat, + schema, + table_name, + ) +else: + # create warehouse in secondary to run table creation statements, guaranteed cleanup + logger.info("Creating warehouse in secondary workspace...") + with managed_warehouse(w_target, size=warehouse_size) as wh_id: + # iterate through all catalogs to share + for cat in catalogs_to_copy: + filtered_tables = ( + system_info.filter( + (system_info.table_catalog == cat) + & (system_info.table_schema != "information_schema") + & (system_info.table_type != "VIEW") + ) + .distinct() + .collect() + ) - try: - _ = w_source.shares.update_permissions(share_name, - changes=[PermissionsChange(add=[Privilege.SELECT], - principal=recipient.name)]) - except BadRequest: - print(f"Could not update permissions for share {share_name}.") - - # build update object with all schemas in the current catalog - updates = [ - SharedDataObjectUpdate(action=SharedDataObjectUpdateAction.ADD, - data_object=SharedDataObject(name=f"{cat}.{schema}", - data_object_type=SharedDataObjectDataObjectType.SCHEMA, - status=SharedDataObjectStatus.ACTIVE)) - for schema in unique_schemas] - - # update the share - try: - _ = w_source.shares.update(share_name, updates=updates) - except Exception as e: - print(f"Error updating share {share_name}: {e}") + unique_schemas = {row["table_schema"] for row in filtered_tables} + all_tables = [row["table_name"] for row in filtered_tables] + all_schemas = [row["table_schema"] for row in filtered_tables] - # create the shared catalog in the target workspace - try: - _ = w_target.catalogs.create(name=f"{cat}_share", provider_name=remote_provider_name, share_name=share_name) - except BadRequest: - print(f"Shared catalog {cat}_share already exists. Skipping creation.") - - with ThreadPoolExecutor(max_workers=num_exec) as executor: - threads = executor.map(clone_table, - repeat(w_target), - repeat(f"{cat}_share"), - repeat(cat), - all_schemas, - all_tables, - repeat(wh_target.id)) - - for thread in threads: - cloned_table_names.append(thread["table_name"]) - cloned_table_schemas.append(thread["schema"]) - cloned_table_catalogs.append(thread["catalog"]) - cloned_table_status.append(thread["status"]) - cloned_table_times.append(thread["creation_time"]) - - if thread["status"] == "SUCCESS": - print("Loaded table {}.{}.{}.".format(thread["catalog"], thread["schema"], thread["table_name"])) - -# create the table statuses as a df and write to a table in dr target -status_df = pd.DataFrame({"catalog": cloned_table_catalogs, - "schema": cloned_table_schemas, - "table": cloned_table_names, - "status": cloned_table_status, - "sync_time": cloned_table_times}) - -# table will get a specific timestamp-based location per run -if write_results: - ts2 = time.time_ns() - (spark.createDataFrame(status_df) - .write.mode("overwrite") - .format("delta") - .save(f"{landing_zone_url}/sync_status_{ts2}")) + # create the share for the current catalog and update permissions + logger.info("Creating share for catalog %s...", cat) + try: + share = w_source.shares.create(name=f"{cat}_share") + share_name = share.name + except BadRequest: + logger.info("Share %s_share already exists. Skipping creation...", cat) + share_name = f"{cat}_share" + + try: + _ = w_source.shares.update_permissions( + share_name, + changes=[ + PermissionsChange( + add=[Privilege.SELECT], principal=recipient.name + ) + ], + ) + except BadRequest: + logger.error("Could not update permissions for share %s.", share_name) + + # build update object with all schemas in the current catalog + updates = [ + SharedDataObjectUpdate( + action=SharedDataObjectUpdateAction.ADD, + data_object=SharedDataObject( + name=f"{cat}.{schema}", + data_object_type=SharedDataObjectDataObjectType.SCHEMA, + status=SharedDataObjectStatus.ACTIVE, + ), + ) + for schema in unique_schemas + ] + + # update the share + try: + _ = w_source.shares.update(share_name, updates=updates) + except Exception as e: + logger.error("Error updating share %s: %s", share_name, e) + + # create the shared catalog in the target workspace + try: + _ = w_target.catalogs.create( + name=f"{cat}_share", + provider_name=remote_provider_name, + share_name=share_name, + ) + except BadRequest: + logger.info( + "Shared catalog %s_share already exists. Skipping creation.", cat + ) + + with ThreadPoolExecutor(max_workers=num_exec) as executor: + threads = executor.map( + clone_table, + repeat(w_target), + repeat(f"{cat}_share"), + repeat(cat), + all_schemas, + all_tables, + repeat(wh_id), + ) + + for thread in threads: + cloned_table_names.append(thread["table_name"]) + cloned_table_schemas.append(thread["schema"]) + cloned_table_catalogs.append(thread["catalog"]) + cloned_table_status.append(thread["status"]) + cloned_table_times.append(thread["creation_time"]) + + if thread["status"] == "SUCCESS": + logger.info( + "Loaded table %s.%s.%s.", + thread["catalog"], + thread["schema"], + thread["table_name"], + ) + + # create the table statuses as a df and write to a table in dr target + status_df = pd.DataFrame( + { + "catalog": cloned_table_catalogs, + "schema": cloned_table_schemas, + "table": cloned_table_names, + "status": cloned_table_status, + "sync_time": cloned_table_times, + } + ) + + # table will get a specific timestamp-based location per run + if write_results: + ts2 = time.time_ns() + ( + spark.createDataFrame(status_df) + .write.mode("overwrite") + .format("delta") + .save(f"{landing_zone_url}/sync_status_{ts2}") + ) + +if __name__ == "__main__": + parser = argparse.ArgumentParser( + description="Sync tables via Delta Sharing to secondary Databricks workspace" + ) + parser.add_argument( + "--dry-run", + action="store_true", + help="Show planned operations without executing", + ) + parser.add_argument( + "--log-level", + default="INFO", + choices=["DEBUG", "INFO", "WARNING", "ERROR"], + help="Set logging level", + ) + args = parser.parse_args() + config.dry_run = args.dry_run + logger = setup_logging(level=args.log_level) diff --git a/sync_tables.py b/sync_tables.py index 994edd5..87b40b8 100644 --- a/sync_tables.py +++ b/sync_tables.py @@ -27,191 +27,149 @@ # warehouse. Table load statuses will be written to the delta table at {landing_zone_url}/sync_status_{time.time_ns()}. +import argparse +import os import time import pandas as pd from itertools import repeat from databricks.sdk import WorkspaceClient -from databricks.sdk.service import sql as dbsql from concurrent.futures import ThreadPoolExecutor -from databricks.sdk.service.sql import Disposition -from databricks.sdk.service.sql import StatementState -from databricks.sdk.service.sql import CreateWarehouseRequestWarehouseType -from databricks.sdk.service.sql import ExecuteStatementRequestOnWaitTimeout -from common import (target_pat, target_host, - source_pat, source_host, - catalogs_to_copy, num_exec, - landing_zone_url, warehouse_size, - response_backoff, manifest_name) +from dr_sync.sql_utils import ( + execute_statement_sync, + managed_warehouse, + drop_table_if_exists, +) +from dr_sync.exceptions import StatementError +from dr_sync.config import DRSyncConfig +from dr_sync.log import setup_logging + +logger = setup_logging() +config = ( + DRSyncConfig.from_env() + if os.environ.get("DR_SYNC_SOURCE_HOST") + else DRSyncConfig.from_common_module() +) +target_host = config.target_host +target_pat = config.target_token +source_host = config.source_host +source_pat = config.source_token +catalogs_to_copy = config.catalogs_to_copy +num_exec = config.num_exec +landing_zone_url = config.landing_zone_url +warehouse_size = config.warehouse_size +response_backoff = config.response_backoff +manifest_name = config.manifest_name # helper function to copy tables def copy_table(w, catalog, schema, table_name, table_type, bucket, warehouse): try: sqlstring = f"CREATE OR REPLACE TABLE delta.`{bucket}/{catalog}_{schema}_{table_name}` DEEP CLONE {catalog}.{schema}.{table_name}" - resp = w.statement_execution.execute_statement(warehouse_id=warehouse, - wait_timeout="0s", - on_wait_timeout=ExecuteStatementRequestOnWaitTimeout("CONTINUE"), - disposition=Disposition("EXTERNAL_LINKS"), - statement=sqlstring) - - while resp.status.state in {StatementState.PENDING, StatementState.RUNNING}: - resp = w.statement_execution.get_statement(resp.statement_id) - time.sleep(response_backoff) - - if resp.status.state != StatementState.SUCCEEDED: - return {"catalog": catalog, - "schema": schema, - "table_name": table_name, - "table_type": f"COPY_ERROR: {resp.status.error.message}", - "location": "N/A"} + execute_statement_sync(w, warehouse, sqlstring, backoff=response_backoff) # return the table params in dict; used to build manifest - return {"catalog": catalog, - "schema": schema, - "table_name": table_name, - "table_type": table_type, - "location": bucket} + return { + "catalog": catalog, + "schema": schema, + "table_name": table_name, + "table_type": table_type, + "location": bucket, + } + + except StatementError as e: + return { + "catalog": catalog, + "schema": schema, + "table_name": table_name, + "table_type": f"COPY_ERROR: {e}", + "location": "N/A", + } except Exception as e: - return {"catalog": catalog, - "schema": schema, - "table_name": table_name, - "table_type": f"COPY_ERROR: {e}", - "location": "N/A"} - - -# helper function to drop external tables -def drop_table(w, catalog, schema, table_name, warehouse): - print(f"Dropping table {catalog}.{schema}.{table_name}...") - - try: - sqlstring = f"DROP TABLE IF EXISTS {catalog}.{schema}.{table_name}" - resp = w.statement_execution.execute_statement(warehouse_id=warehouse, - wait_timeout="0s", - on_wait_timeout=ExecuteStatementRequestOnWaitTimeout("CONTINUE"), - disposition=Disposition("EXTERNAL_LINKS"), - statement=sqlstring) - - while resp.status.state in {StatementState.PENDING, StatementState.RUNNING}: - resp = w.statement_execution.get_statement(resp.statement_id) - time.sleep(response_backoff) - - if resp.status.state != StatementState.SUCCEEDED: - return {"status": 0, - "catalog": catalog, - "schema": schema, - "table_name": table_name} - - return {"status": 1, - "catalog": catalog, - "schema": schema, - "table_name": table_name} - - except Exception: - return {"status": 0, - "catalog": catalog, - "schema": schema, - "table_name": table_name} + return { + "catalog": catalog, + "schema": schema, + "table_name": table_name, + "table_type": f"COPY_ERROR: {e}", + "location": "N/A", + } # helper function to load tables from a specified location def load_table(w, catalog, schema, table_name, table_type, location, warehouse): if table_type == "MANAGED": - print(f"Creating MANAGED table {catalog}.{schema}.{table_name}...") + logger.info("Creating MANAGED table %s.%s.%s...", catalog, schema, table_name) try: sqlstring = f"CREATE OR REPLACE TABLE {catalog}.{schema}.{table_name} DEEP CLONE delta.`{location}`" - resp = w.statement_execution.execute_statement(warehouse_id=warehouse, - wait_timeout="0s", - on_wait_timeout=ExecuteStatementRequestOnWaitTimeout( - "CONTINUE"), - disposition=Disposition("EXTERNAL_LINKS"), - statement=sqlstring) - - while resp.status.state in {StatementState.PENDING, StatementState.RUNNING}: - resp = w.statement_execution.get_statement(resp.statement_id) - time.sleep(response_backoff) - - if resp.status.state != StatementState.SUCCEEDED: - return {"catalog": catalog, - "schema": schema, - "table_name": table_name, - "table_type": table_type, - "location": location, - "status": f"FAIL: {resp.status.error.message}", - "creation_time": time.time_ns()} - - return {"catalog": catalog, - "schema": schema, - "table_name": table_name, - "table_type": table_type, - "location": location, - "status": "SUCCESS", - "creation_time": time.time_ns()} + execute_statement_sync(w, warehouse, sqlstring, backoff=response_backoff) + + return { + "catalog": catalog, + "schema": schema, + "table_name": table_name, + "table_type": table_type, + "location": location, + "status": "SUCCESS", + "creation_time": time.time_ns(), + } except Exception as e: - return {"catalog": catalog, - "schema": schema, - "table_name": table_name, - "table_type": table_type, - "location": location, - "status": f"FAIL: {e}", - "creation_time": time.time_ns()} + return { + "catalog": catalog, + "schema": schema, + "table_name": table_name, + "table_type": table_type, + "location": location, + "status": f"FAIL: {e}", + "creation_time": time.time_ns(), + } elif table_type == "EXTERNAL": - print(f"Creating EXTERNAL table {catalog}.{schema}.{table_name}...") + logger.info("Creating EXTERNAL table %s.%s.%s...", catalog, schema, table_name) try: # must drop table if it exists; CREATE_OR_REPLACE does not work when specifying external location sqlstring = f"CREATE TABLE {catalog}.{schema}.{table_name} USING delta LOCATION '{location}'" - resp = w.statement_execution.execute_statement(warehouse_id=warehouse, - wait_timeout="0s", - on_wait_timeout=ExecuteStatementRequestOnWaitTimeout("CONTINUE"), - disposition=Disposition("EXTERNAL_LINKS"), - statement=sqlstring) - - while resp.status.state in {StatementState.PENDING, StatementState.RUNNING}: - resp = w.statement_execution.get_statement(resp.statement_id) - time.sleep(response_backoff) - - if resp.status.state != StatementState.SUCCEEDED: - return {"catalog": catalog, - "schema": schema, - "table_name": table_name, - "table_type": table_type, - "location": location, - "status": f"FAIL: {resp.status.error.message}", - "creation_time": time.time_ns()} - - return {"catalog": catalog, - "schema": schema, - "table_name": table_name, - "table_type": table_type, - "location": location, - "status": "SUCCESS", - "creation_time": time.time_ns()} + execute_statement_sync(w, warehouse, sqlstring, backoff=response_backoff) - except Exception as e: - return {"catalog": catalog, - "schema": schema, - "table_name": table_name, - "table_type": table_type, - "location": location, - "status": f"FAIL: {e}", - "creation_time": time.time_ns()} + return { + "catalog": catalog, + "schema": schema, + "table_name": table_name, + "table_type": table_type, + "location": location, + "status": "SUCCESS", + "creation_time": time.time_ns(), + } - else: - print(f"Skipping table {catalog}.{schema}.{table_name}; please check manifest file.") - return {"catalog": catalog, + except Exception as e: + return { + "catalog": catalog, "schema": schema, "table_name": table_name, "table_type": table_type, "location": location, - "status": "FAILURE", - "creation_time": "N/A"} + "status": f"FAIL: {e}", + "creation_time": time.time_ns(), + } + else: + logger.warning( + "Skipping table %s.%s.%s; please check manifest file.", + catalog, + schema, + table_name, + ) + return { + "catalog": catalog, + "schema": schema, + "table_name": table_name, + "table_type": table_type, + "location": location, + "status": "FAILURE", + "creation_time": "N/A", + } -# other parameters -wh_type = CreateWarehouseRequestWarehouseType("PRO") # required for serverless warehouse # initialize lists copied_table_names = [] @@ -223,142 +181,262 @@ def load_table(w, catalog, schema, table_name, table_type, location, warehouse): # create the WorkspaceClient pointed at the source WS w_source = WorkspaceClient(host=source_host, token=source_pat) -# create warehouse in the primary workspace -print("Creating warehouse in primary workspace...") -wh_source = w_source.warehouses.create(name=f'sdk-{time.time_ns()}', - cluster_size=warehouse_size, - max_num_clusters=1, - auto_stop_mins=10, - warehouse_type=wh_type, - enable_serverless_compute=True, - tags=dbsql.EndpointTags( - custom_tags=[ - dbsql.EndpointTagPair(key="Owner", value="dr-sync-tool")])).result() - -system_info = spark.sql("SELECT * FROM system.information_schema.tables") - -# loop through all catalogs to copy, then copy all tables excluding system tables. -# we also skip views; these need to be created separately since they cannot be cloned. -for cat in catalogs_to_copy: - filtered_tables = system_info.filter( - (system_info.table_catalog == cat) & - (system_info.table_schema != "information_schema") & - (system_info.table_type != "VIEW")).collect() - - # get schemas, tables and types in list form - schemas = [row['table_schema'] for row in filtered_tables] - table_names = [row['table_name'] for row in filtered_tables] - table_types = [row['table_type'] for row in filtered_tables] - - # use ThreadPool to copy tables in parallel - with ThreadPoolExecutor(max_workers=num_exec) as executor: - threads = executor.map(copy_table, - repeat(w_source), - repeat(cat), - schemas, - table_names, - table_types, - repeat(landing_zone_url), - repeat(wh_source.id)) - - # wait for threads to execute and build lists for manifest - for thread in threads: - copied_table_names.append(thread["table_name"]) - copied_table_types.append(thread["table_type"]) - copied_table_schemas.append(thread["schema"]) - copied_table_catalogs.append(thread["catalog"]) - copied_table_locations.append( - "{}/{}_{}_{}".format(thread["location"], thread["catalog"], thread["schema"], thread["table_name"])) - print("Copied table {}.{}.{}.".format(thread["catalog"], thread["schema"], thread["table_name"])) - -# create the manifest as a df and write to a table in dr target -# this contains catalog, schema, table and location -manifest_df = pd.DataFrame({"catalog": copied_table_catalogs, - "schema": copied_table_schemas, - "table": copied_table_names, - "location": copied_table_locations, - "type": copied_table_types}) - -# write the manifest to the target bucket in case it needs to be accessed later -ts1 = time.time_ns() -(spark.createDataFrame(manifest_df) - .write.mode("overwrite") - .format("delta") - .save(f"{landing_zone_url}/{manifest_name}-{ts1}")) - -# create the WorkspaceClient pointed at the target WS -w_target = WorkspaceClient(host=target_host, token=target_pat) - -# create warehouse to run table creation statements -print("Creating warehouse in secondary workspace...") -wh_target = w_target.warehouses.create(name=f'sdk-{time.time_ns()}', - cluster_size=warehouse_size, - max_num_clusters=1, - auto_stop_mins=10, - warehouse_type=wh_type, - enable_serverless_compute=True, - tags=dbsql.EndpointTags( - custom_tags=[ - dbsql.EndpointTagPair(key="Owner", value="dr-sync-tool")])).result() - -# initialize lists for status tracking -loaded_table_names = [] -loaded_table_types = [] -loaded_table_schemas = [] -loaded_table_catalogs = [] -loaded_table_locations = [] -loaded_table_status = [] -loaded_table_times = [] - -# drop external tables before loading due to CREATE TABLE restrictions -external_df = manifest_df[manifest_df['type'] == 'EXTERNAL'] -with ThreadPoolExecutor(max_workers=num_exec) as executor: - threads = executor.map(drop_table, - repeat(w_target), - list(external_df['catalog']), - list(external_df['schema']), - list(external_df['table']), - repeat(wh_target.id)) - - for thread in threads: - if thread["status"]: - print("Dropped table {}.{}.{}.".format(thread["catalog"], thread["schema"], thread["table_name"])) - else: - print("Error dropping table {}.{}.{}.".format(thread["catalog"], thread["schema"], thread["table_name"])) - -# load all tables -with ThreadPoolExecutor(max_workers=num_exec) as executor: - threads = executor.map(load_table, - repeat(w_target), - list(manifest_df['catalog']), - list(manifest_df['schema']), - list(manifest_df['table']), - list(manifest_df['type']), - list(manifest_df['location']), - repeat(wh_target.id)) - - for thread in threads: - loaded_table_names.append(thread["table_name"]) - loaded_table_types.append(thread["table_type"]) - loaded_table_schemas.append(thread["schema"]) - loaded_table_catalogs.append(thread["catalog"]) - loaded_table_locations.append(thread["location"]) - loaded_table_status.append(thread["status"]) - loaded_table_times.append(thread["creation_time"]) - print("Loaded table {}.{}.{}.".format(thread["catalog"], thread["schema"], thread["table_name"])) - -# create the table statuses as a df and write to a table in dr target -status_df = pd.DataFrame({"catalog": loaded_table_catalogs, - "schema": loaded_table_schemas, - "table": loaded_table_names, - "location": loaded_table_locations, - "type": loaded_table_types, - "status": loaded_table_status, - "sync_time": loaded_table_times}) - -# table will get a specific timestamp-based location per run -ts2 = time.time_ns() -(spark.createDataFrame(status_df) - .write.mode("overwrite") - .format("delta") - .save(f"{landing_zone_url}/sync_status_{ts2}")) +if config.dry_run: + # In dry-run mode, log what would happen without creating warehouses or executing SQL + for cat in catalogs_to_copy: + filtered_tables = spark.sql( + f""" + SELECT table_schema, table_name, table_type + FROM system.information_schema.tables + WHERE table_catalog = '{cat}' + AND table_schema != 'information_schema' + AND table_type != 'VIEW' + """ + ).collect() + + schemas = [row["table_schema"] for row in filtered_tables] + table_names = [row["table_name"] for row in filtered_tables] + table_types = [row["table_type"] for row in filtered_tables] + + logger.info( + "[DRY RUN] Phase 1: Would copy %d tables from catalog %s to landing zone %s", + len(table_names), + cat, + landing_zone_url, + ) + for schema, table_name, table_type in zip(schemas, table_names, table_types): + logger.info( + "[DRY RUN] Would deep clone %s table %s.%s.%s to %s/%s_%s_%s", + table_type, + cat, + schema, + table_name, + landing_zone_url, + cat, + schema, + table_name, + ) + + logger.info( + "[DRY RUN] Phase 2: Would create warehouse in secondary workspace and load tables from landing zone" + ) + for cat in catalogs_to_copy: + filtered_tables = spark.sql( + f""" + SELECT table_schema, table_name, table_type + FROM system.information_schema.tables + WHERE table_catalog = '{cat}' + AND table_schema != 'information_schema' + AND table_type != 'VIEW' + """ + ).collect() + + schemas = [row["table_schema"] for row in filtered_tables] + table_names = [row["table_name"] for row in filtered_tables] + table_types = [row["table_type"] for row in filtered_tables] + + for schema, table_name, table_type in zip(schemas, table_names, table_types): + if table_type == "EXTERNAL": + logger.info( + "[DRY RUN] Would drop and recreate external table %s.%s.%s", + cat, + schema, + table_name, + ) + else: + logger.info( + "[DRY RUN] Would deep clone %s table %s.%s.%s from landing zone", + table_type, + cat, + schema, + table_name, + ) +else: + # Phase 1: copy tables from source to landing zone + logger.info("Creating warehouse in primary workspace...") + with managed_warehouse(w_source, size=warehouse_size) as wh_source_id: + # loop through all catalogs to copy, then copy all tables excluding system tables. + # we also skip views; these need to be created separately since they cannot be cloned. + for cat in catalogs_to_copy: + filtered_tables = spark.sql( + f""" + SELECT table_schema, table_name, table_type + FROM system.information_schema.tables + WHERE table_catalog = '{cat}' + AND table_schema != 'information_schema' + AND table_type != 'VIEW' + """ + ).collect() + + # get schemas, tables and types in list form + schemas = [row["table_schema"] for row in filtered_tables] + table_names = [row["table_name"] for row in filtered_tables] + table_types = [row["table_type"] for row in filtered_tables] + + # use ThreadPool to copy tables in parallel + with ThreadPoolExecutor(max_workers=num_exec) as executor: + threads = executor.map( + copy_table, + repeat(w_source), + repeat(cat), + schemas, + table_names, + table_types, + repeat(landing_zone_url), + repeat(wh_source_id), + ) + + # wait for threads to execute and build lists for manifest + for thread in threads: + copied_table_names.append(thread["table_name"]) + copied_table_types.append(thread["table_type"]) + copied_table_schemas.append(thread["schema"]) + copied_table_catalogs.append(thread["catalog"]) + copied_table_locations.append( + "{}/{}_{}_{}".format( + thread["location"], + thread["catalog"], + thread["schema"], + thread["table_name"], + ) + ) + logger.info( + "Copied table %s.%s.%s.", + thread["catalog"], + thread["schema"], + thread["table_name"], + ) + + # create the manifest as a df and write to a table in dr target + # this contains catalog, schema, table and location + manifest_df = pd.DataFrame( + { + "catalog": copied_table_catalogs, + "schema": copied_table_schemas, + "table": copied_table_names, + "location": copied_table_locations, + "type": copied_table_types, + } + ) + + # write the manifest to the target bucket in case it needs to be accessed later + ts1 = time.time_ns() + ( + spark.createDataFrame(manifest_df) + .write.mode("overwrite") + .format("delta") + .save(f"{landing_zone_url}/{manifest_name}-{ts1}") + ) + + # Phase 2: load tables from landing zone to target + # create the WorkspaceClient pointed at the target WS + w_target = WorkspaceClient(host=target_host, token=target_pat) + + # initialize lists for status tracking + loaded_table_names = [] + loaded_table_types = [] + loaded_table_schemas = [] + loaded_table_catalogs = [] + loaded_table_locations = [] + loaded_table_status = [] + loaded_table_times = [] + + # create warehouse to run table creation statements, guaranteed cleanup + logger.info("Creating warehouse in secondary workspace...") + with managed_warehouse(w_target, size=warehouse_size) as wh_target_id: + # drop external tables before loading due to CREATE TABLE restrictions + external_df = manifest_df[manifest_df["type"] == "EXTERNAL"] + with ThreadPoolExecutor(max_workers=num_exec) as executor: + threads = executor.map( + drop_table_if_exists, + repeat(w_target), + repeat(wh_target_id), + list(external_df["catalog"]), + list(external_df["schema"]), + list(external_df["table"]), + ) + + for thread in threads: + if thread["status"]: + logger.info( + "Dropped table %s.%s.%s.", + thread["catalog"], + thread["schema"], + thread["table_name"], + ) + else: + logger.error( + "Error dropping table %s.%s.%s.", + thread["catalog"], + thread["schema"], + thread["table_name"], + ) + + # load all tables + with ThreadPoolExecutor(max_workers=num_exec) as executor: + threads = executor.map( + load_table, + repeat(w_target), + list(manifest_df["catalog"]), + list(manifest_df["schema"]), + list(manifest_df["table"]), + list(manifest_df["type"]), + list(manifest_df["location"]), + repeat(wh_target_id), + ) + + for thread in threads: + loaded_table_names.append(thread["table_name"]) + loaded_table_types.append(thread["table_type"]) + loaded_table_schemas.append(thread["schema"]) + loaded_table_catalogs.append(thread["catalog"]) + loaded_table_locations.append(thread["location"]) + loaded_table_status.append(thread["status"]) + loaded_table_times.append(thread["creation_time"]) + logger.info( + "Loaded table %s.%s.%s.", + thread["catalog"], + thread["schema"], + thread["table_name"], + ) + + # create the table statuses as a df and write to a table in dr target + status_df = pd.DataFrame( + { + "catalog": loaded_table_catalogs, + "schema": loaded_table_schemas, + "table": loaded_table_names, + "location": loaded_table_locations, + "type": loaded_table_types, + "status": loaded_table_status, + "sync_time": loaded_table_times, + } + ) + + # table will get a specific timestamp-based location per run + ts2 = time.time_ns() + ( + spark.createDataFrame(status_df) + .write.mode("overwrite") + .format("delta") + .save(f"{landing_zone_url}/sync_status_{ts2}") + ) + +if __name__ == "__main__": + parser = argparse.ArgumentParser( + description="Sync tables from primary to secondary Databricks workspace via deep clone" + ) + parser.add_argument( + "--dry-run", + action="store_true", + help="Show planned operations without executing", + ) + parser.add_argument( + "--log-level", + default="INFO", + choices=["DEBUG", "INFO", "WARNING", "ERROR"], + help="Set logging level", + ) + args = parser.parse_args() + config.dry_run = args.dry_run + logger = setup_logging(level=args.log_level) diff --git a/sync_uc_models.py b/sync_uc_models.py index fd469bb..4b01781 100644 --- a/sync_uc_models.py +++ b/sync_uc_models.py @@ -1,36 +1,74 @@ +import argparse +import os from itertools import repeat from databricks.sdk import WorkspaceClient from concurrent.futures import ThreadPoolExecutor from databricks.sdk.errors.platform import ResourceAlreadyExists -from common import (target_pat, target_host, - source_pat, source_host, - catalogs_to_copy, num_exec) +from dr_sync.config import DRSyncConfig +from dr_sync.log import setup_logging + +config = ( + DRSyncConfig.from_env() + if os.environ.get("DR_SYNC_SOURCE_HOST") + else DRSyncConfig.from_common_module() +) +logger = setup_logging() +target_host = config.target_host +target_pat = config.target_token +source_host = config.source_host +source_pat = config.source_token +catalogs_to_copy = config.catalogs_to_copy +num_exec = config.num_exec # helper function to create models and set appropriate owner def create_model(w, catalog_name, schema_name, model_name, location, owner, comment): - print(f"Creating model {model_name} in {catalog_name}.{schema_name}...") + logger.info("Creating model %s in %s.%s...", model_name, catalog_name, schema_name) + + # dry-run guard: log what would be created without executing + if config.dry_run: + logger.info( + "[DRY RUN] Would create model %s in %s.%s", + model_name, + catalog_name, + schema_name, + ) + return { + "model": f"{catalog_name}.{schema_name}.{model_name}", + "status": "dry_run", + } # try creating new model try: - model = w.registered_models.create(catalog_name=catalog_name, - schema_name=schema_name, - name=model_name, - comment=comment, - storage_location=location) + model = w.registered_models.create( + catalog_name=catalog_name, + schema_name=schema_name, + name=model_name, + comment=comment, + storage_location=location, + ) - _ = w.registered_models.update(full_name=model.full_name, comment=comment, owner=owner) + _ = w.registered_models.update( + full_name=model.full_name, comment=comment, owner=owner + ) return {"model": model.full_name, "status": "success"} # if model already exists, just update the owner (in case it has changed) except ResourceAlreadyExists: - _ = w.registered_models.update(full_name=f"{catalog_name}.{schema_name}.{model_name}", owner=owner) - return {"model": f"{catalog_name}.{schema_name}.{model_name}", "status": "already_exists"} + _ = w.registered_models.update( + full_name=f"{catalog_name}.{schema_name}.{model_name}", owner=owner + ) + return { + "model": f"{catalog_name}.{schema_name}.{model_name}", + "status": "already_exists", + } # for any other exception, return the error except Exception as e: - return {"model": f"{catalog_name}.{schema_name}.{model_name}", "status": f"ERROR: {e}"} - + return { + "model": f"{catalog_name}.{schema_name}.{model_name}", + "status": f"ERROR: {e}", + } # create the WorkspaceClient pointed at the target WS @@ -48,9 +86,11 @@ def create_model(w, catalog_name, schema_name, model_name, location, owner, comm # models and dealing with the "already_exists" errors. We attempt to update owners and comments # in case the model already exists but the owner has changed. for cat in catalogs_to_copy: - filtered_models = [model for model in registered_models - if model.catalog_name == cat - and model.schema_name != "information_schema"] + filtered_models = [ + model + for model in registered_models + if model.catalog_name == cat and model.schema_name != "information_schema" + ] # get schemas, tables and locations in list form schema_names = [model.schema_name for model in filtered_models] @@ -60,19 +100,46 @@ def create_model(w, catalog_name, schema_name, model_name, location, owner, comm model_comments = [model.comment for model in filtered_models] with ThreadPoolExecutor(max_workers=num_exec) as executor: - threads = executor.map(create_model, - repeat(w_target), - repeat(cat), - schema_names, - model_names, - model_locs, - model_owners, - model_comments) + threads = executor.map( + create_model, + repeat(w_target), + repeat(cat), + schema_names, + model_names, + model_locs, + model_owners, + model_comments, + ) for thread in threads: if thread["status"] == "success": - print("Created model {}.".format(thread["model"])) + logger.info("Created model %s.", thread["model"]) elif thread["status"] == "already_exists": - print("Skipped model {} because it already exists.".format(thread["model"])) + logger.warning( + "Skipped model %s because it already exists.", thread["model"] + ) else: - print("Could not create model {}; error: {}".format(thread["model"], thread["status"])) + logger.error( + "Could not create model %s; error: %s", + thread["model"], + thread["status"], + ) + +if __name__ == "__main__": + parser = argparse.ArgumentParser( + description="Sync Unity Catalog registered models between workspaces" + ) + parser.add_argument( + "--dry-run", + action="store_true", + help="Show planned operations without executing", + ) + parser.add_argument( + "--log-level", + default="INFO", + choices=["DEBUG", "INFO", "WARNING", "ERROR"], + help="Set logging level", + ) + args = parser.parse_args() + config.dry_run = args.dry_run + logger = setup_logging(level=args.log_level) diff --git a/sync_views.py b/sync_views.py index fa5f307..0ed7ca4 100644 --- a/sync_views.py +++ b/sync_views.py @@ -1,4 +1,4 @@ -# sync_tables.py +# sync_views.py # # *EXAMPLE* Script to sync views between workspaces. This will very likely need to be altered in your environment to # match the use cases, syntax styles, etc. that you use. Please do NOT expect this to work directly. @@ -24,80 +24,76 @@ # warehouse. Table load statuses will be written to the delta table at {landing_zone_url}/sync_status_{time.time_ns()}. +import argparse +import os import time import pandas as pd from itertools import repeat from databricks.sdk import WorkspaceClient -from databricks.sdk.service import sql as dbsql from concurrent.futures import ThreadPoolExecutor -from databricks.sdk.service.sql import Disposition -from databricks.sdk.service.sql import StatementState -from databricks.sdk.service.sql import CreateWarehouseRequestWarehouseType -from databricks.sdk.service.sql import ExecuteStatementRequestOnWaitTimeout -from common import (target_pat, target_host, - catalogs_to_copy, num_exec, - landing_zone_url, warehouse_size, - response_backoff) +from dr_sync.sql_utils import execute_statement_sync, managed_warehouse +from dr_sync.exceptions import StatementError +from dr_sync.config import DRSyncConfig +from dr_sync.log import setup_logging + +logger = setup_logging() +config = ( + DRSyncConfig.from_env() + if os.environ.get("DR_SYNC_SOURCE_HOST") + else DRSyncConfig.from_common_module() +) +target_host = config.target_host +target_pat = config.target_token +catalogs_to_copy = config.catalogs_to_copy +num_exec = config.num_exec +landing_zone_url = config.landing_zone_url +warehouse_size = config.warehouse_size +response_backoff = config.response_backoff # helper function to create a view def create_view(w, catalog, schema, view_name, warehouse): try: - view_stmt = spark.sql(f"show create table {catalog}.{schema}.{view_name}").collect()[0]["createtab_stmt"] - - resp = w.statement_execution.execute_statement(warehouse_id=warehouse, - wait_timeout="0s", - on_wait_timeout=ExecuteStatementRequestOnWaitTimeout("CONTINUE"), - disposition=Disposition("EXTERNAL_LINKS"), - statement=view_stmt) - - while resp.status.state in {StatementState.PENDING, StatementState.RUNNING}: - resp = w.statement_execution.get_statement(resp.statement_id) - time.sleep(response_backoff) - - if resp.status.state != StatementState.SUCCEEDED: - return {"catalog": catalog, - "schema": schema, - "view_name": view_name, - "status": f"FAIL: {resp.status.error.message}", - "creation_time": time.time_ns()} - - return {"catalog": catalog, - "schema": schema, - "view_name": view_name, - "status": "SUCCESS", - "creation_time": time.time_ns()} + view_stmt = spark.sql( + f"show create table {catalog}.{schema}.{view_name}" + ).collect()[0]["createtab_stmt"] + + execute_statement_sync(w, warehouse, view_stmt, backoff=response_backoff) + + return { + "catalog": catalog, + "schema": schema, + "view_name": view_name, + "status": "SUCCESS", + "creation_time": time.time_ns(), + } + + except StatementError as e: + return { + "catalog": catalog, + "schema": schema, + "view_name": view_name, + "status": f"FAIL: {e}", + "creation_time": time.time_ns(), + } except Exception as e: - return {"catalog": catalog, - "schema": schema, - "view_name": view_name, - f"status": f"FAIL: {e}", - "creation_time": time.time_ns()} + return { + "catalog": catalog, + "schema": schema, + "view_name": view_name, + "status": f"FAIL: {e}", + "creation_time": time.time_ns(), + } -# other parameters -wh_type = CreateWarehouseRequestWarehouseType("PRO") # required for serverless warehouse - # pull all views from source ws all_views = spark.sql("SELECT * FROM system.information_schema.views") # create the WorkspaceClient pointed at the target WS w_target = WorkspaceClient(host=target_host, token=target_pat) -# create warehouse to run view creation statements -print("Creating warehouse in secondary workspace...") -wh_target = w_target.warehouses.create(name=f'sdk-{time.time_ns()}', - cluster_size=warehouse_size, - max_num_clusters=1, - auto_stop_mins=10, - warehouse_type=wh_type, - enable_serverless_compute=True, - tags=dbsql.EndpointTags( - custom_tags=[ - dbsql.EndpointTagPair(key="Owner", value="dr-sync-tool")])).result() - # initialize lists for status tracking loaded_view_names = [] loaded_view_schemas = [] @@ -105,42 +101,95 @@ def create_view(w, catalog, schema, view_name, warehouse): loaded_view_status = [] loaded_view_times = [] -# load all views per catalog -for cat in catalogs_to_copy: - filtered_views = all_views.filter( - (all_views.table_catalog == cat) & - all_views.table_schema != "information_schema").collect() - - # get schemas and view names - schemas = [row['table_schema'] for row in filtered_views] - view_names = [row['table_name'] for row in filtered_views] - - with ThreadPoolExecutor(max_workers=num_exec) as executor: - threads = executor.map(create_view, - repeat(w_target), - repeat(cat), - schemas, - view_names, - repeat(wh_target.id)) - - for thread in threads: - loaded_view_names.append(thread["view_name"]) - loaded_view_schemas.append(thread["schema"]) - loaded_view_catalogs.append(thread["catalog"]) - loaded_view_status.append(thread["status"]) - loaded_view_times.append(thread["creation_time"]) - print("Loaded view {}.{}.{}.".format(thread["catalog"], thread["schema"], thread["view_name"])) - -# create the table statuses as a df and write to a table in dr target -status_df = pd.DataFrame({"catalog": loaded_view_catalogs, - "schema": loaded_view_schemas, - "table": loaded_view_names, - "status": loaded_view_status, - "sync_time": loaded_view_times}) - -# table will get a specific timestamp-based location per run -ts = time.time_ns() -(spark.createDataFrame(status_df) - .write.mode("overwrite") - .format("delta") - .save(f"{landing_zone_url}/view_sync_status_{ts}")) +if config.dry_run: + # In dry-run mode, log what would happen without creating warehouses or executing SQL + for cat in catalogs_to_copy: + filtered_views = all_views.filter( + (all_views.table_catalog == cat) + & (all_views.table_schema != "information_schema") + ).collect() + + schemas = [row["table_schema"] for row in filtered_views] + view_names = [row["table_name"] for row in filtered_views] + + logger.info( + "[DRY RUN] Would create %d views in catalog %s", len(view_names), cat + ) + for schema, view_name in zip(schemas, view_names): + logger.info("[DRY RUN] Would create view %s.%s.%s", cat, schema, view_name) +else: + # create warehouse to run view creation statements, guaranteed cleanup + logger.info("Creating warehouse in secondary workspace...") + with managed_warehouse(w_target, size=warehouse_size) as wh_id: + # load all views per catalog + for cat in catalogs_to_copy: + filtered_views = all_views.filter( + (all_views.table_catalog == cat) + & (all_views.table_schema != "information_schema") + ).collect() + + # get schemas and view names + schemas = [row["table_schema"] for row in filtered_views] + view_names = [row["table_name"] for row in filtered_views] + + with ThreadPoolExecutor(max_workers=num_exec) as executor: + threads = executor.map( + create_view, + repeat(w_target), + repeat(cat), + schemas, + view_names, + repeat(wh_id), + ) + + for thread in threads: + loaded_view_names.append(thread["view_name"]) + loaded_view_schemas.append(thread["schema"]) + loaded_view_catalogs.append(thread["catalog"]) + loaded_view_status.append(thread["status"]) + loaded_view_times.append(thread["creation_time"]) + logger.info( + "Loaded view %s.%s.%s.", + thread["catalog"], + thread["schema"], + thread["view_name"], + ) + + # create the table statuses as a df and write to a table in dr target + status_df = pd.DataFrame( + { + "catalog": loaded_view_catalogs, + "schema": loaded_view_schemas, + "table": loaded_view_names, + "status": loaded_view_status, + "sync_time": loaded_view_times, + } + ) + + # table will get a specific timestamp-based location per run + ts = time.time_ns() + ( + spark.createDataFrame(status_df) + .write.mode("overwrite") + .format("delta") + .save(f"{landing_zone_url}/view_sync_status_{ts}") + ) + +if __name__ == "__main__": + parser = argparse.ArgumentParser( + description="Sync views between primary and secondary Databricks workspaces" + ) + parser.add_argument( + "--dry-run", + action="store_true", + help="Show planned operations without executing", + ) + parser.add_argument( + "--log-level", + default="INFO", + choices=["DEBUG", "INFO", "WARNING", "ERROR"], + help="Set logging level", + ) + args = parser.parse_args() + config.dry_run = args.dry_run + logger = setup_logging(level=args.log_level)