duckrun 0.2.1__py3-none-any.whl → 0.2.3__py3-none-any.whl

duckrun/files.py

@@ -0,0 +1,251 @@
+"""
+File operations functionality for duckrun - OneLake Files copy and download
+"""
+import os
+from typing import Optional, List
+import obstore as obs
+from obstore.store import AzureStore
+
+
+def copy(duckrun_instance, local_folder: str, remote_folder: str, 
+         file_extensions: Optional[List[str]] = None, 
+         overwrite: bool = False) -> bool:
+    """
+    Copy files from a local folder to OneLake Files section.
+    
+    Args:
+        duckrun_instance: The Duckrun connection instance
+        local_folder: Path to local folder containing files to upload
+        remote_folder: Target subfolder path in OneLake Files (e.g., "reports/daily") - REQUIRED
+        file_extensions: Optional list of file extensions to filter (e.g., ['.csv', '.parquet'])
+        overwrite: Whether to overwrite existing files (default: False)
+        
+    Returns:
+        True if all files uploaded successfully, False otherwise
+        
+    Examples:
+        # Upload all files from local folder to a target folder
+        dr.copy("./local_data", "uploaded_data")
+        
+        # Upload only CSV files to a specific subfolder
+        dr.copy("./reports", "daily_reports", ['.csv'])
+        
+        # Upload with overwrite enabled
+        dr.copy("./backup", "backups", overwrite=True)
+    """
+    if not os.path.exists(local_folder):
+        print(f"❌ Local folder not found: {local_folder}")
+        return False
+        
+    if not os.path.isdir(local_folder):
+        print(f"❌ Path is not a directory: {local_folder}")
+        return False
+        
+    # Get Azure token
+    token = duckrun_instance._get_storage_token()
+    if token == "PLACEHOLDER_TOKEN_TOKEN_NOT_AVAILABLE":
+        print("Authenticating with Azure for file upload (trying CLI, will fallback to browser if needed)...")
+        from azure.identity import AzureCliCredential, InteractiveBrowserCredential, ChainedTokenCredential
+        credential = ChainedTokenCredential(AzureCliCredential(), InteractiveBrowserCredential())
+        token_obj = credential.get_token("https://storage.azure.com/.default")
+        token = token_obj.token
+        os.environ["AZURE_STORAGE_TOKEN"] = token
+    
+    # Setup OneLake Files URL (not Tables)
+    files_base_url = f'abfss://{duckrun_instance.workspace}@{duckrun_instance.storage_account}.dfs.fabric.microsoft.com/{duckrun_instance.lakehouse_name}.Lakehouse/Files/'
+    store = AzureStore.from_url(files_base_url, bearer_token=token)
+    
+    # Collect files to upload
+    files_to_upload = []
+    for root, dirs, files in os.walk(local_folder):
+        for file in files:
+            local_file_path = os.path.join(root, file)
+            
+            # Filter by extensions if specified
+            if file_extensions:
+                _, ext = os.path.splitext(file)
+                if ext.lower() not in [e.lower() for e in file_extensions]:
+                    continue
+            
+            # Calculate relative path from local_folder
+            rel_path = os.path.relpath(local_file_path, local_folder)
+            
+            # Build remote path in OneLake Files (remote_folder is now mandatory)
+            remote_path = f"{remote_folder.strip('/')}/{rel_path}".replace("\\", "/")
+            
+            files_to_upload.append((local_file_path, remote_path))
+    
+    if not files_to_upload:
+        print(f"No files found to upload in {local_folder}")
+        if file_extensions:
+            print(f"  (filtered by extensions: {file_extensions})")
+        return True
+    
+    print(f"📁 Uploading {len(files_to_upload)} files from '{local_folder}' to OneLake Files...")
+    print(f"   Target folder: {remote_folder}")
+    
+    uploaded_count = 0
+    failed_count = 0
+    
+    for local_path, remote_path in files_to_upload:
+        try:
+            # Check if file exists (if not overwriting)
+            if not overwrite:
+                try:
+                    obs.head(store, remote_path)
+                    print(f"  ⏭ Skipped (exists): {remote_path}")
+                    continue
+                except Exception:
+                    # File doesn't exist, proceed with upload
+                    pass
+            
+            # Read local file
+            with open(local_path, 'rb') as f:
+                file_data = f.read()
+            
+            # Upload to OneLake Files
+            obs.put(store, remote_path, file_data)
+            
+            file_size = len(file_data)
+            size_mb = file_size / (1024 * 1024) if file_size > 1024*1024 else file_size / 1024
+            size_unit = "MB" if file_size > 1024*1024 else "KB"
+            
+            print(f"  ✓ Uploaded: {local_path} → {remote_path} ({size_mb:.1f} {size_unit})")
+            uploaded_count += 1
+            
+        except Exception as e:
+            print(f"  ❌ Failed: {local_path} → {remote_path} | Error: {str(e)[:100]}")
+            failed_count += 1
+    
+    print(f"\n{'='*60}")
+    if failed_count == 0:
+        print(f"✅ Successfully uploaded all {uploaded_count} files to OneLake Files")
+    else:
+        print(f"⚠ Uploaded {uploaded_count} files, {failed_count} failed")
+    print(f"{'='*60}")
+    
+    return failed_count == 0
+
+
+def download(duckrun_instance, remote_folder: str = "", local_folder: str = "./downloaded_files",
+             file_extensions: Optional[List[str]] = None,
+             overwrite: bool = False) -> bool:
+    """
+    Download files from OneLake Files section to a local folder.
+    
+    Args:
+        duckrun_instance: The Duckrun connection instance
+        remote_folder: Optional subfolder path in OneLake Files to download from
+        local_folder: Local folder path to download files to (default: "./downloaded_files")
+        file_extensions: Optional list of file extensions to filter (e.g., ['.csv', '.parquet'])
+        overwrite: Whether to overwrite existing local files (default: False)
+        
+    Returns:
+        True if all files downloaded successfully, False otherwise
+        
+    Examples:
+        # Download all files from OneLake Files root
+        dr.download()
+        
+        # Download only CSV files from a specific subfolder
+        dr.download("daily_reports", "./reports", ['.csv'])
+    """
+    # Get Azure token
+    token = duckrun_instance._get_storage_token()
+    if token == "PLACEHOLDER_TOKEN_TOKEN_NOT_AVAILABLE":
+        print("Authenticating with Azure for file download (trying CLI, will fallback to browser if needed)...")
+        from azure.identity import AzureCliCredential, InteractiveBrowserCredential, ChainedTokenCredential
+        credential = ChainedTokenCredential(AzureCliCredential(), InteractiveBrowserCredential())
+        token_obj = credential.get_token("https://storage.azure.com/.default")
+        token = token_obj.token
+        os.environ["AZURE_STORAGE_TOKEN"] = token
+    
+    # Setup OneLake Files URL (not Tables)
+    files_base_url = f'abfss://{duckrun_instance.workspace}@{duckrun_instance.storage_account}.dfs.fabric.microsoft.com/{duckrun_instance.lakehouse_name}.Lakehouse/Files/'
+    store = AzureStore.from_url(files_base_url, bearer_token=token)
+    
+    # Create local directory
+    os.makedirs(local_folder, exist_ok=True)
+    
+    # List files in OneLake Files
+    print(f"📁 Discovering files in OneLake Files...")
+    if remote_folder:
+        print(f"   Source folder: {remote_folder}")
+        prefix = f"{remote_folder.strip('/')}/"
+    else:
+        prefix = ""
+    
+    try:
+        list_stream = obs.list(store, prefix=prefix)
+        files_to_download = []
+        
+        for batch in list_stream:
+            for obj in batch:
+                remote_path = obj["path"]
+                
+                # Filter by extensions if specified
+                if file_extensions:
+                    _, ext = os.path.splitext(remote_path)
+                    if ext.lower() not in [e.lower() for e in file_extensions]:
+                        continue
+                
+                # Calculate local path
+                if remote_folder:
+                    rel_path = os.path.relpath(remote_path, remote_folder.strip('/'))
+                else:
+                    rel_path = remote_path
+                
+                local_path = os.path.join(local_folder, rel_path).replace('/', os.sep)
+                files_to_download.append((remote_path, local_path))
+        
+        if not files_to_download:
+            print(f"No files found to download")
+            if file_extensions:
+                print(f"  (filtered by extensions: {file_extensions})")
+            return True
+        
+        print(f"📥 Downloading {len(files_to_download)} files to '{local_folder}'...")
+        
+        downloaded_count = 0
+        failed_count = 0
+        
+        for remote_path, local_path in files_to_download:
+            try:
+                # Check if local file exists (if not overwriting)
+                if not overwrite and os.path.exists(local_path):
+                    print(f"  ⏭ Skipped (exists): {local_path}")
+                    continue
+                
+                # Ensure local directory exists
+                os.makedirs(os.path.dirname(local_path), exist_ok=True)
+                
+                # Download file
+                data = obs.get(store, remote_path).bytes()
+                
+                # Write to local file
+                with open(local_path, 'wb') as f:
+                    f.write(data)
+                
+                file_size = len(data)
+                size_mb = file_size / (1024 * 1024) if file_size > 1024*1024 else file_size / 1024
+                size_unit = "MB" if file_size > 1024*1024 else "KB"
+                
+                print(f"  ✓ Downloaded: {remote_path} → {local_path} ({size_mb:.1f} {size_unit})")
+                downloaded_count += 1
+                
+            except Exception as e:
+                print(f"  ❌ Failed: {remote_path} → {local_path} | Error: {str(e)[:100]}")
+                failed_count += 1
+        
+        print(f"\n{'='*60}")
+        if failed_count == 0:
+            print(f"✅ Successfully downloaded all {downloaded_count} files from OneLake Files")
+        else:
+            print(f"⚠ Downloaded {downloaded_count} files, {failed_count} failed")
+        print(f"{'='*60}")
+        
+        return failed_count == 0
+        
+    except Exception as e:
+        print(f"❌ Error listing files from OneLake: {e}")
+        return False

duckrun/runner.py

@@ -0,0 +1,287 @@
+"""
+Pipeline execution functionality for duckrun
+"""
+import os
+import requests
+import importlib.util
+from typing import List, Tuple, Dict, Optional, Callable, Any
+from string import Template
+from deltalake import DeltaTable, write_deltalake
+# Row Group configuration for optimal Delta Lake performance
+RG = 8_000_000
+
+
+def _build_write_deltalake_args(path, df, mode, schema_mode=None, partition_by=None):
+    """
+    Build arguments for write_deltalake based on requirements:
+    - If schema_mode='merge': use rust engine (no row group params)
+    - Otherwise: use pyarrow engine with row group optimization
+    """
+    args = {
+        'table_or_uri': path,
+        'data': df,
+        'mode': mode
+    }
+    
+    # Add partition_by if specified
+    if partition_by:
+        args['partition_by'] = partition_by
+    
+    # Engine selection based on schema_mode
+    if schema_mode == 'merge':
+        # Use rust engine for schema merging (no row group params supported)
+        args['schema_mode'] = 'merge'
+        args['engine'] = 'rust'
+    else:
+        # Use pyarrow engine with row group optimization (default)
+        args['max_rows_per_file'] = RG
+        args['max_rows_per_group'] = RG
+        args['min_rows_per_group'] = RG
+    
+    return args
+
+
+def run(duckrun_instance, pipeline: List[Tuple]) -> bool:
+    """
+    Execute pipeline of tasks.
+    
+    Task formats:
+        - Python: ('function_name', (arg1, arg2, ...))
+        - SQL:    ('table_name', 'mode') or ('table_name', 'mode', {sql_params})
+        - SQL with Delta options: ('table_name', 'mode', {sql_params}, {delta_options})
+    
+    Returns:
+        True if all tasks succeeded
+        False if any task failed (exception) or Python task returned 0 (early exit)
+    """
+    if duckrun_instance.sql_folder is None:
+        raise RuntimeError("sql_folder is not configured. Cannot run pipelines.")
+    
+    for i, task in enumerate(pipeline, 1):
+        print(f"\n{'='*60}")
+        print(f"Task {i}/{len(pipeline)}: {task[0]}")
+        print('='*60)
+        
+        try:
+            result = None
+            
+            if len(task) == 2:
+                name, second = task
+                if isinstance(second, str) and second in {'overwrite', 'append', 'ignore'}:
+                    result = _run_sql(duckrun_instance, name, second, {}, {})
+                else:
+                    args = second if isinstance(second, (tuple, list)) else (second,)
+                    result = _run_python(duckrun_instance, name, tuple(args))
+                
+            elif len(task) == 3:
+                table, mode, params = task
+                if not isinstance(params, dict):
+                    raise ValueError(f"Expected dict for params, got {type(params)}")
+                result = _run_sql(duckrun_instance, table, mode, params, {})
+                
+            elif len(task) == 4:
+                table, mode, params, delta_options = task
+                if not isinstance(params, dict):
+                    raise ValueError(f"Expected dict for SQL params, got {type(params)}")
+                if not isinstance(delta_options, dict):
+                    raise ValueError(f"Expected dict for Delta options, got {type(delta_options)}")
+                result = _run_sql(duckrun_instance, table, mode, params, delta_options)
+                
+            else:
+                raise ValueError(f"Invalid task format: {task}")
+            
+            # Check if Python task returned 0 (early exit condition)
+            # Only check for Python tasks as SQL tasks return table names (strings) and only stop on exceptions
+            if (len(task) == 2 and 
+                not isinstance(task[1], str) and 
+                result == 0):
+                print(f"\n⏹️  Python task {i} returned 0 - stopping pipeline execution")
+                print(f"   Remaining tasks ({len(pipeline) - i}) will not be executed")
+                return False
+                
+        except Exception as e:
+            print(f"\n❌ Task {i} failed: {e}")
+            return False
+
+    print(f"\n{'='*60}")
+    print("✅ All tasks completed successfully")
+    print('='*60)
+    return True
+
+
+def _run_python(duckrun_instance, name: str, args: tuple) -> Any:
+    """Execute Python task, return result"""
+    duckrun_instance._create_onelake_secret()
+    func = _load_py_function(duckrun_instance, name)
+    if not func:
+        raise RuntimeError(f"Python function '{name}' not found")
+    
+    print(f"Running Python: {name}{args}")
+    result = func(*args)
+    print(f"✅ Python '{name}' completed")
+    return result
+
+
+def _run_sql(duckrun_instance, table: str, mode: str, params: Dict, delta_options: Dict = None) -> str:
+    """Execute SQL task, write to Delta, return normalized table name"""
+    duckrun_instance._create_onelake_secret()
+    
+    if mode not in {'overwrite', 'append', 'ignore'}:
+        raise ValueError(f"Invalid mode '{mode}'. Use: overwrite, append, or ignore")
+
+    sql = _read_sql_file(duckrun_instance, table, params)
+    if sql is None:
+        raise RuntimeError(f"Failed to read SQL file for '{table}'")
+
+    normalized_table = _normalize_table_name(table)
+    path = f"{duckrun_instance.table_base_url}{duckrun_instance.schema}/{normalized_table}"
+
+    # Extract Delta Lake specific options from delta_options
+    delta_options = delta_options or {}
+    merge_schema = delta_options.get('mergeSchema')
+    schema_mode = 'merge' if str(merge_schema).lower() in ('true', '1') else None
+    partition_by = delta_options.get('partitionBy') or delta_options.get('partition_by')
+
+    if mode == 'overwrite':
+        duckrun_instance.con.sql(f"DROP VIEW IF EXISTS {normalized_table}")
+        df = duckrun_instance.con.sql(sql).record_batch()
+        
+        write_args = _build_write_deltalake_args(
+            path, df, 'overwrite', 
+            schema_mode=schema_mode, 
+            partition_by=partition_by
+        )
+        write_deltalake(**write_args)
+        
+        duckrun_instance.con.sql(f"CREATE OR REPLACE VIEW {normalized_table} AS SELECT * FROM delta_scan('{path}')")
+        dt = DeltaTable(path)
+        dt.vacuum(retention_hours=0, dry_run=False, enforce_retention_duration=False)
+        dt.cleanup_metadata()
+
+    elif mode == 'append':
+        df = duckrun_instance.con.sql(sql).record_batch()
+        
+        write_args = _build_write_deltalake_args(
+            path, df, 'append', 
+            schema_mode=schema_mode, 
+            partition_by=partition_by
+        )
+        write_deltalake(**write_args)
+        
+        duckrun_instance.con.sql(f"CREATE OR REPLACE VIEW {normalized_table} AS SELECT * FROM delta_scan('{path}')")
+        dt = DeltaTable(path)
+        if len(dt.file_uris()) > duckrun_instance.compaction_threshold:
+            print(f"Compacting {normalized_table} ({len(dt.file_uris())} files)")
+            dt.optimize.compact()
+            dt.vacuum(dry_run=False)
+            dt.cleanup_metadata()
+
+    elif mode == 'ignore':
+        try:
+            DeltaTable(path)
+            print(f"Table {normalized_table} exists. Skipping (mode='ignore')")
+        except Exception:
+            print(f"Table {normalized_table} doesn't exist. Creating...")
+            duckrun_instance.con.sql(f"DROP VIEW IF EXISTS {normalized_table}")
+            df = duckrun_instance.con.sql(sql).record_batch()
+            
+            write_args = _build_write_deltalake_args(
+                path, df, 'overwrite', 
+                schema_mode=schema_mode, 
+                partition_by=partition_by
+            )
+            write_deltalake(**write_args)
+            
+            duckrun_instance.con.sql(f"CREATE OR REPLACE VIEW {normalized_table} AS SELECT * FROM delta_scan('{path}')")
+            dt = DeltaTable(path)
+            dt.vacuum(dry_run=False)
+            dt.cleanup_metadata()
+
+    engine_info = f" (engine=rust, schema_mode=merge)" if schema_mode == 'merge' else " (engine=pyarrow)"
+    partition_info = f" partitioned by {partition_by}" if partition_by else ""
+    print(f"✅ SQL '{table}' → '{normalized_table}' ({mode}){engine_info}{partition_info}")
+    return normalized_table
+
+
+def _normalize_table_name(name: str) -> str:
+    """Extract base table name before first '__'"""
+    return name.split('__', 1)[0] if '__' in name else name
+
+
+def _read_sql_file(duckrun_instance, table_name: str, params: Optional[Dict] = None) -> Optional[str]:
+    if duckrun_instance.sql_folder is None:
+        raise RuntimeError("sql_folder is not configured. Cannot read SQL files.")
+    
+    is_url = duckrun_instance.sql_folder.startswith("http")
+    if is_url:
+        url = f"{duckrun_instance.sql_folder.rstrip('/')}/{table_name}.sql".strip()
+        try:
+            resp = requests.get(url)
+            resp.raise_for_status()
+            content = resp.text
+        except Exception as e:
+            print(f"Failed to fetch SQL from {url}: {e}")
+            return None
+    else:
+        path = os.path.join(duckrun_instance.sql_folder, f"{table_name}.sql")
+        try:
+            with open(path, 'r') as f:
+                content = f.read()
+        except Exception as e:
+            print(f"Failed to read SQL file {path}: {e}")
+            return None
+
+    if not content.strip():
+        print(f"SQL file is empty: {table_name}.sql")
+        return None
+
+    full_params = {
+        'ws': duckrun_instance.workspace,
+        'lh': duckrun_instance.lakehouse_name,
+        'schema': duckrun_instance.schema,
+        'storage_account': duckrun_instance.storage_account
+    }
+    if params:
+        full_params.update(params)
+
+    try:
+        template = Template(content)
+        content = template.substitute(full_params)
+    except KeyError as e:
+        print(f"Missing parameter in SQL file: ${e}")
+        return None
+    except Exception as e:
+        print(f"Error during SQL template substitution: {e}")
+        return None
+
+    return content
+
+
+def _load_py_function(duckrun_instance, name: str) -> Optional[Callable]:
+    if duckrun_instance.sql_folder is None:
+        raise RuntimeError("sql_folder is not configured. Cannot load Python functions.")
+    
+    is_url = duckrun_instance.sql_folder.startswith("http")
+    try:
+        if is_url:
+            url = f"{duckrun_instance.sql_folder.rstrip('/')}/{name}.py".strip()
+            resp = requests.get(url)
+            resp.raise_for_status()
+            code = resp.text
+            namespace = {}
+            exec(code, namespace)
+            func = namespace.get(name)
+            return func if callable(func) else None
+        else:
+            path = os.path.join(duckrun_instance.sql_folder, f"{name}.py")
+            if not os.path.isfile(path):
+                print(f"Python file not found: {path}")
+                return None
+            spec = importlib.util.spec_from_file_location(name, path)
+            mod = importlib.util.module_from_spec(spec)
+            spec.loader.exec_module(mod)
+            func = getattr(mod, name, None)
+            return func if callable(func) else None
+    except Exception as e:
+        print(f"Error loading Python function '{name}': {e}")
+        return None