Add a flag to allow to specify database for pg extractor (#107)

databuilder/extractor/postgres_metadata_extractor.py

@@ -41,6 +41,7 @@ class PostgresMetadataExtractor(Extractor):
     WHERE_CLAUSE_SUFFIX_KEY = 'where_clause_suffix'
     CLUSTER_KEY = 'cluster_key'
     USE_CATALOG_AS_CLUSTER_NAME = 'use_catalog_as_cluster_name'
+    DATABASE_KEY = 'database_key'
 
     # Default values
     DEFAULT_CLUSTER_NAME = 'master'
@@ -59,6 +60,9 @@ class PostgresMetadataExtractor(Extractor):
         else:
             cluster_source = "'{}'".format(self._cluster)
 
+        self._database = conf.get_string(PostgresMetadataExtractor.DATABASE_KEY,
+                                         default='postgres').encode('utf-8', 'ignore')
+
         self.sql_stmt = PostgresMetadataExtractor.SQL_STATEMENT.format(
             where_clause_suffix=conf.get_string(PostgresMetadataExtractor.WHERE_CLAUSE_SUFFIX_KEY),
             cluster_source=cluster_source
@@ -100,7 +104,7 @@ class PostgresMetadataExtractor(Extractor):
                 columns.append(ColumnMetadata(row['col_name'], row['col_description'],
                                               row['col_type'], row['col_sort_order']))
 
-            yield TableMetadata('postgres', last_row['cluster'],
+            yield TableMetadata(self._database, last_row['cluster'],
                                 last_row['schema_name'],
                                 last_row['name'],
                                 last_row['description'],

setup.py

 from setuptools import setup, find_packages
 
 
-__version__ = '1.3.6'
+__version__ = '1.3.7'
 
 
 setup(

tests/unit/extractor/test_postgres_metadata_extractor.py

@@ -21,7 +21,9 @@ class TestPostgresMetadataExtractor(unittest.TestCase):
             'extractor.postgres_metadata.{}'.format(PostgresMetadataExtractor.CLUSTER_KEY):
             'MY_CLUSTER',
             'extractor.postgres_metadata.{}'.format(PostgresMetadataExtractor.USE_CATALOG_AS_CLUSTER_NAME):
-            False
+            False,
+            'extractor.postgres_metadata.{}'.format(PostgresMetadataExtractor.DATABASE_KEY):
+            'postgres'
         }
         self.conf = ConfigFactory.from_dict(config_dict)
 
@@ -94,6 +96,7 @@ class TestPostgresMetadataExtractor(unittest.TestCase):
                                       ColumnMetadata('source', 'description of source', 'varchar', 3),
                                       ColumnMetadata('etl_created_at', 'description of etl_created_at', 'timestamp', 4),
                                       ColumnMetadata('ds', None, 'varchar', 5)])
+
             self.assertEqual(expected.__repr__(), actual.__repr__())
             self.assertIsNone(extractor.extract())