SQLServerToPostgres.py

from sqlalchemy import create_engine, MetaData, Table, Column, Integer, BigInteger, SmallInteger, Boolean, String, Text, Date, Time, DateTime, Numeric, Float, LargeBinary
from sqlalchemy.dialects.postgresql import UUID
from sqlalchemy.sql import text
import subprocess
from pathlib import Path
import os
from collections import defaultdict    

import logging

class SQLServerToPostgres:
    TYPE_MAP = {
        "int": Integer,
        "bigint": BigInteger,
        "smallint": SmallInteger,
        "tinyint": SmallInteger,
        "bit": Boolean,
        "varchar": String,
        "nvarchar": String,
        "char": String,
        "nchar": String,
        "text": Text,
        "ntext": Text,
        "datetime": DateTime,
        "datetime2": DateTime,
        "smalldatetime": DateTime,
        "date": Date,
        "time": Time,
        "decimal": Numeric,
        "numeric": Numeric,
        "money": Numeric,
        "smallmoney": Numeric,
        "float": Float,
        "real": Float,
        "uniqueidentifier": UUID,
        "binary": LargeBinary,
        "varbinary": LargeBinary,
    }
    
    psql_path=None

    def __init__(self, sqlserver_url, postgres_url, source_schema="dbo", target_schema="public", out_dir="exports"):
        self.sqlserver_engine = create_engine(sqlserver_url)
        self.postgres_engine = create_engine(postgres_url)
        self.source_schema = source_schema
        self.target_schema = target_schema
        self.out_dir = out_dir
        
        self.source_database = self.sqlserver_engine.url.database
        self.source_server = self.sqlserver_engine.url.host
        self.target_database = self.postgres_engine.url.database
        self.target_server = self.postgres_engine.url.host
        self.target_user = self.postgres_engine.url.username
        self.target_password = self.postgres_engine.url.password
        self.target_port = self.postgres_engine.url.port
        
        self.logger = logging.getLogger(self.__class__.__name__)
        
        
    def test_connections(self) :
        s = len(self.get_table_names() )
        self.logger.info(f'{self.sqlserver_engine.url} {s}' )
        
        s = len(self.get_table_names_postgres() )
        self.logger.info(f'{self.postgres_engine.url} {s}' )
    


    def map_type(self, col):
        t = col["type"].lower()
        mapped = self.TYPE_MAP.get(t, Text)
        if t in ("varchar", "nvarchar", "char", "nchar") and col["length"]:
            return mapped(col["length"])
        if t in ("decimal", "numeric", "money", "smallmoney"):
            return mapped(col["precision"], col["scale"])
        return mapped

    def get_columns(self, table_name):
        query = text("""
            SELECT
                COLUMN_NAME = lower(c.COLUMN_NAME),
                c.DATA_TYPE,
                c.CHARACTER_MAXIMUM_LENGTH,
                c.NUMERIC_PRECISION,
                c.NUMERIC_SCALE,
                c.IS_NULLABLE,
                COLUMNPROPERTY(
                    OBJECT_ID(c.TABLE_SCHEMA + '.' + c.TABLE_NAME),
                    c.COLUMN_NAME,
                    'IsIdentity'
                ) AS IS_IDENTITY
            FROM INFORMATION_SCHEMA.COLUMNS c
            WHERE c.TABLE_SCHEMA = :schema
              AND c.TABLE_NAME = :table
            ORDER BY c.ORDINAL_POSITION
        """)
        with self.sqlserver_engine.connect() as conn:
            return conn.execute(query, {"schema": self.source_schema, "table": table_name}).mappings().all()

    def get_primary_keys(self, table_name):
        query = text("""
            SELECT COLUMN_NAME = lower(k.COLUMN_NAME)
            FROM INFORMATION_SCHEMA.TABLE_CONSTRAINTS t
            JOIN INFORMATION_SCHEMA.KEY_COLUMN_USAGE k
              ON t.CONSTRAINT_NAME = k.CONSTRAINT_NAME
            WHERE t.TABLE_SCHEMA = :schema
              AND t.TABLE_NAME = :table
              AND t.CONSTRAINT_TYPE = 'PRIMARY KEY'
            ORDER BY k.ORDINAL_POSITION
        """)
        with self.sqlserver_engine.connect() as conn:
            return set(conn.execute(query, {"schema": self.source_schema, "table": table_name}).scalars().all())

    def create_table(self, table_name):
        '''
        If in the tables, if there is column called urls, and creative_name, it it is not nullable, pleaes change it to nullable because PostRegesql cannot handle NUL(empty string in SQL servr), so, it change to empty string, and it is NULL in PostRegesql. 
        
        Also, if the table is:  
        creative_names
        creative_names_backup
        urls
        urls_backup
        
        Change column "name" is nullable.         

        Parameters
        ----------
        table_name : str
            the table name, it will be lower case in PostGreSQL

        Returns
        -------

            '''
        columns = self.get_columns(table_name)
        pk_columns = self.get_primary_keys(table_name)
        metadata = MetaData(schema=self.target_schema)

        sa_columns = []
        for col in columns:
            sa_type = self.map_type({
                "type": col["DATA_TYPE"],
                "length": col["CHARACTER_MAXIMUM_LENGTH"],
                "precision": col["NUMERIC_PRECISION"],
                "scale": col["NUMERIC_SCALE"],
            })
            
            # Determine nullable status
            is_nullable = col["IS_NULLABLE"] == "YES"
            col_name = col["COLUMN_NAME"].lower()
            
            # If column is "urls" or "creative_name" and NOT nullable, make it nullable
            # (PostgreSQL treats empty strings as NULL)
            if col_name in ("url", "creative_name", 'advertiser_app_id','publisher_app_id') and not is_nullable:
                is_nullable = True
                self.logger.info(f"Changed column {col_name} to nullable in table {table_name} (empty string handling)")
            
            # For specific tables, make "name" column nullable
            if table_name.lower() in ("creative_names", "creative_names_backup", "urls", "urls_backup"):
                if col_name == "name" and not is_nullable:
                    is_nullable = True
                    self.logger.info(f"Changed column 'name' to nullable in table {table_name}")
            
            kwargs = {
                "nullable": is_nullable,
                "primary_key": col["COLUMN_NAME"] in pk_columns
            }
            if col["IS_IDENTITY"]:
                kwargs["autoincrement"] = True

            sa_columns.append(Column(col["COLUMN_NAME"], sa_type, **kwargs))

        Table(table_name, metadata, *sa_columns)
        metadata.create_all(self.postgres_engine)
        self.logger.info(f"Table {self.target_schema}.{table_name} created successfully.")


    
    def get_table_names(self ):
        """
        Return all table names in the current SQL Server database
        for the given schema.
    
        The table name is not case sensitive in SQL, but it is in Postg
    
        :param engine: SQLAlchemy Engine (connected to target DB)
        :param source_schema: Schema name (e.g. 'dbo')
        :return: List[str] of table names
        """
        sql = text("""
            SELECT name = lower(t.name)
            FROM sys.tables t
            JOIN sys.schemas s ON t.schema_id = s.schema_id
            WHERE s.name = :schema
            ORDER BY t.name
        """)
    
        with self.sqlserver_engine.connect() as conn:
            return conn.execute(sql, {"schema": self.source_schema}).scalars().all()


    def get_table_names_postgres(self):
        """
        Return all table names in the current PostgreSQL database
        for the given schema.
        """
        sql = text("""
            SELECT tablename
            FROM pg_catalog.pg_tables
            WHERE schemaname = :schema
            ORDER BY tablename
        """)
    
        with self.postgres_engine.connect() as conn:
            return conn.execute(sql, {"schema": self.target_schema}).scalars().all()
    
    #create all tables 
    def created_all_tables_and_primary_key(self) :
        for table_name in self.get_table_names():
            self.create_table(table_name)
            
    #create all tables 
    def copy_all_tables_data(self) :
        
        for table_name in self.get_table_names():
            lower_table_name = table_name.lower()
            
            # if lower_table_name not in (
            #         # 'ads'
            #         # , 
            #         # 'pie_stage'
            #         # , 
                    # 'pie'
                    # ) :
                # continue
            

            self.logger.info(f"bcp process table: {table_name}")
            
            self.copy_one_table(table_name, lower_table_name)
                
    def copy_one_table(self, table_name, lower_table_name) :
        try : 
            self.logger.info(f"bcp process table: {table_name}")
            csv_file_new =''
            limit = -1
            raw_count = self.get_source_table_row_count(table_name)
            csv_file=self.bcp_export_top(table_name, limit=limit)
            # csv_file=self.bcp_export(table_name)
            
   
            # self.psql_copy_csv_to_postgres(lower_table_name, csv_file)
            count = self.get_target_table_row_count(lower_table_name) 
            
            if count != raw_count and count != limit :
                raise Exception(f'rowcount do not match after batch copy {count} {raw_count} or 100')
            self.logger.info(f"completed copying: {table_name} to {lower_table_name} {raw_count} vs {count} ")
        except Exception as ex:
            self.logger.error(f"error process table: {table_name} {raw_count} {csv_file_new}")
            self.logger.exception(ex)

    def copy_all_tables_data_bcp_psql(self):
        """Copy all tables: BCP export from SQL Server then load into PostgreSQL via psql."""
        for table_name in self.get_table_names():
            lower_table_name = table_name.lower()
            self.logger.info(f"BCP then PSQL: {table_name}")
            csv_file = self.bcp_export_top(table_name, limit=-1)
            self.psql_copy_csv_to_postgres(lower_table_name, csv_file)
            raw_count = self.get_source_table_row_count(table_name)
            count = self.get_target_table_row_count(lower_table_name)
            self.logger.info(f"completed: {table_name} -> {lower_table_name} ({raw_count} vs {count})")

    def bcp_export_all_tables_only(self):
        """BCP output data only for all tables (no PSQL load)."""
        for table_name in self.get_table_names():
            self.logger.info(f"BCP export only: {table_name}")
            self.bcp_export(table_name)
            
    def bcp_export(self, table):
        Path(self.out_dir).mkdir(exist_ok=True)
        output_file = Path(self.out_dir) / f"{table}.csv"

        cmd = [
            "bcp",
            f"{self.source_database}.{self.source_schema}.{table}",
            "out",
            output_file,
            "-c",
            "-t,",
            "-r\n",
            "-C", "65001",
            "-T",
            "-k",
            "-S", self.source_server
        ]
    
        self._run_subprocess(cmd)
        self.logger.info('saved file:' + str(output_file))
        return str(output_file)

    def _build_isnull_empty_select_list(self, table_name):
        """
        Build a SELECT list that handles string columns for CSV encoding.
        
        For string columns:
        - Empty strings are converted to a space ' '
        - Only wrap in quotes and escape if the string contains special characters
          (comma, double quote, newline, or carriage return)
        - If no special characters, leave as-is
        
        For binary/varbinary columns:
        - If NULL, convert to empty string ''
        - If NOT NULL but length is 0, convert to number 0
        - Otherwise use as-is
        
        For other non-string columns, they are used as-is.
        
        Parameters
        ----------
        table_name : str
            The table name to build the select list for
            
        Returns
        -------
        str
            Comma-separated SELECT list with CSV encoding for string columns
        """
        # Get columns with original case preserved
        query = text("""
            SELECT
                c.COLUMN_NAME,
                c.DATA_TYPE
            FROM INFORMATION_SCHEMA.COLUMNS c
            WHERE c.TABLE_SCHEMA = :schema
              AND c.TABLE_NAME = :table
            ORDER BY c.ORDINAL_POSITION
        """)
        
        with self.sqlserver_engine.connect() as conn:
            columns = conn.execute(query, {"schema": self.source_schema, "table": table_name}).mappings().all()
        
        string_types = {'varchar', 'nvarchar', 'char', 'nchar', 'text', 'ntext'}
        binary_types = {'binary', 'varbinary'}
        
        select_parts = []
        for col in columns:
            col_name = col["COLUMN_NAME"]
            data_type = col["DATA_TYPE"].lower()
            
            if data_type in string_types:
                # For string columns:
                # 1. Convert empty/NULL to space
                # 2. Check if contains special chars (comma, quote, newline, carriage return)
                # 3. Only wrap in quotes and escape if special chars are present
                # Special chars that need CSV encoding: , " \n \r
                select_parts.append(
                    f"CASE "
                    f"WHEN [{col_name}] = '' THEN ' ' "
                    f"WHEN CHARINDEX(',', [{col_name}]) > 0 "
                    f"     OR CHARINDEX('\"', [{col_name}]) > 0 "
                    f"     OR CHARINDEX(CHAR(10), [{col_name}]) > 0 "
                    f"     OR CHARINDEX(CHAR(13), [{col_name}]) > 0 "
                    f"THEN '\"' + REPLACE([{col_name}], '\"', '\"\"') + '\"' "
                    f"ELSE [{col_name}] "
                    f"END AS [{col_name}]"
                )
            elif data_type in binary_types:
                # For binary/varbinary columns:
                # If NULL, return empty string
                # If NOT NULL but length is 0, return 0
                # Otherwise use as-is
                select_parts.append(
                    f"CASE "
                    f"WHEN [{col_name}] IS NULL THEN '' "
                    f"WHEN DATALENGTH([{col_name}]) = 0 THEN 0 "
                    f"ELSE [{col_name}] "
                    f"END AS [{col_name}]"
                )
            else:
                # For non-string columns, use as-is
                select_parts.append(f"[{col_name}]")
        
        return ", ".join(select_parts)

    def bcp_export_top(self, table_name, limit=100):
        
        Path(self.out_dir).mkdir(exist_ok=True)
        outfile = Path(self.out_dir) / f"{table_name}_top{limit}.csv"
    
        select_list = self._build_isnull_empty_select_list(table_name)
    
        select_top_sentence = f'TOP ({limit})' if limit >0 else ''
            
        query = (
            f"SELECT {select_top_sentence} {select_list} "
            f"FROM [{self.source_database}].[{self.source_schema}].[{table_name}] "
        )
    
        cmd = [
            "bcp",
            query,
            "queryout",
            str(outfile),
            "-c",
            "-t,",
            "-r\n",
            "-C", "65001",
            "-T",
            "-S", self.source_server
        ]
    
        self._run_subprocess(cmd)
        self.logger.info(f"Saved file: {outfile}")
        return str(outfile)
    

            
    def _run_subprocess(self, cmd ) :
        self.logger.info(f"run subprocess {cmd}")
        text = ' '.join([str(x) for x in cmd] ) 
        self.logger.info(f"run subprocess cmd: {text}")

        process = subprocess.Popen(
            cmd,
            stdout=subprocess.PIPE,
            stderr=subprocess.STDOUT,   # 👈 merge stderr
            text=True,
            bufsize=1
        )

        for line in iter(process.stdout.readline, ''):
            self.logger.info(line)
    
        process.stdout.close()
        rc = process.wait()

        if rc != 0:
            raise RuntimeError(f"Command failed with exit code {rc}")

        
        self.logger.info(f'subprocess {process.returncode}')
        if process.returncode != 0:
            raise RuntimeError(process.stderr)
        
        return process.stdout
 
    def psql_copy_csv_to_postgres( self,
         pg_table: str,
         csv_file: str,
         has_header: bool = False
             ):
        self.logger.info('loading file:' + csv_file)
        csv_file = Path(csv_file)
    
        copy_cmd = (
            f"\\copy {self.target_schema}.{pg_table} "
            f"FROM '{csv_file.as_posix()}' "
            f"CSV {'HEADER' if has_header else ''}"
        )
    
        if self.target_password :
            os.environ["PGPASSWORD"] = self.target_password
        cmd = [
            self.psql_path if self.psql_path else "psql",
            "-h", self.target_server,
            "-p", str(self.target_port),
            "-U", self.target_user,
            "-d", self.target_database,
            "-c", copy_cmd,
        ]
    
        self._run_subprocess(cmd)
    
    def get_source_table_row_count(self, table_name) :
        return self.get_table_row_count_mssql(self.sqlserver_engine, table_name, self.source_schema)
    
    def get_target_table_row_count(self, table_name) :
        return self.get_table_row_count(self.postgres_engine, table_name, self.target_schema)


    def get_table_row_count_mssql(self, engine, table_name: str, schema: str = "dbo") -> int:
        """
        Get the total number of rows in a table.
    
        Args:
            engine (Engine): SQLAlchemy engine connected to the database.
            table_name (str): Name of the table.
            schema (str): Schema name. Defaults to 'dbo'.
    
        Returns:
            int: Total number of rows in the table.
        """
        query = text(f"SELECT COUNT_BIG(*) AS cnt FROM [{schema}].[{table_name}]")
        with engine.connect() as conn:
            result = conn.execute(query)
            row_count = result.scalar()  # Returns the single integer
        return row_count
    
    def get_table_row_count(self, engine, table_name: str, schema: str = "dbo") -> int:
        """
        Get the total number of rows in a table.
    
        Args:
            engine (Engine): SQLAlchemy engine connected to the database.
            table_name (str): Name of the table.
            schema (str): Schema name. Defaults to 'dbo'.
    
        Returns:
            int: Total number of rows in the table.
        """
        query = text(f'SELECT COUNT(*) AS cnt FROM {schema}."{table_name}"')
        with engine.connect() as conn:
            result = conn.execute(query)
            row_count = result.scalar()  # Returns the single integer
        return row_count

    def check_table_row_counts(self, table_names=None, show_matching=True):
        """
        Check and compare row counts for tables in SQL Server and PostgreSQL.
        
        Args:
            table_names (list, optional): List of table names to check. 
                                         If None, checks all tables in source schema.
            show_matching (bool): If True, shows tables with matching counts.
                                If False, only shows mismatches. Defaults to True.
        
        Returns:
            dict: Dictionary with table names as keys and dicts containing:
                - 'source': row count in SQL Server
                - 'target': row count in PostgreSQL
                - 'match': boolean indicating if counts match
                - 'difference': difference in row counts (target - source)
        """
        if table_names is None:
            table_names = self.get_table_names()
        
        results = {}
        
        self.logger.info(f"Checking row counts for {len(table_names)} table(s)...")
        
        for table_name in table_names:
            try:
                lower_table_name = table_name.lower()
                
                # Get source count
                source_count = self.get_source_table_row_count(table_name)
                
                # Get target count (may not exist)
                try:
                    target_count = self.get_target_table_row_count(lower_table_name)
                except Exception as e:
                    # Table might not exist in target
                    target_count = None
                    self.logger.warning(f"Table {lower_table_name} not found in PostgreSQL: {e}")
                
                match = (target_count == source_count) if target_count is not None else False
                difference = (target_count - source_count) if target_count is not None else None
                
                results[table_name] = {
                    'source': source_count,
                    'target': target_count,
                    'match': match,
                    'difference': difference
                }
                
                # Log the result
                if target_count is None:
                    status = "NOT FOUND"
                    self.logger.info(f"{table_name:50s} | SQL Server: {source_count:>15,} | PostgreSQL: {status:>15s}")
                elif match:
                    if show_matching:
                        self.logger.info(f"{table_name:50s} | SQL Server: {source_count:>15,} | PostgreSQL: {target_count:>15,} | MATCH")
                else:
                    self.logger.warning(f"{table_name:50s} | SQL Server: {source_count:>15,} | PostgreSQL: {target_count:>15,} | XXX MISMATCH (diff: {difference:+,})")
                    
            except Exception as e:
                self.logger.error(f"Error checking table {table_name}: {e}")
                results[table_name] = {
                    'source': None,
                    'target': None,
                    'match': False,
                    'difference': None,
                    'error': str(e)
                }
        
        # Summary
        total_tables = len(results)
        matching = sum(1 for r in results.values() if r.get('match', False))
        mismatching = sum(1 for r in results.values() if r.get('target') is not None and not r.get('match', False))
        not_found = sum(1 for r in results.values() if r.get('target') is None)
        
        self.logger.info(f"\nSummary:")
        self.logger.info(f"  Total tables checked: {total_tables}")
        self.logger.info(f"  Matching counts: {matching}")
        self.logger.info(f"  Mismatching counts: {mismatching}")
        self.logger.info(f"  Not found in target: {not_found}")
        
        return results


    def drop_all_tables_in_target(self) :
        for table_name in self.get_table_names_postgres() :
            self.drop_a_table_in_target(table_name)

    def drop_a_table_in_target(self, table_name):
        """
        Drops all tables in the specified schema.
        
        Args:
            engine (Engine): SQLAlchemy engine connected to the database.
            schema (str): The schema from which to drop all tables.
        """
        table_row_count = self.get_target_table_row_count(table_name)
        if table_row_count > 0:
            # Table is not empty: require user to type table name and server name to confirm
            expected_server = self.target_server or ""
            self.logger.warning(
                f"Table '{table_name}' is not empty ({table_row_count} row(s)). "
                f"To confirm deletion, type the table name and target server name "
                f"(table name: {table_name}, server: {expected_server})."
            )
            typed_table = input(f"Type the table name to confirm: ").strip()
            typed_server = input(f"Type the target server name to confirm: ").strip()
            if typed_table != table_name or typed_server != expected_server:
                raise Exception(
                    f"Deletion aborted: confirmation failed. "
                    f"Expected table '{table_name}' and server '{expected_server}', "
                    f"got table '{typed_table}' and server '{typed_server}'."
                )
            self.logger.info("Confirmation accepted. Proceeding with drop.")
            
        with self.postgres_engine.connect() as conn:
            # Disable FK constraints temporarily
    
            self.logger.info(f'Dropping table {self.target_schema}."{table_name}"')
            conn.execute(text(f'DROP TABLE IF EXISTS "{self.target_schema}"."{table_name}" CASCADE;'))
    
            self.logger.info(f"Table {table_name} in schema '{self.target_schema}' have been dropped.")
            
    
    
    def copy_all_tables_data_directly(self) :
        
 
        for table_name in self.get_table_names():
 

            lower_table_name = table_name.lower()
            # if lower_table_name not in ('urls'
            #                             , 'urls_backup'
            #                             ,'device_properties'

            #         ) :
                # continue 
 
            limit = 100000
            
            self.transfer_table(lower_table_name, lower_table_name, limit)

            raw_count = self.get_source_table_row_count(table_name)
            count = self.get_target_table_row_count(lower_table_name) 
            
            self.logger.info(f"completed copying: {table_name} to {lower_table_name} {raw_count} vs {count} ")
            if count != raw_count  :
                self.logger.error(f"error copying: {table_name} to {lower_table_name} {raw_count} vs {count} ")
                # raise Exception(f'rowcount do not match after batch copy {count} {raw_count} or 100')


    def stream_sqlserver_rows(self, query, chunk_size=10000):
        with self.sqlserver_engine.connect().execution_options(stream_results=True) as conn:
            result = conn.execute(text(query))
            while True:
                rows = result.fetchmany(chunk_size)
                if not rows:
                    break
                yield rows, result.keys()
    
    def insert_into_postgres(self, table_name, columns, rows):
        placeholders = ", ".join([f":{c}" for c in columns])
        col_list = ", ".join(columns)
        
        sql = text(f"""
            INSERT INTO {self.target_schema}.{table_name} ({col_list})
            VALUES ({placeholders})
        """)
        
        data = [dict(zip(columns, row)) for row in rows]
        
        with self.postgres_engine.begin() as conn:
            conn.execute(sql, data)

    
    
    def transfer_table(
            self,
            source_table,
            target_table,
            chunk_size=10000
        ):
        query = f"SELECT * FROM {self.source_schema}.{source_table}"
        
        total = 0
        
        for rows, columns in self.stream_sqlserver_rows(query, chunk_size):
            self.insert_into_postgres(
                target_table,
                columns,
                rows
            )
            total += len(rows)
            self.logger.info(f"Inserted {total:,} rows")
        
        self.logger.info("Transfer complete")


    
    def add_all_indexes_and_foreign_key(self, table_name =None):
        """Create all indexes on target tables (after data is copied). Override or extend to implement."""
        self.add_all_indexes(table_name)
        self.add_all_foreign_keys(table_name)


    def add_all_indexes(self, table_name):
        """
        change table name and column name to lower 
        Recreate all non-PK indexes from SQL Server into PostgreSQL.
        """
        self.logger.info(f"Adding indexes to target database: {table_name}")
        
        where_clause = "" if not table_name else f" AND t.name =  '{table_name}' "
        
    
        sql = f""" SELECT
    s.name  AS schema_name,
    t.name  AS table_name,
    i.name  AS index_name,
    i.is_unique,
    ic.key_ordinal,
    c.name  AS column_name
FROM sys.indexes i
JOIN sys.index_columns ic
    ON i.object_id = ic.object_id AND i.index_id = ic.index_id
JOIN sys.columns c
    ON ic.object_id = c.object_id AND ic.column_id = c.column_id
JOIN sys.tables t
    ON i.object_id = t.object_id
JOIN sys.schemas s
    ON t.schema_id = s.schema_id
WHERE
    i.is_primary_key = 0
    AND i.is_disabled = 0
    AND ic.is_included_column = 0
    AND s.name = '{self.source_schema}'
    {where_clause}
ORDER BY
    s.name, t.name, i.name, ic.key_ordinal;
"""
 
    
        with self.sqlserver_engine.connect() as src, \
             self.postgres_engine.begin() as tgt:
    
            rows = src.execute(sql).fetchall()
    
            indexes = defaultdict(list)
    
            for r in rows:
                key = (r.schema_name, r.table_name, r.index_name, r.is_unique)
                indexes[key].append(r.column_name)
    

            for (schema, table, index, is_unique), cols in indexes.items():
                col_list = ", ".join(f'"{c.lower()}"' for c in cols)
                unique = "UNIQUE" if is_unique else ""
                table = table.lower()
    
                pg_index = f"{index}".lower()
    
                stmt = f"""
                    CREATE {unique} INDEX IF NOT EXISTS "{pg_index}"
                    ON "{self.target_schema}"."{table}" ({col_list});
                """
    
                self.logger.info(stmt.strip())
                tgt.execute(stmt)
    
        self.logger.info("Indexes added successfully")

    def add_all_foreign_keys(self, table_name):
        """
        Recreate all foreign key constraints in PostgreSQL.
        """
        self.logger.info("Adding foreign keys to target database")
        where_clause = "" if not table_name else f" AND tab1.name =  '{table_name}' "
    
        sql = f""" 
        SELECT
    fk.name AS fk_name,
    sch1.name AS schema_name,
    tab1.name AS table_name,
    col1.name AS column_name,
    sch2.name AS ref_schema,
    tab2.name AS ref_table,
    col2.name AS ref_column,
    fk.delete_referential_action_desc AS on_delete,
    fk.update_referential_action_desc AS on_update
FROM sys.foreign_keys fk
JOIN sys.foreign_key_columns fkc
    ON fk.object_id = fkc.constraint_object_id
JOIN sys.tables tab1
    ON fkc.parent_object_id = tab1.object_id
JOIN sys.schemas sch1
    ON tab1.schema_id = sch1.schema_id
JOIN sys.columns col1
    ON fkc.parent_object_id = col1.object_id
   AND fkc.parent_column_id = col1.column_id
JOIN sys.tables tab2
    ON fkc.referenced_object_id = tab2.object_id
JOIN sys.schemas sch2
    ON tab2.schema_id = sch2.schema_id
JOIN sys.columns col2
    ON fkc.referenced_object_id = col2.object_id
   AND fkc.referenced_column_id = col2.column_id
WHERE 
    sch1.name = '{self.source_schema}'
   {where_clause}
ORDER BY fk.name, fkc.constraint_column_id;

"""
    
        with self.sqlserver_engine.connect() as src, \
             self.postgres_engine.begin() as tgt:
    
            rows = src.execute(sql).fetchall()
    
            fks = defaultdict(list)
            actions = {}
    
            for r in rows:
                key = (r.schema_name, r.table_name, r.fk_name)
                fks[key].append((r.column_name, r.ref_schema, r.ref_table, r.ref_column))
                actions[key] = (r.on_delete, r.on_update)
    
            for (schema, table, fk), cols in fks.items():
                table = table.lower()
                cols_local = ", ".join(f'"{c[0].lower()}"' for c in cols)
                cols_ref = ", ".join(f'"{c[3].lower()}"' for c in cols)
    
                ref_schema = cols[0][1]
                ref_table = cols[0][2]
    
                on_delete, on_update = actions[(schema, table, fk)]
    
                stmt = f"""
                    ALTER TABLE "{self.target_schema}"."{table}"
                    ADD CONSTRAINT "{fk.lower()}"
                    FOREIGN KEY ({cols_local})
                    REFERENCES "{ref_schema}"."{ref_table.lower()}" ({cols_ref})
                    ON DELETE {on_delete.replace('_', ' ')}
                    ON UPDATE {on_update.replace('_', ' ')};
                """
    
                self.logger.info(stmt.strip())
                tgt.execute(stmt)
    
        self.logger.info("Foreign keys added successfully")
    
    
 

    def get_sqlserver_indexes(self):
        sql = """
        SELECT
            sch.name AS schema_name,
            t.name AS table_name,
            i.name AS index_name,
            i.is_unique,
            ic.key_ordinal,
            c.name AS column_name
        FROM sys.indexes i
        JOIN sys.index_columns ic
            ON i.object_id = ic.object_id
           AND i.index_id = ic.index_id
        JOIN sys.columns c
            ON ic.object_id = c.object_id
           AND ic.column_id = c.column_id
        JOIN sys.tables t
            ON i.object_id = t.object_id
        JOIN sys.schemas sch
            ON t.schema_id = sch.schema_id
        WHERE sch.name = ?
          AND i.is_primary_key = 0
          AND i.is_disabled = 0
        ORDER BY sch.name, t.name, i.name, ic.key_ordinal
        """
    
        rows = self.sqlserver_engine.execute(sql, self.source_schema).fetchall()
    
        index_map = defaultdict(lambda: {
            "schema": None,
            "table_name": None,
            "index_name": None,
            "is_unique": False,
            "columns": []
        })
    
        for r in rows:
            key = (r.schema_name, r.table_name, r.index_name)
            idx = index_map[key]
            idx["schema"] = r.schema_name
            idx["table_name"] = r.table_name
            idx["index_name"] = r.index_name
            idx["is_unique"] = bool(r.is_unique)
            idx["columns"].append(r.column_name)
    
        return list(index_map.values())


    def get_sqlserver_foreign_keys(self):
        sql = """
        SELECT
            sch.name AS schema_name,
            t.name AS table_name,
            fk.name AS constraint_name,
            c.name AS column_name,
            ref_sch.name AS ref_schema,
            rt.name AS ref_table,
            rc.name AS ref_column,
            fkc.constraint_column_id
        FROM sys.foreign_keys fk
        JOIN sys.foreign_key_columns fkc
            ON fk.object_id = fkc.constraint_object_id
        JOIN sys.tables t
            ON fk.parent_object_id = t.object_id
        JOIN sys.schemas sch
            ON t.schema_id = sch.schema_id
        JOIN sys.columns c
            ON fkc.parent_object_id = c.object_id
           AND fkc.parent_column_id = c.column_id
        JOIN sys.tables rt
            ON fk.referenced_object_id = rt.object_id
        JOIN sys.schemas ref_sch
            ON rt.schema_id = ref_sch.schema_id
        JOIN sys.columns rc
            ON fkc.referenced_object_id = rc.object_id
           AND fkc.referenced_column_id = rc.column_id
        WHERE sch.name = ?
        ORDER BY fk.name, fkc.constraint_column_id
        """
    
        rows = self.sqlserver_engine.execute(sql, self.source_schema).fetchall()
    
        fk_map = defaultdict(lambda: {
            "schema": None,
            "table_name": None,
            "constraint_name": None,
            "columns": [],
            "ref_schema": None,
            "ref_table": None,
            "ref_columns": []
        })
    
        for r in rows:
            key = (r.schema_name, r.table_name, r.constraint_name)
            fk = fk_map[key]
            fk["schema"] = r.schema_name
            fk["table_name"] = r.table_name
            fk["constraint_name"] = r.constraint_name
            fk["ref_schema"] = r.ref_schema
            fk["ref_table"] = r.ref_table
            fk["columns"].append(r.column_name)
            fk["ref_columns"].append(r.ref_column)
    
        return list(fk_map.values())
    

    def generate_pg_indexes_and_fks(self,
        output_file: str,
        schema_mapping: dict=None,
        default_schema: str = "public"
    ):
        """
        Generate PostgreSQL SQL for indexes and foreign keys
        using schema mapping (SQL Server -> PostgreSQL).
        All identifiers are lower-case.
        """
    
    
        if not dict :
            dict['dbo'] ='public'
     
        index_defs =self.get_sqlserver_indexes()
        fk_defs  =self.get_sqlserver_foreign_keys()
        
        
        def map_schema(src_schema: str) -> str:
            if not src_schema:
                return self.target_schema
            return schema_mapping.get(src_schema.lower(), self.target_schema).lower()
    
        with open(output_file, "w", encoding="utf-8") as f:
    
            # =====================
            # Indexes
            # =====================
            f.write("-- =====================\n")
            f.write("-- Indexes\n")
            f.write("-- =====================\n\n")
    
            for idx in index_defs:
                src_schema = idx.get("schema", "")
                tgt_schema = map_schema(src_schema)
    
                table = idx["table_name"].lower()
                index_name = idx["index_name"].lower()
                self.logger.info(f'Index: {table} {index_name} ')
                columns = ", ".join(c.lower() for c in idx["columns"])
                unique = "UNIQUE " if idx.get("is_unique") else ""
    
                sql = (
                    f"CREATE {unique}INDEX IF NOT EXISTS {index_name} "
                    f"ON {tgt_schema}.{table} ({columns});\n"
                )
                f.write(sql)
    
            # =====================
            # Foreign Keys
            # =====================
            f.write("\n-- =====================\n")
            f.write("-- Foreign Keys\n")
            f.write("-- =====================\n\n")
    
            for fk in fk_defs:
                src_schema = fk.get("schema", "")
                ref_schema = fk.get("ref_schema", "")
    
                tgt_schema = map_schema(src_schema)
                tgt_ref_schema = map_schema(ref_schema)
    
                table = fk["table_name"].lower()
                constraint = fk["constraint_name"].lower()
                ref_table = fk["ref_table"].lower()
                self.logger.info(f'Foreign key: {table} {constraint} ')
                src_cols = ", ".join(c.lower() for c in fk["columns"])
                ref_cols = ", ".join(c.lower() for c in fk["ref_columns"])
    
                sql = (
                    f"ALTER TABLE {tgt_schema}.{table}\n"
                    f"  ADD CONSTRAINT {constraint}\n"
                    f"  FOREIGN KEY ({src_cols})\n"
                    f"  REFERENCES {tgt_ref_schema}.{ref_table} ({ref_cols});\n\n"
                )
                f.write(sql)