Enhance DumpCleaner and import_to_supabase functions with target schema support, improve logging, and refine error handling

Author: Domains18

cloudsql_to_supabase/clean.py

@@ -7,18 +7,43 @@
 logger = logging.getLogger('cloudsql_to_supabase.clean')
 
 class DumpCleaner:
-    def __init__(self, input_file: Path = None, output_file: Path = None) -> None:
+    def __init__(self, input_file: Path = None, output_file: Path = None, target_schema: str = None) -> None:
         self.input_file = input_file or Path(config.OUTPUT_DUMP)
         self.output_file = output_file or Path(config.CLEANED_DUMP)
+        self.target_schema = target_schema or config.SUPABASE_SCHEMA
+        
         self.skip_patterns = [
             r'^(CREATE|ALTER) ROLE',  # Skip role creation/alteration
             r'^COMMENT ON EXTENSION',  # Skip extension comments
         ]
+        
+        # Regular ownership replacement
+        owner_replacement = 'OWNER TO public;'
+        # If using non-public schema and we have permissions, keep the target schema as owner
+        if self.target_schema != "public":
+            owner_replacement = f'OWNER TO {self.target_schema};'
+            
         self.replacement_rules = [
-            (r'OWNER TO .*?;', 'OWNER TO public;'),  # Change ownership to public
+            (r'OWNER TO .*?;', owner_replacement),
             (r'CREATE SCHEMA .*?;', '-- Schema creation removed'),  # Comment out schema creation
         ]
 
+        # If we're targeting a non-public schema, add schema modifications
+        if self.target_schema != "public":
+            # Add rules to update schema references in the dump
+            self.replacement_rules.extend([
+                # Update SET search_path statements
+                (r'SET search_path = public', f'SET search_path = {self.target_schema}'),
+                # Update schema references in table names
+                (r'CREATE TABLE public\.', f'CREATE TABLE {self.target_schema}.'),
+                # Update schema references in sequence names
+                (r'CREATE SEQUENCE public\.', f'CREATE SEQUENCE {self.target_schema}.'),
+                # Update schema references in views
+                (r'CREATE VIEW public\.', f'CREATE VIEW {self.target_schema}.'),
+                # Update schema references in functions
+                (r'CREATE FUNCTION public\.', f'CREATE FUNCTION {self.target_schema}.'),
+            ])
+        
     def clean_dump_file(self) -> Path:
         """
         Clean the SQL dump file for Supabase import by removing/modifying
@@ -57,9 +82,14 @@ def clean_dump_file(self) -> Path:
         logger.info(f"Cleaned dump saved as {self.output_file}")
         return self.output_file
 
-def clean_dump_file(input_file: Optional[Path] = None, output_file: Optional[Path] = None) -> Path:
+def clean_dump_file(input_file: Optional[Path] = None, output_file: Optional[Path] = None, target_schema: Optional[str] = None) -> Path:
     """
     Convenience function to clean a SQL dump file
+    
+    Args:
+        input_file: Path to the input SQL dump file
+        output_file: Path to the output cleaned SQL file
+        target_schema: Schema to use in Supabase. If None, uses the one from config
     """
-    cleaner = DumpCleaner(input_file, output_file)
-    return cleaner.clean_dump_file()
+    cleaner = DumpCleaner(input_file, output_file, target_schema)
+    return cleaner.clean_dump_file()

cloudsql_to_supabase/config.py

@@ -1,50 +1,54 @@
-from dotenv import load_dotenv
 import os
 from pathlib import Path
-import logging 
-
-logging.basicConfig(level=logging.INFO,format='%(asctime)s - %(levelname)s - %(message)s')
-
+from dotenv import load_dotenv
+import logging
 
+# Configure logging
+logging.basicConfig(
+    level=logging.INFO,
+    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
+)
 logger = logging.getLogger('cloudsql_to_supabase')
 
+# Load environment variables
 load_dotenv()
 
-
+# Database configurations
 CLOUDSQL_USER = os.getenv("CLOUDSQL_USER")
 CLOUDSQL_HOST = os.getenv("CLOUDSQL_HOST")
 CLOUDSQL_DB = os.getenv("CLOUDSQL_DB")
 CLOUDSQL_PORT = int(os.getenv("CLOUDSQL_PORT", 5432))
 CLOUDSQL_SSL_MODE = os.getenv("CLOUDSQL_SSL_MODE", "prefer")
+CLOUDSQL_SCHEMA = os.getenv("CLOUDSQL_SCHEMA", "public")
 
 SUPABASE_USER = os.getenv("SUPABASE_USER", "postgres")
 SUPABASE_HOST = os.getenv("SUPABASE_HOST")
 SUPABASE_DB = os.getenv("SUPABASE_DB", "postgres")
 SUPABASE_PASSWORD = os.getenv("SUPABASE_PASSWORD")
 SUPABASE_PORT = int(os.getenv("SUPABASE_PORT", 5432))
 SUPABASE_SSL_MODE = os.getenv("SUPABASE_SSL_MODE", "require")
+SUPABASE_SCHEMA = os.getenv("SUPABASE_SCHEMA", "public")
 
-OUTPUT_DIR = Path(os.getenv('OUTPUT_DIR', '.'))
+# Output file paths
+OUTPUT_DIR = Path(os.getenv("OUTPUT_DIR", "."))
 OUTPUT_DUMP = OUTPUT_DIR / os.getenv("OUTPUT_DUMP", "backup.sql")
 CLEANED_DUMP = OUTPUT_DIR / os.getenv("CLEANED_DUMP", "cleaned_backup.sql")
 
-
+# Validation function
 def validate_config():
-    """
-    Validate the configuration values.
-    """
+    """Validate that the required environment variables are set."""
     required_vars = {
-         "CLOUDSQL_USER": CLOUDSQL_USER,
+        "CLOUDSQL_USER": CLOUDSQL_USER,
         "CLOUDSQL_HOST": CLOUDSQL_HOST,
         "CLOUDSQL_DB": CLOUDSQL_DB,
         "SUPABASE_HOST": SUPABASE_HOST,
         "SUPABASE_PASSWORD": SUPABASE_PASSWORD,
     }
 
-    
     missing = [var for var, value in required_vars.items() if not value]
 
     if missing:
         raise ValueError(f"Missing required environment variables: {', '.join(missing)}")
 
-    OUTPUT_DIR.mkdir(exist_ok=True)
+    # Create output directory if it doesn't exist
+    OUTPUT_DIR.mkdir(exist_ok=True)

cloudsql_to_supabase/import_.py

@@ -4,46 +4,65 @@
 from typing import Optional
 from . import config, utils
 
-
-
 logger = logging.getLogger('cloudsql_to_supabase.import')
 
-
-
-def import_to_supabase(input_file: Optional[Path] = None, password: Optional[str] = None) -> None:
+def import_to_supabase(input_file: Optional[Path] = None, password: Optional[str] = None, schema: Optional[str] = None) -> None:
     """
-        import cleaned sql dump into supabase
-        
-        Args:
-            input_file: Path to the sql dump file to import. if None, uses the default.
-            password: supabase database password. if none uses from config
-            
+    Import a cleaned SQL dump file into Supabase
+    
+    Args:
+        input_file: Path to the SQL dump file to import. If None, uses the default.
+        password: Supabase database password. If None, uses from config.
+        schema: Target schema to import into. If None, uses the one from config.
     """
     config.validate_config()
     dump_file = input_file or Path(config.CLEANED_DUMP)
+    target_schema = schema or config.SUPABASE_SCHEMA
+    
     if not dump_file.exists():
-        raise FileNotFoundError(f"Dump not found: {dump_file}")
+        raise FileNotFoundError(f"Dump file not found: {dump_file}")
 
-    logger.info(f'importing data into supabase database: {config.SUPABASE_DB}')
+    logger.info(f"Importing into Supabase database: {config.SUPABASE_DB}, schema: {target_schema}")
 
+    # Set up environment with password
     env = os.environ.copy()
     env['PGPASSWORD'] = password or config.SUPABASE_PASSWORD
 
+    # Create schema if it doesn't exist (and it's not 'public')
+    if target_schema != "public":
+        create_schema_cmd = (
+            f"psql -h {config.SUPABASE_HOST} "
+            f"-p {config.SUPABASE_PORT} "
+            f"-U {config.SUPABASE_USER} "
+            f"-d {config.SUPABASE_DB} "
+            f"--sslmode={config.SUPABASE_SSL_MODE} "
+            f"-c \"CREATE SCHEMA IF NOT EXISTS {target_schema};\""
+        )
+        try:
+            logger.info(f"Creating schema if it doesn't exist: {target_schema}")
+            utils.run_command(create_schema_cmd, env)
+        except Exception as e:
+            logger.warning(f"Failed to create schema, it may already exist: {e}")
 
+    # Build psql command
     cmd = (
         f"psql -h {config.SUPABASE_HOST} "
         f"-p {config.SUPABASE_PORT} "
         f"-U {config.SUPABASE_USER} "
         f"-d {config.SUPABASE_DB} "
         f"--set ON_ERROR_STOP=on "  # Stop on first error
         f"--single-transaction "     # Run as a single transaction
-        f"-f {dump_file}"
     )
 
+    # Set search_path if using non-public schema
+    if target_schema != "public":
+        cmd += f"--set search_path={target_schema} "
+    
+    cmd += f"--sslmode={config.SUPABASE_SSL_MODE} -f {dump_file}"
 
     try:
         utils.run_command(cmd, env)
-        logger.info("import completed successfully")
+        logger.info("Import completed successfully")
     except Exception as e:
-        logger.error(f'import failed: {e}')
+        logger.error(f"Import failed: {e}")
         raise

cloudsql_to_supabase/utils.py

@@ -3,14 +3,11 @@
 from typing import Dict, Optional, List, Union
 import logging
 
+logger = logging.getLogger('cloudsql_to_supabase.utils')
 
-
-logger = logging.getLogger('cloud_to_supabase.utils')
-
-
-def run_command(cmd: str, env: Optional[Dict] = None, show_output: bool = True)-> subprocess.CompletedProcess:
+def run_command(cmd: str, env: Optional[Dict] = None, show_output: bool = True) -> subprocess.CompletedProcess:
     """
-        Run a shell command safely with proper logging and error handling.
+    Run a shell command safely with proper logging and error handling.
     
     Args:
         cmd: Command to run
@@ -20,36 +17,34 @@ def run_command(cmd: str, env: Optional[Dict] = None, show_output: bool = True)-
     Returns:
         CompletedProcess instance
     """
-    
+    # Log the command (with password redacted if present)
     safe_cmd = cmd
     if "PGPASSWORD" in str(env):
-        safe_cmd = cmd.replace(env.get('PGPASSWORD', ''), '************')
-        logger.info(f'Running command: {safe_cmd}')
+        safe_cmd = cmd.replace(env.get('PGPASSWORD', ''), '********')
+    
+    logger.info(f"Running command: {safe_cmd}")
+    
+    try:
+        # Use shlex.split for safer command execution
+        args = shlex.split(cmd)
+        result = subprocess.run(
+            args, 
+            env=env, 
+            stdout=subprocess.PIPE, 
+            stderr=subprocess.PIPE,
+            text=True,
+            check=False  # We'll handle errors ourselves
+        )
 
+        if result.returncode != 0:
+            logger.error(f"Command failed with exit code {result.returncode}")
+            logger.error(result.stderr)
+            raise RuntimeError(f"Command failed: {result.stderr}")
 
-        try: 
-            args = shlex.split(cmd)
-            result = subprocess.run(
-                args,
-                env=env,
-                stdout=subprocess.PIPE,
-                stderr=subprocess.PIPE,
-                text=True,
-                check=False
-            )
-            
-            if result.returncode != 0:
-                logger.error(f'command failed with exit code {result.returncode}')
-                logger.error(result.stderr)
-                raise RuntimeError(f'command failed: {result.stderr}')
+        if show_output and result.stdout:
+            logger.info(result.stdout)
 
-            
-            if show_output and result.stdout:
-                logger.info(result.stdout)
-                
-                
-            return result
-        
-        except Exception as e:
-            logger.exception(f'Error executing command: {e}')
-            raise
+        return result
+    except Exception as e:
+        logger.exception(f"Error executing command: {e}")
+        raise