Refactor: bug fixes, resource safety, shared utilities, logging, validation, and performance

.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://<primary-workspace-hostname>
+DR_SYNC_SOURCE_TOKEN=<primary-workspace-pat>
+
+# Target (secondary) workspace
+DR_SYNC_TARGET_HOST=https://<secondary-workspace-hostname>
+DR_SYNC_TARGET_TOKEN=<secondary-workspace-pat>
+
+# 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=<secondary-metastore-id>
+DR_SYNC_MANIFEST_NAME=manifest
+
+# Runtime flags
+DR_SYNC_DRY_RUN=false

.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

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 = "<primary-ws-hostname>"               # source hostname, including https://
-source_pat = "<primary-ws-pat>"                     # source PAT
-target_host = "<secondary-ws-hostname>"             # target hostname, including https://
-target_pat = "<primary-ws-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 = "<secondary-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 = "<primary-ws-hostname>"  # source hostname, including https://
+source_pat = "<primary-ws-pat>"  # source PAT
+target_host = "<secondary-ws-hostname>"  # target hostname, including https://
+target_pat = "<primary-ws-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 = "<secondary-metastore-id>"  # global metastore ID for secondary metastore
+manifest_name = "manifest"  # name of the manifest file, if written

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",
+]

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