Allow override of db name (#89)

Author: BrianLondon

druzhba/table.py

@@ -174,6 +174,10 @@ class TableConfig(object):
     include_comments : boolean, optional
         flag to specify whether or not to ingest table and column comments
         when building or rebuilding the target table.
+    override_db_name : str, optional
+        override database name to use in the index tracking table instead of
+        the actual database name. This allows tracking multiple databases or
+        environments under a common name for index purposes.
 
     Attributes
     ----------
@@ -241,6 +245,7 @@ def __init__(
         include_comments=True,
         monitor_tables_config=None,
         lookback_value=0,
+        override_db_name=None,
     ):
         self.database_alias = database_alias
         self.db_host = db_connection_params.host
@@ -282,6 +287,7 @@ def __init__(
         self.monitor_tables_config = monitor_tables_config
         self.lookback_value = lookback_value
         self._lookback_index_value = "notset"
+        self.override_db_name = override_db_name
 
         self.date_key = datetime.datetime.strftime(
             datetime.datetime.utcnow(), "%Y%m%dT%H%M%S"
@@ -301,6 +307,14 @@ def __init__(
         self.logger = logging.getLogger(f"druzhba.{database_alias}.{source_table_name}")
         self.s3 = Session().client("s3")
 
+    @property
+    def index_db_name(self):
+        """Returns the database name to use for index tracking operations.
+        
+        Returns override_db_name if set, otherwise returns the actual db_name.
+        """
+        return self.override_db_name if self.override_db_name else self.db_name
+
     @classmethod
     def _clean_type_map(cls, type_map):
         if not type_map:
@@ -579,7 +593,7 @@ def _load_old_index_value(self):
         self.logger.debug("Querying Redshift for last updated index")
         with get_redshift().cursor() as cur:
             cur.execute(
-                query, (self.database_alias, self.db_name, self.source_table_name)
+                query, (self.database_alias, self.index_db_name, self.source_table_name)
             )
             index_value = cur.fetchone()
 
@@ -627,7 +641,7 @@ def _load_lookback_index_value(self):
         self.logger.debug("Querying Redshift for nth last updated index to lookback")
         with get_redshift().cursor() as cur:
             cur.execute(
-                query, (self.database_alias, self.db_name, self.source_table_name)
+                query, (self.database_alias, self.index_db_name, self.source_table_name)
             )
             index_value = cur.fetchone()
 
@@ -849,7 +863,7 @@ def set_last_updated_index(self):
         with get_redshift().cursor() as cur:
             args = (
                 self.database_alias,
-                self.db_name,
+                self.index_db_name,
                 self.source_table_name,
                 new_index_value,
             )

test/unit/test_table.py

@@ -1087,3 +1087,51 @@ def test_parse_invalid(self):
         x = "something"
         output = Permissions.parse(x)
         self.assertEqual(output, None)
+
+    def test_override_db_name(self):
+        """Test that override_db_name parameter works correctly for index tracking"""
+        from druzhba.table import TableConfig
+        from druzhba.db import ConnectionParams
+        
+        connection_params = ConnectionParams(
+            name="test_db",
+            host="localhost", 
+            port=5432,
+            user="test_user",
+            password="test_pass",
+            additional={}
+        )
+        
+        # Test without override - should use actual db_name
+        table_config = TableConfig(
+            database_alias="test_alias",
+            db_connection_params=connection_params,
+            destination_table_name="dest_table",
+            destination_schema_name="public",
+            source_table_name="source_table",
+            index_schema="idx_schema",
+            index_table="idx_table"
+        )
+        
+        self.assertEqual(table_config.index_db_name, "test_db")
+        self.assertEqual(table_config.db_name, "test_db")
+        
+        # Test with override - should use override_db_name
+        table_config_with_override = TableConfig(
+            database_alias="test_alias",
+            db_connection_params=connection_params,
+            destination_table_name="dest_table",
+            destination_schema_name="public",
+            source_table_name="source_table",
+            index_schema="idx_schema",
+            index_table="idx_table",
+            override_db_name="override_name"
+        )
+        
+        self.assertEqual(table_config_with_override.index_db_name, "override_name")
+        self.assertEqual(table_config_with_override.db_name, "test_db")
+        self.assertEqual(table_config_with_override.override_db_name, "override_name")
+
+
+if __name__ == "__main__":
+    unittest.main()