wehub/open-metadata--openmetadata
1
0
Fork 0
Code Issues Pull Requests Actions 19 Deployments Models Agent Notes Packages Projects Releases Wiki Activity
Files
main
open-metadata--openmetadata/ingestion/tests/unit/test_source_connection.py
T
wehub-resource-sync bf2343b7e4
Integration Tests - MySQL + Elasticsearch / Detect Changes (push) Has been cancelled
Details
Integration Tests - MySQL + Elasticsearch / integration-tests-mysql-elasticsearch (push) Has been cancelled
Details
Integration Tests - PostgreSQL + Elasticsearch + Redis / Detect Changes (push) Has been cancelled
Details
Integration Tests - PostgreSQL + Elasticsearch + Redis / integration-tests-postgres-elasticsearch-redis (push) Has been cancelled
Details
Integration Tests - PostgreSQL + OpenSearch / Detect Changes (push) Has been cancelled
Details
Integration Tests - PostgreSQL + OpenSearch / integration-tests-postgres-opensearch (push) Has been cancelled
Details
Java Checkstyle / java-checkstyle (push) Has been cancelled
Details
Maven Collate Tests / maven-collate-ci (push) Has been cancelled
Details
OpenMetadata Service Unit Tests / openmetadata-service-unit-tests-status (push) Has been cancelled
Details
Publish Package to Maven Central Repository / publish-maven-packages (push) Has been cancelled
Details
OpenMetadata Service Unit Tests / Detect Changes (push) Has been cancelled
Details
OpenMetadata Service Unit Tests / openmetadata-service-unit-tests (push) Has been cancelled
Details
OpenMetadata Service Unit Tests / k8s_operator-unit-tests (push) Has been cancelled
Details
chore: import upstream snapshot with attribution
2026-07-13 13:35:45 +08:00

879 lines
35 KiB
Python
Raw Permalink Blame History

# Copyright 2025 Collate
# Licensed under the Collate Community License, Version 1.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
# https://github.com/open-metadata/OpenMetadata/blob/main/ingestion/LICENSE
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
from unittest import TestCase
from trino.auth import BasicAuthentication, JWTAuthentication, OAuth2Authentication
from metadata.generated.schema.entity.services.connections.database.clickhouseConnection import (
ClickhouseConnection,
ClickhouseScheme,
)
from metadata.generated.schema.entity.services.connections.database.common import (
noConfigAuthenticationTypes,
)
from metadata.generated.schema.entity.services.connections.database.common.basicAuth import (
BasicAuth,
)
from metadata.generated.schema.entity.services.connections.database.common.jwtAuth import (
JwtAuth,
)
from metadata.generated.schema.entity.services.connections.database.databricks.databricksOAuth import (
DatabricksOauth,
)
from metadata.generated.schema.entity.services.connections.database.databricks.personalAccessToken import (
PersonalAccessToken,
)
from metadata.generated.schema.entity.services.connections.database.databricksConnection import (
DatabricksConnection,
DatabricksScheme,
)
from metadata.generated.schema.entity.services.connections.database.db2Connection import (
Db2Connection,
Db2Scheme,
)
from metadata.generated.schema.entity.services.connections.database.mariaDBConnection import (
MariaDBConnection,
MariaDBScheme,
)
from metadata.generated.schema.entity.services.connections.database.mssqlConnection import (
MssqlConnection,
MssqlScheme,
)
from metadata.generated.schema.entity.services.connections.database.mysqlConnection import (
MysqlConnection,
MySQLScheme,
)
from metadata.generated.schema.entity.services.connections.database.oracleConnection import (
OracleConnection as OracleConnectionConfig,
)
from metadata.generated.schema.entity.services.connections.database.oracleConnection import (
OracleDatabaseSchema,
OracleScheme,
OracleServiceName,
OracleTNSConnection,
)
from metadata.generated.schema.entity.services.connections.database.postgresConnection import (
PostgresConnection,
PostgresScheme,
)
from metadata.generated.schema.entity.services.connections.database.redshiftConnection import (
RedshiftConnection,
RedshiftScheme,
)
from metadata.generated.schema.entity.services.connections.database.singleStoreConnection import (
SingleStoreConnection,
SingleStoreScheme,
)
from metadata.generated.schema.entity.services.connections.database.snowflakeConnection import (
SnowflakeConnection as SnowflakeConnectionConfig,
)
from metadata.generated.schema.entity.services.connections.database.snowflakeConnection import (
SnowflakeScheme,
)
from metadata.generated.schema.entity.services.connections.database.trinoConnection import (
TrinoConnection as TrinoConnectionConfig,
)
from metadata.generated.schema.entity.services.connections.database.trinoConnection import (
TrinoScheme,
)
from metadata.ingestion.connections.builders import (
get_connection_args_common,
get_connection_url_common,
)
from metadata.ingestion.source.database.oracle.connection import OracleConnection
from metadata.ingestion.source.database.snowflake.connection import SnowflakeConnection
from metadata.ingestion.source.database.trino.connection import TrinoConnection
# pylint: disable=import-outside-toplevel
class SourceConnectionTest(TestCase):
def test_databricks_url_without_db(self):
from metadata.ingestion.source.database.databricks.connection import (
get_connection_url,
)
expected_result = "databricks://1.1.1.1:443"
databricks_conn_obj = DatabricksConnection(
scheme=DatabricksScheme.databricks,
hostPort="1.1.1.1:443",
authType=PersonalAccessToken(token="KlivDTACWXKmZVfN1qIM"),
httpPath="/sql/1.0/warehouses/abcdedfg",
)
assert expected_result == get_connection_url(databricks_conn_obj)
def test_databricks_url_with_db(self):
from metadata.ingestion.source.database.databricks.connection import (
get_connection_url,
)
expected_result = "databricks://1.1.1.1:443?catalog=main"
databricks_conn_obj = DatabricksConnection(
scheme=DatabricksScheme.databricks,
hostPort="1.1.1.1:443",
authType=DatabricksOauth(
clientId="d40e2905-88ef-42ab-8898-fbefff2d071d",
clientSecret="secret-value",
),
httpPath="/sql/1.0/warehouses/abcdedfg",
catalog="main",
)
assert expected_result == get_connection_url(databricks_conn_obj)
def test_databricks_pipeline_url(self):
from metadata.generated.schema.entity.services.connections.pipeline.databricksPipelineConnection import (
DatabricksPipelineConnection,
)
from metadata.ingestion.source.pipeline.databrickspipeline.connection import (
get_connection_url,
)
conn_obj = DatabricksPipelineConnection(
hostPort="my-workspace.cloud.databricks.com:443",
token="dapi1234567890",
)
url = get_connection_url(conn_obj)
assert url == "databricks://token:dapi1234567890@my-workspace.cloud.databricks.com:443"
assert "databricks+connector" not in url
def test_databricks_url_with_special_chars_in_catalog(self):
from metadata.ingestion.source.database.databricks.connection import (
get_connection_url,
)
databricks_conn_obj = DatabricksConnection(
scheme=DatabricksScheme.databricks,
hostPort="1.1.1.1:443",
authType=PersonalAccessToken(token="KlivDTACWXKmZVfN1qIM"),
httpPath="/sql/1.0/warehouses/abcdedfg",
catalog="my catalog&name=val",
)
url = get_connection_url(databricks_conn_obj)
assert url == "databricks://1.1.1.1:443?catalog=my+catalog%26name%3Dval"
def test_unity_catalog_url_without_catalog(self):
from metadata.generated.schema.entity.services.connections.database.unityCatalogConnection import (
DatabricksScheme as UCDatabricksScheme,
)
from metadata.generated.schema.entity.services.connections.database.unityCatalogConnection import (
UnityCatalogConnection,
)
from metadata.ingestion.source.database.unitycatalog.connection import (
get_connection_url,
)
conn_obj = UnityCatalogConnection(
scheme=UCDatabricksScheme.databricks,
hostPort="my-workspace.cloud.databricks.com:443",
authType=PersonalAccessToken(token="dapi1234567890"),
httpPath="/sql/1.0/warehouses/abc",
)
url = get_connection_url(conn_obj)
assert url == "databricks://my-workspace.cloud.databricks.com:443"
def test_unity_catalog_url_with_catalog(self):
from metadata.generated.schema.entity.services.connections.database.unityCatalogConnection import (
DatabricksScheme as UCDatabricksScheme,
)
from metadata.generated.schema.entity.services.connections.database.unityCatalogConnection import (
UnityCatalogConnection,
)
from metadata.ingestion.source.database.unitycatalog.connection import (
get_connection_url,
)
conn_obj = UnityCatalogConnection(
scheme=UCDatabricksScheme.databricks,
hostPort="my-workspace.cloud.databricks.com:443",
authType=PersonalAccessToken(token="dapi1234567890"),
httpPath="/sql/1.0/warehouses/abc",
catalog="production",
)
url = get_connection_url(conn_obj)
assert url == "databricks://my-workspace.cloud.databricks.com:443?catalog=production"
def test_unity_catalog_url_with_special_chars_in_catalog(self):
from metadata.generated.schema.entity.services.connections.database.unityCatalogConnection import (
DatabricksScheme as UCDatabricksScheme,
)
from metadata.generated.schema.entity.services.connections.database.unityCatalogConnection import (
UnityCatalogConnection,
)
from metadata.ingestion.source.database.unitycatalog.connection import (
get_connection_url,
)
conn_obj = UnityCatalogConnection(
scheme=UCDatabricksScheme.databricks,
hostPort="my-workspace.cloud.databricks.com:443",
authType=PersonalAccessToken(token="dapi1234567890"),
httpPath="/sql/1.0/warehouses/abc",
catalog="my catalog&name=val",
)
url = get_connection_url(conn_obj)
assert url == "databricks://my-workspace.cloud.databricks.com:443?catalog=my+catalog%26name%3Dval"
def test_trino_url_without_params(self):
expected_url = "trino://username@localhost:443/catalog"
trino_conn_obj = TrinoConnectionConfig(
scheme=TrinoScheme.trino,
hostPort="localhost:443",
username="username",
authType=BasicAuth(password="pass"),
catalog="catalog",
)
trino_connection = TrinoConnection(trino_conn_obj)
assert expected_url == str(trino_connection.client.url)
# Passing @ in username and password
expected_url = "trino://username%40444@localhost:443/catalog"
trino_conn_obj = TrinoConnectionConfig(
scheme=TrinoScheme.trino,
hostPort="localhost:443",
username="username@444",
authType=BasicAuth(password="pass@111"),
catalog="catalog",
)
trino_connection = TrinoConnection(trino_conn_obj)
assert expected_url == str(trino_connection.client.url)
def test_trino_conn_arguments(self):
# connection arguments without connectionArguments and without proxies
expected_args = {
"auth": BasicAuthentication("user", None),
"http_scheme": "https",
}
trino_conn_obj = TrinoConnectionConfig(
username="user",
authType=BasicAuth(password=None),
hostPort="localhost:443",
catalog="tpcds",
connectionArguments=None,
scheme=TrinoScheme.trino,
)
trino_connection = TrinoConnection(trino_conn_obj)
assert expected_args == trino_connection.build_connection_args(trino_conn_obj).root
# connection arguments with connectionArguments and without proxies
expected_args = {
"user": "user-to-be-impersonated",
"auth": BasicAuthentication("user", None),
"http_scheme": "https",
}
trino_conn_obj = TrinoConnectionConfig(
username="user",
authType=BasicAuth(password=None),
hostPort="localhost:443",
catalog="tpcds",
connectionArguments={"user": "user-to-be-impersonated"},
scheme=TrinoScheme.trino,
)
trino_connection = TrinoConnection(trino_conn_obj)
assert expected_args == trino_connection.build_connection_args(trino_conn_obj).root
# connection arguments without connectionArguments and with proxies
expected_args = {
"auth": BasicAuthentication("user", None),
"http_scheme": "https",
}
trino_conn_obj = TrinoConnectionConfig(
username="user",
authType=BasicAuth(password=None),
hostPort="localhost:443",
catalog="tpcds",
connectionArguments=None,
proxies={"http": "foo.bar:3128", "http://host.name": "foo.bar:4012"},
scheme=TrinoScheme.trino,
)
trino_connection = TrinoConnection(trino_conn_obj)
conn_args = trino_connection.build_connection_args(trino_conn_obj).root
assert "http_session" in conn_args
conn_args.pop("http_session")
assert expected_args == conn_args
# connection arguments with connectionArguments and with proxies
expected_args = {
"user": "user-to-be-impersonated",
"auth": BasicAuthentication("user", None),
"http_scheme": "https",
}
trino_conn_obj = TrinoConnectionConfig(
username="user",
authType=BasicAuth(password=None),
hostPort="localhost:443",
catalog="tpcds",
connectionArguments={"user": "user-to-be-impersonated"},
proxies={"http": "foo.bar:3128", "http://host.name": "foo.bar:4012"},
scheme=TrinoScheme.trino,
)
trino_connection = TrinoConnection(trino_conn_obj)
conn_args = trino_connection.build_connection_args(trino_conn_obj).root
assert "http_session" in conn_args
conn_args.pop("http_session")
assert expected_args == conn_args
def test_trino_url_with_params(self):
expected_url = "trino://username@localhost:443/catalog?param=value"
trino_conn_obj = TrinoConnectionConfig(
scheme=TrinoScheme.trino,
hostPort="localhost:443",
username="username",
authType=BasicAuth(password="pass"),
catalog="catalog",
connectionOptions={"param": "value"},
)
trino_connection = TrinoConnection(trino_conn_obj)
assert expected_url == str(trino_connection.client.url)
def test_trino_url_with_jwt_auth(self):
expected_url = "trino://username@localhost:443/catalog"
expected_args = {
"auth": JWTAuthentication("jwt_token_value"),
"http_scheme": "https",
}
trino_conn_obj = TrinoConnectionConfig(
scheme=TrinoScheme.trino,
hostPort="localhost:443",
username="username",
authType=JwtAuth(jwt="jwt_token_value"),
catalog="catalog",
)
trino_connection = TrinoConnection(trino_conn_obj)
assert expected_url == str(trino_connection.client.url)
assert expected_args == trino_connection.build_connection_args(trino_conn_obj).root
def test_trino_with_proxies(self):
test_proxies = {"http": "http_proxy", "https": "https_proxy"}
trino_conn_obj = TrinoConnectionConfig(
scheme=TrinoScheme.trino,
hostPort="localhost:443",
username="username",
authType=BasicAuth(password="pass"),
catalog="catalog",
proxies=test_proxies,
)
trino_connection = TrinoConnection(trino_conn_obj)
assert test_proxies == trino_connection.build_connection_args(trino_conn_obj).root.get("http_session").proxies
def test_trino_without_catalog(self):
# Test trino url without catalog
expected_url = "trino://username@localhost:443"
trino_conn_obj = TrinoConnectionConfig(
scheme=TrinoScheme.trino,
hostPort="localhost:443",
username="username",
authType=BasicAuth(password="pass"),
)
trino_connection = TrinoConnection(trino_conn_obj)
assert expected_url == str(trino_connection.client.url)
def test_trino_without_catalog(self): # noqa: F811
# Test trino url without catalog
expected_url = "trino://username@localhost:443"
trino_conn_obj = TrinoConnectionConfig(
scheme=TrinoScheme.trino,
hostPort="localhost:443",
username="username",
authType=BasicAuth(password="pass"),
)
trino_connection = TrinoConnection(trino_conn_obj)
assert expected_url == str(trino_connection.client.url)
def test_trino_with_oauth2(self):
# Test trino url without catalog
expected_url = "trino://username@localhost:443" # noqa: F841
trino_conn_obj = TrinoConnectionConfig(
scheme=TrinoScheme.trino,
hostPort="localhost:443",
username="username",
authType=noConfigAuthenticationTypes.NoConfigAuthenticationTypes.OAuth2,
)
trino_connection = TrinoConnection(trino_conn_obj)
assert trino_connection.build_connection_args(trino_conn_obj).root.get("auth") == OAuth2Authentication()
def test_mysql_url(self):
# connection arguments without db
expected_url = "mysql+pymysql://openmetadata_user:@localhost:3306"
mysql_conn_obj = MysqlConnection(
username="openmetadata_user",
hostPort="localhost:3306",
scheme=MySQLScheme.mysql_pymysql,
)
assert expected_url == get_connection_url_common(mysql_conn_obj)
# connection arguments with db
expected_url = "mysql+pymysql://openmetadata_user:@localhost:3306"
mysql_conn_obj = MysqlConnection(
username="openmetadata_user",
hostPort="localhost:3306",
scheme=MySQLScheme.mysql_pymysql,
)
assert expected_url == get_connection_url_common(mysql_conn_obj)
def test_clickhouse_url(self):
# connection arguments without db
expected_url = "clickhouse+http://username:@localhost:8123"
clickhouse_conn_obj = ClickhouseConnection(
username="username",
hostPort="localhost:8123",
scheme=ClickhouseScheme.clickhouse_http,
databaseSchema=None,
)
assert expected_url == get_connection_url_common(clickhouse_conn_obj)
# connection arguments with db
expected_url = "clickhouse+http://username:@localhost:8123/default"
clickhouse_conn_obj = ClickhouseConnection(
username="username",
hostPort="localhost:8123",
scheme=ClickhouseScheme.clickhouse_http,
databaseSchema="default",
)
assert expected_url == get_connection_url_common(clickhouse_conn_obj)
expected_url = "clickhouse+http://username:@localhost:8123/default?protocol=https"
clickhouse_conn_obj = ClickhouseConnection(
username="username",
hostPort="localhost:8123",
scheme=ClickhouseScheme.clickhouse_http,
connectionOptions=dict(protocol="https"), # noqa: C408
databaseSchema="default",
)
assert expected_url == get_connection_url_common(clickhouse_conn_obj)
def test_mariadb_url(self):
# connection arguments without db
expected_url = "mysql+pymysql://openmetadata_user:@localhost:3306"
mariadb_conn_obj = MariaDBConnection(
username="openmetadata_user",
hostPort="localhost:3306",
scheme=MariaDBScheme.mysql_pymysql,
)
assert expected_url == get_connection_url_common(mariadb_conn_obj)
# connection arguments with db
expected_url = "mysql+pymysql://openmetadata_user:@localhost:3306"
mariadb_conn_obj = MariaDBConnection(
username="openmetadata_user",
hostPort="localhost:3306",
scheme=MariaDBScheme.mysql_pymysql,
)
assert expected_url == get_connection_url_common(mariadb_conn_obj)
def test_postgres_url(self):
# connection arguments with db
expected_url = "postgresql+psycopg2://openmetadata_user:@localhost:5432/default"
postgres_conn_obj = PostgresConnection(
username="openmetadata_user",
hostPort="localhost:5432",
database="default",
scheme=PostgresScheme.postgresql_psycopg2,
)
assert expected_url == get_connection_url_common(postgres_conn_obj)
def test_redshift_url(self):
# connection arguments witho db
expected_url = (
"redshift+psycopg2://username:strong_password@cluster.name.region.redshift.amazonaws.com:5439/dev"
)
redshift_conn_obj = RedshiftConnection(
username="username",
authType=BasicAuth(password="strong_password"),
hostPort="cluster.name.region.redshift.amazonaws.com:5439",
scheme=RedshiftScheme.redshift_psycopg2,
database="dev",
)
assert expected_url == get_connection_url_common(redshift_conn_obj)
def test_singleStore_url(self): # noqa: N802
# connection arguments without db
expected_url = "mysql+pymysql://openmetadata_user:@localhost:5432"
singleStore_conn_obj = SingleStoreConnection( # noqa: N806
username="openmetadata_user",
hostPort="localhost:5432",
scheme=SingleStoreScheme.mysql_pymysql,
)
assert expected_url == get_connection_url_common(singleStore_conn_obj)
# connection arguments with db
expected_url = "mysql+pymysql://openmetadata_user:@localhost:5432"
singleStore_conn_obj = SingleStoreConnection( # noqa: N806
username="openmetadata_user",
hostPort="localhost:5432",
scheme=SingleStoreScheme.mysql_pymysql,
)
assert expected_url == get_connection_url_common(singleStore_conn_obj)
def test_db2_url(self):
# connection arguments with db
expected_url = "db2+ibm_db://openmetadata_user:@localhost:50000/testdb"
db2_conn_obj = Db2Connection(
scheme=Db2Scheme.db2_ibm_db,
username="openmetadata_user",
hostPort="localhost:50000",
database="testdb",
)
assert expected_url == get_connection_url_common(db2_conn_obj)
def test_snowflake_url(self):
# Passing @ in username and password
expected_url = "snowflake://coding%40444:Abhi%40123@ue18849.us-east-2.aws?account=ue18849.us-east-2.aws&warehouse=COMPUTE_WH"
snowflake_conn_obj = SnowflakeConnectionConfig(
scheme=SnowflakeScheme.snowflake,
username="coding@444",
password="Abhi@123",
warehouse="COMPUTE_WH",
account="ue18849.us-east-2.aws",
)
assert expected_url == SnowflakeConnection.get_connection_url(snowflake_conn_obj)
# connection arguments with db
expected_url = (
"snowflake://coding:Abhi@ue18849.us-east-2.aws/testdb?account=ue18849.us-east-2.aws&warehouse=COMPUTE_WH"
)
snowflake_conn_obj = SnowflakeConnectionConfig(
scheme=SnowflakeScheme.snowflake,
username="coding",
password="Abhi",
database="testdb",
warehouse="COMPUTE_WH",
account="ue18849.us-east-2.aws",
)
assert expected_url == SnowflakeConnection.get_connection_url(snowflake_conn_obj)
def test_mysql_conn_arguments(self):
# connection arguments without connectionArguments
expected_args = {}
mysql_conn_obj = MysqlConnection(
username="user",
authType=BasicAuth(
password=None,
),
hostPort="localhost:443",
connectionArguments=None,
scheme=MySQLScheme.mysql_pymysql,
)
assert expected_args == get_connection_args_common(mysql_conn_obj)
# connection arguments with connectionArguments
expected_args = {"user": "user-to-be-impersonated"}
mysql_conn_obj = MysqlConnection(
username="user",
authType=BasicAuth(
password=None,
),
hostPort="localhost:443",
connectionArguments={"user": "user-to-be-impersonated"},
scheme=MySQLScheme.mysql_pymysql,
)
assert expected_args == get_connection_args_common(mysql_conn_obj)
def test_clickhouse_conn_arguments(self):
# connection arguments without connectionArguments
expected_args = {}
clickhouse_conn_obj = ClickhouseConnection(
username="user",
password=None,
hostPort="localhost:443",
databaseSchema=None,
connectionArguments=None,
scheme=ClickhouseScheme.clickhouse_http,
)
assert expected_args == get_connection_args_common(clickhouse_conn_obj)
# connection arguments with connectionArguments
expected_args = {"user": "user-to-be-impersonated"}
clickhouse_conn_obj = ClickhouseConnection(
username="user",
password=None,
hostPort="localhost:443",
databaseSchema="tiny",
connectionArguments={"user": "user-to-be-impersonated"},
scheme=ClickhouseScheme.clickhouse_http,
)
assert expected_args == get_connection_args_common(clickhouse_conn_obj)
def test_mariadb_conn_arguments(self):
# connection arguments without connectionArguments
expected_args = {}
mariadb_conn_obj = MariaDBConnection(
username="user",
password=None,
hostPort="localhost:443",
connectionArguments=None,
scheme=MariaDBScheme.mysql_pymysql,
)
assert expected_args == get_connection_args_common(mariadb_conn_obj)
# connection arguments with connectionArguments
expected_args = {"user": "user-to-be-impersonated"}
mariadb_conn_obj = MariaDBConnection(
username="user",
password=None,
hostPort="localhost:443",
connectionArguments={"user": "user-to-be-impersonated"},
scheme=MariaDBScheme.mysql_pymysql,
)
assert expected_args == get_connection_args_common(mariadb_conn_obj)
def test_postgres_conn_arguments(self):
# connection arguments without connectionArguments
expected_args = {}
postgres_conn_obj = PostgresConnection(
username="user",
authType=BasicAuth(
password=None,
),
database="postgres",
hostPort="localhost:443",
connectionArguments=None,
scheme=PostgresScheme.postgresql_psycopg2,
)
assert expected_args == get_connection_args_common(postgres_conn_obj)
# connection arguments with connectionArguments
expected_args = {"user": "user-to-be-impersonated"}
postgres_conn_obj = PostgresConnection(
authType=BasicAuth(
password=None,
),
username="user",
hostPort="localhost:443",
database="tiny",
connectionArguments={"user": "user-to-be-impersonated"},
scheme=PostgresScheme.postgresql_psycopg2,
)
assert expected_args == get_connection_args_common(postgres_conn_obj)
def test_redshift_conn_arguments(self):
# connection arguments without connectionArguments
expected_args = {}
redshift_conn_obj = RedshiftConnection(
username="user",
authType=BasicAuth(password=None),
hostPort="localhost:443",
database="tiny",
connectionArguments=None,
scheme=RedshiftScheme.redshift_psycopg2,
)
assert expected_args == get_connection_args_common(redshift_conn_obj)
# connection arguments with connectionArguments
expected_args = {"user": "user-to-be-impersonated"}
redshift_conn_obj = RedshiftConnection(
username="user",
authType=BasicAuth(password=None),
hostPort="localhost:443",
database="tiny",
connectionArguments={"user": "user-to-be-impersonated"},
scheme=RedshiftScheme.redshift_psycopg2,
)
assert expected_args == get_connection_args_common(redshift_conn_obj)
def test_singleStore_conn_arguments(self): # noqa: N802
# connection arguments without connectionArguments
expected_args = {}
singleStore_conn_obj = SingleStoreConnection( # noqa: N806
username="user",
password=None,
hostPort="localhost:443",
connectionArguments=None,
scheme=SingleStoreScheme.mysql_pymysql,
)
assert expected_args == get_connection_args_common(singleStore_conn_obj)
# connection arguments with connectionArguments
expected_args = {"user": "user-to-be-impersonated"}
singleStore_conn_obj = SingleStoreConnection( # noqa: N806
username="user",
password=None,
hostPort="localhost:443",
connectionArguments={"user": "user-to-be-impersonated"},
scheme=SingleStoreScheme.mysql_pymysql,
)
assert expected_args == get_connection_args_common(singleStore_conn_obj)
def test_db2_conn_arguments(self):
# connection arguments without connectionArguments
expected_args = {}
db2_conn_obj = Db2Connection(
username="user",
password=None,
hostPort="localhost:443",
connectionArguments=None,
scheme=Db2Scheme.db2_ibm_db,
database="testdb",
)
assert expected_args == get_connection_args_common(db2_conn_obj)
# connection arguments with connectionArguments
expected_args = {"user": "user-to-be-impersonated"}
db2_conn_obj = Db2Connection(
username="user",
password=None,
hostPort="localhost:443",
connectionArguments={"user": "user-to-be-impersonated"},
scheme=Db2Scheme.db2_ibm_db,
database="testdb",
)
assert expected_args == get_connection_args_common(db2_conn_obj)
def test_snowflake_conn_arguments(self):
# connection arguments without connectionArguments
expected_args = {}
snowflake_conn_obj = SnowflakeConnectionConfig(
username="user",
password="test-pwd",
database="tiny",
warehouse="COMPUTE_WH",
scheme=SnowflakeScheme.snowflake,
account="account.region_name.cloud_service",
)
assert expected_args == get_connection_args_common(snowflake_conn_obj)
# connection arguments with connectionArguments
expected_args = {"user": "user-to-be-impersonated"}
snowflake_conn_obj = SnowflakeConnectionConfig(
username="user",
password="test-pwd",
database="tiny",
warehouse="COMPUTE_WH",
connectionArguments={"user": "user-to-be-impersonated"},
scheme=SnowflakeScheme.snowflake,
account="account.region_name.cloud_service",
)
assert expected_args == get_connection_args_common(snowflake_conn_obj)
def test_mssql_url(self):
from metadata.ingestion.source.database.mssql.connection import (
get_connection_url,
)
# connection arguments without db
expected_url = "mssql+pytds://sa:password@localhost:1433"
mssql_conn_obj = MssqlConnection(
username="sa",
password="password",
hostPort="localhost:1433",
scheme=MssqlScheme.mssql_pytds,
database=None,
)
assert expected_url == get_connection_url(mssql_conn_obj)
def test_mssql_url(self): # noqa: F811
from metadata.ingestion.source.database.mssql.connection import (
get_connection_url,
)
# Passing @ in username and password
expected_url = "mssql+pytds://sa%40123:password%40444@localhost:1433/master"
mssql_conn_obj = MssqlConnection(
username="sa@123",
password="password@444",
hostPort="localhost:1433",
scheme=MssqlScheme.mssql_pytds,
database="master",
)
assert expected_url == get_connection_url(mssql_conn_obj)
# connection arguments witho db
expected_url = "mssql+pytds://sa:password@localhost:1433/catalog_test"
mssql_conn_obj = MssqlConnection(
username="sa",
password="password",
hostPort="localhost:1433",
scheme=MssqlScheme.mssql_pytds,
database="catalog_test",
)
assert expected_url == get_connection_url(mssql_conn_obj)
def test_oracle_url(self):
# oracle with db
expected_url = "oracle+cx_oracle://admin:password@localhost:1541/testdb"
oracle_conn_obj = OracleConnectionConfig(
username="admin",
password="password",
hostPort="localhost:1541",
scheme=OracleScheme.oracle_cx_oracle,
oracleConnectionType=OracleDatabaseSchema(databaseSchema="testdb"),
)
assert expected_url == OracleConnection.get_connection_url(oracle_conn_obj)
# oracle with service name
expected_url = "oracle+cx_oracle://admin:password@localhost:1541/?service_name=testdb"
oracle_conn_obj = OracleConnectionConfig(
username="admin",
password="password",
hostPort="localhost:1541",
scheme=OracleScheme.oracle_cx_oracle,
oracleConnectionType=OracleServiceName(oracleServiceName="testdb"),
)
assert expected_url == OracleConnection.get_connection_url(oracle_conn_obj)
# oracle with db & connection options
expected_url = [
"oracle+cx_oracle://admin:password@localhost:1541/testdb?test_key_2=test_value_2&test_key_1=test_value_1",
"oracle+cx_oracle://admin:password@localhost:1541/testdb?test_key_1=test_value_1&test_key_2=test_value_2",
]
oracle_conn_obj = OracleConnectionConfig(
username="admin",
password="password",
hostPort="localhost:1541",
scheme=OracleScheme.oracle_cx_oracle,
oracleConnectionType=OracleDatabaseSchema(databaseSchema="testdb"),
connectionOptions=dict(test_key_1="test_value_1", test_key_2="test_value_2"), # noqa: C408
)
assert OracleConnection.get_connection_url(oracle_conn_obj) in expected_url
# oracle with service name & connection options
expected_url = [
"oracle+cx_oracle://admin:password@localhost:1541/?service_name=testdb&test_key_2=test_value_2&test_key_1=test_value_1",
"oracle+cx_oracle://admin:password@localhost:1541/?service_name=testdb&test_key_1=test_value_1&test_key_2=test_value_2",
]
oracle_conn_obj = OracleConnectionConfig(
username="admin",
password="password",
hostPort="localhost:1541",
scheme=OracleScheme.oracle_cx_oracle,
oracleConnectionType=OracleServiceName(oracleServiceName="testdb"),
connectionOptions=dict(test_key_1="test_value_1", test_key_2="test_value_2"), # noqa: C408
)
assert OracleConnection.get_connection_url(oracle_conn_obj) in expected_url
tns_connection = (
"(DESCRIPTION=(ADDRESS_LIST=(ADDRESS=(PROTOCOL=TCP)"
"(HOST=myhost)(PORT=1530)))(CONNECT_DATA=(SID=MYSERVICENAME)))"
)
expected_url = f"oracle+cx_oracle://admin:password@{tns_connection}"
oracle_conn_obj = OracleConnectionConfig(
username="admin",
password="password",
hostPort="localhost:1541", # We will ignore it here
oracleConnectionType=OracleTNSConnection(oracleTNSConnection=tns_connection),
)
assert OracleConnection.get_connection_url(oracle_conn_obj) == expected_url
Reference in New Issue View Git Blame Copy Permalink