feat: connection_string config property

tap_postgres/__init__.py

@@ -1,6 +1,7 @@
 import argparse
 import itertools
 import copy
+import urllib.parse
 import psycopg2
 import psycopg2.extras
 import psycopg2.extensions
@@ -20,15 +21,21 @@
 from tap_postgres.stream_utils import (
     dump_catalog, clear_state_on_replication_change,
     is_selected_via_metadata, refresh_streams_schema, any_logical_streams)
+from tap_postgres.config_utils import check_config_for_key_groups
 
 LOGGER = singer.get_logger('tap_postgres')
 
-REQUIRED_CONFIG_KEYS = [
-    'dbname',
-    'host',
-    'port',
-    'user',
-    'password'
+REQUIRED_CONFIG_KEYS_GROUPS = [
+    [
+        'dbname',
+        'host',
+        'port',
+        'user',
+        'password'
+    ],
+    [
+        'connection_string'
+    ]
 ]
 
 
@@ -381,23 +388,39 @@ def parse_args(required_config_keys):
         setattr(args, 'catalog_path', args.catalog)
         args.catalog = Catalog.load(args.catalog)
 
-    utils.check_config(args.config, required_config_keys)
+    check_config_for_key_groups(args.config, required_config_keys)
 
     return args
 
 
-def main_impl():
+def main_impl(): # pragma: no cover
     """
     Main method
     """
-    args = parse_args(REQUIRED_CONFIG_KEYS)
+    args = parse_args(REQUIRED_CONFIG_KEYS_GROUPS)
+
+    if 'connection_uri' in args.config:
+        parsed_conn_str = urllib.parse.urlparse(args.config['connection_uri'])
+
+        host = parsed_conn_str.hostname
+        dbname = parsed_conn_str.path[1:] # Strip starting "/"
+        user = parsed_conn_str.username
+        password = parsed_conn_str.password
+        port = parsed_conn_str.port
+    else:
+        host = args.config['host']
+        dbname = args.config['dbname']
+        user = args.config['user']
+        password = args.config['password']
+        port = args.config['port']
+
     conn_config = {
         # Required config keys
-        'host': args.config['host'],
-        'user': args.config['user'],
-        'password': args.config['password'],
-        'port': args.config['port'],
-        'dbname': args.config['dbname'],
+        'host': host,
+        'user': user,
+        'password': password,
+        'port': port,
+        'dbname': dbname,
 
         # Optional config keys
         'tap_id': args.config.get('tap_id'),

tap_postgres/config_utils.py

@@ -0,0 +1,18 @@
+def check_config_for_key_groups(config, required_key_groups):
+    """
+    Checks if any required group of keys is present in config.
+    If none of the groups is present in full - raise Exception.
+    """
+    missing_key_groups = []
+
+    for required_keys in required_key_groups:
+        missing_keys = [key for key in required_keys if key not in config]
+        if missing_keys:
+            missing_key_groups.append(missing_keys)
+
+        # At least one group is complete no need to check further
+        return
+
+    if missing_key_groups:
+        missing_keys_msg = " or ".join(str(i) for i in missing_key_groups)
+        raise Exception(f"Config is missing required keys: {missing_keys_msg}")