From 7c3ad3d2931cf6e19cdd94c65e44f684de85f874 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Marcin=20Szyma=C5=84ski?= <ms32035@gmail.com>
Date: Fri, 8 Apr 2022 04:11:52 +0100
Subject: [PATCH] feat(ingest): enable connection string for all sqlalchemy
 datasources (#4508)

* feat(ingest): enable connection string for all sqlalchemy datasources

* Update sql_common.py

* fix types

* update docs

* rename variable to sqlalchemy_uri

Co-authored-by: Harshal Sheth <hsheth2@gmail.com>
---
 metadata-ingestion/source_docs/clickhouse.md       |  1 +
 metadata-ingestion/source_docs/druid.md            |  1 +
 metadata-ingestion/source_docs/mysql.md            |  1 +
 metadata-ingestion/source_docs/oracle.md           |  1 +
 metadata-ingestion/source_docs/postgres.md         |  1 +
 metadata-ingestion/source_docs/redshift.md         |  1 +
 metadata-ingestion/source_docs/snowflake.md        |  1 +
 metadata-ingestion/source_docs/trino.md            |  1 +
 .../src/datahub/ingestion/source/sql/sql_common.py | 14 +++++++++-----
 9 files changed, 17 insertions(+), 5 deletions(-)

diff --git a/metadata-ingestion/source_docs/clickhouse.md b/metadata-ingestion/source_docs/clickhouse.md
index d3054bf051..6c776a609f 100644
--- a/metadata-ingestion/source_docs/clickhouse.md
+++ b/metadata-ingestion/source_docs/clickhouse.md
@@ -87,6 +87,7 @@ Note that a `.` is used to denote nested fields in the YAML recipe.
 | `password`                  |          |                                                                            | ClickHouse password.                                                                                                                                                                    |
 | `host_port`                 | ✅        |                                                                            | ClickHouse host URL.                                                                                                                                                                    |
 | `database`                  |          |                                                                            | ClickHouse database to connect.                                                                                                                                                         |
+| `sqlalchemy_uri`            |          |          | URI of database to connect to. See https://docs.sqlalchemy.org/en/14/core/engines.html#database-urls. Takes precedence over other connection parameters. |
 | `env`                       |          | `"PROD"`                                                                   | Environment to use in namespace when constructing URNs.                                                                                                                                 |
 | `platform_instance`         |          | None                                                                       | The Platform instance to use while constructing URNs.                                                                                                                                   |
 | `options.<option>`          |          |                                                                            | Any options specified here will be passed to SQLAlchemy's `create_engine` as kwargs.<br />See https://docs.sqlalchemy.org/en/14/core/engines.html#sqlalchemy.create_engine for details. |
diff --git a/metadata-ingestion/source_docs/druid.md b/metadata-ingestion/source_docs/druid.md
index bafa26d8a6..22a744e712 100644
--- a/metadata-ingestion/source_docs/druid.md
+++ b/metadata-ingestion/source_docs/druid.md
@@ -54,6 +54,7 @@ As a SQL-based service, the Athena integration is also supported by our SQL prof
 | `password`                  |          |                         | Database password.                                                                                                                                                                      |
 | `host_port`                 | ✅       |                         | Host URL and port to connect to.                                                                                                                                                        |
 | `database`                  |          |                         | Database to ingest.                                                                                                                                                                     |
+| `sqlalchemy_uri`            |          |          | URI of database to connect to. See https://docs.sqlalchemy.org/en/14/core/engines.html#database-urls. Takes precedence over other connection parameters. |
 | `database_alias`            |          |                         | Alias to apply to database when ingesting.                                                                                                                                              |
 | `env`                       |          | `"PROD"`                | Environment to use in namespace when constructing URNs.                                                                                                                                 |
 | `platform_instance`         |          | None             | The Platform instance to use while constructing URNs.         |
diff --git a/metadata-ingestion/source_docs/mysql.md b/metadata-ingestion/source_docs/mysql.md
index beb71d795d..89854fb590 100644
--- a/metadata-ingestion/source_docs/mysql.md
+++ b/metadata-ingestion/source_docs/mysql.md
@@ -57,6 +57,7 @@ Note that a `.` is used to denote nested fields in the YAML recipe.
 | `password`                     |          |                    | MySQL password.                                                                                                                                                                         |
 | `host_port`                    |          | `"localhost:3306"` | MySQL host URL.                                                                                                                                                                         |
 | `database`                     |          |                    | MySQL database.                                                                                                                                                                         |
+| `sqlalchemy_uri`               |          |                    | URI of database to connect to. See https://docs.sqlalchemy.org/en/14/core/engines.html#database-urls. Takes precedence over other connection parameters. |
 | `database_alias`               |          |                    | Alias to apply to database when ingesting. A table with urn `urn:li:dataset:(urn:li:dataPlatform:mysql,<database>.<table>,PROD)` will turn into `urn:li:dataset:(urn:li:dataPlatform:mysql,<database_alias>.<table>,PROD)`.                                                                                                                                             |
 | `env`                          |          | `"PROD"`           | Environment to use in namespace when constructing URNs.                                                                                                                                 |
 | `platform_instance`            |          | None               | The Platform instance to use while constructing URNs.                                                                                                                                   |
diff --git a/metadata-ingestion/source_docs/oracle.md b/metadata-ingestion/source_docs/oracle.md
index 84423c000c..3c911e8f1e 100644
--- a/metadata-ingestion/source_docs/oracle.md
+++ b/metadata-ingestion/source_docs/oracle.md
@@ -61,6 +61,7 @@ As a SQL-based service, the Athena integration is also supported by our SQL prof
 | `host_port`                    |                                |          | Oracle host URL.                                                                                                                                                                                                                                                                  |
 | `database`                     | If `service_name` is not set   |          | If using, omit `service_name`.                                                                                                                                                                                                                                                    |
 | `service_name`                 | If `database_alias` is not set |          | Oracle service name. If using, omit `database`.                                                                                                                                                                                                                                   |
+| `sqlalchemy_uri`               |                                |          | URI of database to connect to. See https://docs.sqlalchemy.org/en/14/core/engines.html#database-urls. Takes precedence over other connection parameters. |
 | `database_alias`               |                                |          | Alias to apply to database when ingesting.                                                                                                                                                                                                                                        |
 | `env`                          |                                | `"PROD"` | Environment to use in namespace when constructing URNs.                                                                                                                                                                                                                           |
 | `options.<option>`             |                                |          | Any options specified here will be passed to SQLAlchemy's `create_engine` as kwargs.<br />See https://docs.sqlalchemy.org/en/14/core/engines.html#sqlalchemy.create_engine for details.                                                                                           |
diff --git a/metadata-ingestion/source_docs/postgres.md b/metadata-ingestion/source_docs/postgres.md
index 9db5d15bb6..6f7b3cece5 100644
--- a/metadata-ingestion/source_docs/postgres.md
+++ b/metadata-ingestion/source_docs/postgres.md
@@ -59,6 +59,7 @@ As a SQL-based service, the Athena integration is also supported by our SQL prof
 | `password`                     |          |          | PostgreSQL password.                                                                                                                                                                    |
 | `host_port`                    | ✅       |          | PostgreSQL host URL.                                                                                                                                                                    |
 | `database`                     |          |          | PostgreSQL database.                                                                                                                                                                    |
+| `sqlalchemy_uri`               |          |          | URI of database to connect to. See https://docs.sqlalchemy.org/en/14/core/engines.html#database-urls. Takes precedence over other connection parameters. |
 | `database_alias`               |          |          | Alias to apply to database when ingesting.                                                                                                                                              |
 | `env`                          |          | `"PROD"` | Environment to use in namespace when constructing URNs.                                                                                                                                 |
 | `platform_instance`            |          | None     | The Platform instance to use while constructing URNs.                                                                                                                                   |
diff --git a/metadata-ingestion/source_docs/redshift.md b/metadata-ingestion/source_docs/redshift.md
index 9b0a0717d1..a334bc1508 100644
--- a/metadata-ingestion/source_docs/redshift.md
+++ b/metadata-ingestion/source_docs/redshift.md
@@ -111,6 +111,7 @@ Note that a `.` is used to denote nested fields in the YAML recipe.
 | `password`                     |          |                    | Redshift password.                                                                                                                                                                      |
 | `host_port`                    | ✅        |                    | Redshift host URL.                                                                                                                                                                      |
 | `database`                     |          |                    | Redshift database.                                                                                                                                                                      |
+| `sqlalchemy_uri`               |          |                    | URI of database to connect to. See https://docs.sqlalchemy.org/en/14/core/engines.html#database-urls. Takes precedence over other connection parameters. |
 | `database_alias`               |          |                    | Alias to apply to database when ingesting.                                                                                                                                              |
 | `env`                          |          | `"PROD"`           | Environment to use in namespace when constructing URNs.                                                                                                                                 |
 | `platform_instance`            |          | None               | The Platform instance to use while constructing URNs.                                                                                                                                   |
diff --git a/metadata-ingestion/source_docs/snowflake.md b/metadata-ingestion/source_docs/snowflake.md
index aa1228aee1..a5d6bb237e 100644
--- a/metadata-ingestion/source_docs/snowflake.md
+++ b/metadata-ingestion/source_docs/snowflake.md
@@ -168,6 +168,7 @@ Note that a `.` is used to denote nested fields in the YAML recipe.
 | `host_port`                    | ✅        |                                                                            | Snowflake host URL.                                                                                                                                                                     |
 | `warehouse`                    |          |                                                                            | Snowflake warehouse.                                                                                                                                                                    |
 | `role`                         |          |                                                                            | Snowflake role.                                                                                                                                                                         |
+| `sqlalchemy_uri`               |          |                                                                            | URI of database to connect to. See https://docs.sqlalchemy.org/en/14/core/engines.html#database-urls. Takes precedence over other connection parameters. |
 | `env`                          |          | `"PROD"`                                                                   | Environment to use in namespace when constructing URNs.                                                                                                                                 |
 | `platform_instance`            |          | None                                                                       | The Platform instance to use while constructing URNs.                                                                                                                                   |
 | `options.<option>`             |          |                                                                            | Any options specified here will be passed to SQLAlchemy's `create_engine` as kwargs.<br />See https://docs.sqlalchemy.org/en/14/core/engines.html#sqlalchemy.create_engine for details. |
diff --git a/metadata-ingestion/source_docs/trino.md b/metadata-ingestion/source_docs/trino.md
index 6d779e1baa..f716b4c4ef 100644
--- a/metadata-ingestion/source_docs/trino.md
+++ b/metadata-ingestion/source_docs/trino.md
@@ -53,6 +53,7 @@ As a SQL-based service, the Trino integration is also supported by our SQL profi
 | `password`                     |          |          | Trino password.                                                                                                                                                                         |
 | `host_port`                    | ✅        |          | Trino host URL.                                                                                                                                                                         |
 | `database`                     | ✅        |          | Trino database (catalog).                                                                                                                                                               |
+| `sqlalchemy_uri`               |          |          | URI of database to connect to. See https://docs.sqlalchemy.org/en/14/core/engines.html#database-urls. Takes precedence over other connection parameters.
 | `database_alias`               |          |          | Alias to apply to database when ingesting.                                                                                                                                              |
 | `env`                          |          | `"PROD"` | Environment to use in namespace when constructing URNs.                                                                                                                                 |
 | `options.<option>`             |          |          | Any options specified here will be passed to SQLAlchemy's `create_engine` as kwargs.<br />See https://docs.sqlalchemy.org/en/14/core/engines.html#sqlalchemy.create_engine for details. |
diff --git a/metadata-ingestion/src/datahub/ingestion/source/sql/sql_common.py b/metadata-ingestion/src/datahub/ingestion/source/sql/sql_common.py
index f28dc4ee64..42ca8800de 100644
--- a/metadata-ingestion/src/datahub/ingestion/source/sql/sql_common.py
+++ b/metadata-ingestion/src/datahub/ingestion/source/sql/sql_common.py
@@ -249,17 +249,21 @@ class SQLAlchemyConfig(StatefulIngestionConfigBase):
 class BasicSQLAlchemyConfig(SQLAlchemyConfig):
     username: Optional[str] = None
     password: Optional[pydantic.SecretStr] = None
-    host_port: str
+    host_port: Optional[str] = None
     database: Optional[str] = None
     database_alias: Optional[str] = None
-    scheme: str
+    scheme: Optional[str] = None
+    sqlalchemy_uri: Optional[str] = None
 
     def get_sql_alchemy_url(self, uri_opts=None):
-        return make_sqlalchemy_uri(
-            self.scheme,
+        if not ((self.host_port and self.scheme) or self.sqlalchemy_uri):
+            raise ValueError("host_port and schema or connect_uri required.")
+
+        return self.sqlalchemy_uri or make_sqlalchemy_uri(
+            self.scheme,  # type: ignore
             self.username,
             self.password.get_secret_value() if self.password else None,
-            self.host_port,
+            self.host_port,  # type: ignore
             self.database,
             uri_opts=uri_opts,
         )