mirror of
https://github.com/open-metadata/OpenMetadata.git
synced 2025-08-07 08:38:23 +00:00
429 lines
14 KiB
Python
429 lines
14 KiB
Python
# Copyright 2021 Collate
|
|
# Licensed under the Apache License, Version 2.0 (the "License");
|
|
# you may not use this file except in compliance with the License.
|
|
# You may obtain a copy of the License at
|
|
# http://www.apache.org/licenses/LICENSE-2.0
|
|
# 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.
|
|
|
|
"""
|
|
Hosts the singledispatch to build source URLs
|
|
"""
|
|
import os
|
|
from functools import singledispatch
|
|
from urllib.parse import quote_plus
|
|
|
|
from pydantic import SecretStr
|
|
from requests import Session
|
|
|
|
from metadata.generated.schema.entity.services.connections.database.athenaConnection import (
|
|
AthenaConnection,
|
|
)
|
|
from metadata.generated.schema.entity.services.connections.database.azureSQLConnection import (
|
|
AzureSQLConnection,
|
|
)
|
|
from metadata.generated.schema.entity.services.connections.database.bigQueryConnection import (
|
|
BigQueryConnection,
|
|
)
|
|
from metadata.generated.schema.entity.services.connections.database.clickhouseConnection import (
|
|
ClickhouseConnection,
|
|
)
|
|
from metadata.generated.schema.entity.services.connections.database.databricksConnection import (
|
|
DatabricksConnection,
|
|
)
|
|
from metadata.generated.schema.entity.services.connections.database.db2Connection import (
|
|
Db2Connection,
|
|
)
|
|
from metadata.generated.schema.entity.services.connections.database.druidConnection import (
|
|
DruidConnection,
|
|
)
|
|
from metadata.generated.schema.entity.services.connections.database.hiveConnection import (
|
|
HiveConnection,
|
|
)
|
|
from metadata.generated.schema.entity.services.connections.database.mariaDBConnection import (
|
|
MariaDBConnection,
|
|
)
|
|
from metadata.generated.schema.entity.services.connections.database.mssqlConnection import (
|
|
MssqlConnection,
|
|
)
|
|
from metadata.generated.schema.entity.services.connections.database.mysqlConnection import (
|
|
MysqlConnection,
|
|
)
|
|
from metadata.generated.schema.entity.services.connections.database.oracleConnection import (
|
|
OracleConnection,
|
|
OracleDatabaseSchema,
|
|
OracleServiceName,
|
|
)
|
|
from metadata.generated.schema.entity.services.connections.database.pinotDBConnection import (
|
|
PinotDBConnection,
|
|
)
|
|
from metadata.generated.schema.entity.services.connections.database.postgresConnection import (
|
|
PostgresConnection,
|
|
)
|
|
from metadata.generated.schema.entity.services.connections.database.prestoConnection import (
|
|
PrestoConnection,
|
|
)
|
|
from metadata.generated.schema.entity.services.connections.database.redshiftConnection import (
|
|
RedshiftConnection,
|
|
)
|
|
from metadata.generated.schema.entity.services.connections.database.salesforceConnection import (
|
|
SalesforceConnection,
|
|
)
|
|
from metadata.generated.schema.entity.services.connections.database.singleStoreConnection import (
|
|
SingleStoreConnection,
|
|
)
|
|
from metadata.generated.schema.entity.services.connections.database.snowflakeConnection import (
|
|
SnowflakeConnection,
|
|
)
|
|
from metadata.generated.schema.entity.services.connections.database.sqliteConnection import (
|
|
SQLiteConnection,
|
|
)
|
|
from metadata.generated.schema.entity.services.connections.database.trinoConnection import (
|
|
TrinoConnection,
|
|
)
|
|
from metadata.generated.schema.entity.services.connections.database.verticaConnection import (
|
|
VerticaConnection,
|
|
)
|
|
from metadata.generated.schema.security.credentials.gcsCredentials import GCSValues
|
|
|
|
CX_ORACLE_LIB_VERSION = "8.3.0"
|
|
|
|
|
|
def get_connection_url_common(connection):
|
|
url = f"{connection.scheme.value}://"
|
|
|
|
if connection.username:
|
|
url += f"{connection.username}"
|
|
if not connection.password:
|
|
connection.password = SecretStr("")
|
|
url += f":{quote_plus(connection.password.get_secret_value())}"
|
|
url += "@"
|
|
|
|
url += connection.hostPort
|
|
if hasattr(connection, "database"):
|
|
url += f"/{connection.database}" if connection.database else ""
|
|
|
|
elif hasattr(connection, "databaseSchema"):
|
|
url += f"/{connection.databaseSchema}" if connection.databaseSchema else ""
|
|
|
|
options = (
|
|
connection.connectionOptions.dict()
|
|
if connection.connectionOptions
|
|
else connection.connectionOptions
|
|
)
|
|
if options:
|
|
if (hasattr(connection, "database") and not connection.database) or (
|
|
hasattr(connection, "databaseSchema") and not connection.databaseSchema
|
|
):
|
|
url += "/"
|
|
params = "&".join(
|
|
f"{key}={quote_plus(value)}" for (key, value) in options.items() if value
|
|
)
|
|
url = f"{url}?{params}"
|
|
return url
|
|
|
|
|
|
@singledispatch
|
|
def get_connection_url(connection):
|
|
raise NotImplemented(
|
|
f"Connection URL build not implemented for type {type(connection)}: {connection}"
|
|
)
|
|
|
|
|
|
@get_connection_url.register(MariaDBConnection)
|
|
@get_connection_url.register(PostgresConnection)
|
|
@get_connection_url.register(RedshiftConnection)
|
|
@get_connection_url.register(MysqlConnection)
|
|
@get_connection_url.register(SalesforceConnection)
|
|
@get_connection_url.register(ClickhouseConnection)
|
|
@get_connection_url.register(SingleStoreConnection)
|
|
@get_connection_url.register(Db2Connection)
|
|
@get_connection_url.register(VerticaConnection)
|
|
def _(connection):
|
|
return get_connection_url_common(connection)
|
|
|
|
|
|
@get_connection_url.register
|
|
def _(connection: MssqlConnection):
|
|
if connection.scheme.value == connection.scheme.mssql_pyodbc.value:
|
|
return f"{connection.scheme.value}://{connection.uriString}"
|
|
return get_connection_url_common(connection)
|
|
|
|
|
|
@get_connection_url.register
|
|
def _(connection: OracleConnection):
|
|
# Patching the cx_Oracle module with oracledb lib
|
|
# to work take advantage of the thin mode of oracledb
|
|
# which doesn't require the oracle client libs to be installed
|
|
import sys
|
|
|
|
import oracledb
|
|
|
|
oracledb.version = CX_ORACLE_LIB_VERSION
|
|
sys.modules["cx_Oracle"] = oracledb
|
|
|
|
url = f"{connection.scheme.value}://"
|
|
if connection.username:
|
|
url += f"{connection.username}"
|
|
if not connection.password:
|
|
connection.password = SecretStr("")
|
|
url += f":{quote_plus(connection.password.get_secret_value())}"
|
|
url += "@"
|
|
|
|
url += connection.hostPort
|
|
|
|
if isinstance(connection.oracleConnectionType, OracleDatabaseSchema):
|
|
url += (
|
|
f"/{connection.oracleConnectionType.databaseSchema}"
|
|
if connection.oracleConnectionType.databaseSchema
|
|
else ""
|
|
)
|
|
|
|
elif isinstance(connection.oracleConnectionType, OracleServiceName):
|
|
url = f"{url}/?service_name={connection.oracleConnectionType.oracleServiceName}"
|
|
|
|
options = (
|
|
connection.connectionOptions.dict()
|
|
if connection.connectionOptions
|
|
else connection.connectionOptions
|
|
)
|
|
if options:
|
|
params = "&".join(
|
|
f"{key}={quote_plus(value)}" for (key, value) in options.items() if value
|
|
)
|
|
if isinstance(connection.oracleConnectionType, OracleServiceName):
|
|
url = f"{url}&{params}"
|
|
else:
|
|
url = f"{url}?{params}"
|
|
|
|
return url
|
|
|
|
|
|
@get_connection_url.register
|
|
def _(connection: SQLiteConnection):
|
|
"""
|
|
SQLite is only used for testing with the in-memory db
|
|
"""
|
|
|
|
database_mode = connection.databaseMode if connection.databaseMode else ":memory:"
|
|
|
|
return f"{connection.scheme.value}:///{database_mode}"
|
|
|
|
|
|
@get_connection_url.register
|
|
def _(connection: TrinoConnection):
|
|
url = f"{connection.scheme.value}://"
|
|
if connection.username:
|
|
url += f"{quote_plus(connection.username)}"
|
|
if connection.password:
|
|
url += f":{quote_plus(connection.password.get_secret_value())}"
|
|
url += "@"
|
|
url += f"{connection.hostPort}"
|
|
url += f"/{connection.catalog}"
|
|
if connection.params is not None:
|
|
params = "&".join(
|
|
f"{key}={quote_plus(value)}"
|
|
for (key, value) in connection.params.items()
|
|
if value
|
|
)
|
|
url = f"{url}?{params}"
|
|
return url
|
|
|
|
|
|
@get_connection_url.register
|
|
def _(connection: DatabricksConnection):
|
|
url = f"{connection.scheme.value}://token:{connection.token.get_secret_value()}@{connection.hostPort}"
|
|
return url
|
|
|
|
|
|
@get_connection_url.register
|
|
def _(connection: PrestoConnection):
|
|
url = f"{connection.scheme.value}://"
|
|
if connection.username:
|
|
url += f"{quote_plus(connection.username)}"
|
|
if connection.password:
|
|
url += f":{quote_plus(connection.password.get_secret_value())}"
|
|
url += "@"
|
|
url += f"{connection.hostPort}"
|
|
url += f"/{connection.catalog}"
|
|
if connection.databaseSchema:
|
|
url += f"?schema={quote_plus(connection.databaseSchema)}"
|
|
return url
|
|
|
|
|
|
@singledispatch
|
|
def get_connection_args(connection):
|
|
return connection.connectionArguments or {}
|
|
|
|
|
|
@get_connection_args.register
|
|
def _(connection: TrinoConnection):
|
|
if connection.proxies:
|
|
session = Session()
|
|
session.proxies = connection.proxies
|
|
if connection.connectionArguments:
|
|
connection_args = connection.connectionArguments.dict()
|
|
connection_args.update({"http_session": session})
|
|
return connection_args
|
|
else:
|
|
return {"http_session": session}
|
|
else:
|
|
return connection.connectionArguments if connection.connectionArguments else {}
|
|
|
|
|
|
@get_connection_url.register
|
|
def _(connection: SnowflakeConnection):
|
|
url = f"{connection.scheme.value}://"
|
|
|
|
if connection.username:
|
|
url += f"{connection.username}"
|
|
if not connection.password:
|
|
connection.password = SecretStr("")
|
|
url += (
|
|
f":{quote_plus(connection.password.get_secret_value())}"
|
|
if connection
|
|
else ""
|
|
)
|
|
url += "@"
|
|
|
|
url += connection.account
|
|
url += f"/{connection.database}" if connection.database else ""
|
|
|
|
options = (
|
|
connection.connectionOptions.dict()
|
|
if connection.connectionOptions
|
|
else connection.connectionOptions
|
|
)
|
|
if options:
|
|
if not connection.database:
|
|
url += "/"
|
|
params = "&".join(
|
|
f"{key}={quote_plus(value)}" for (key, value) in options.items() if value
|
|
)
|
|
url = f"{url}?{params}"
|
|
options = {
|
|
"account": connection.account,
|
|
"warehouse": connection.warehouse,
|
|
"role": connection.role,
|
|
}
|
|
params = "&".join(f"{key}={value}" for (key, value) in options.items() if value)
|
|
if params:
|
|
url = f"{url}?{params}"
|
|
return url
|
|
|
|
|
|
@get_connection_url.register
|
|
def _(connection: HiveConnection):
|
|
url = f"{connection.scheme.value}://"
|
|
if (
|
|
connection.username
|
|
and connection.connectionArguments
|
|
and hasattr(connection.connectionArguments, "auth")
|
|
and connection.connectionArguments.auth in ("LDAP", "CUSTOM")
|
|
):
|
|
url += f"{connection.username}"
|
|
if not connection.password:
|
|
connection.password = SecretStr("")
|
|
url += f":{quote_plus(connection.password.get_secret_value())}"
|
|
url += "@"
|
|
|
|
url += connection.hostPort
|
|
url += f"/{connection.databaseSchema}" if connection.databaseSchema else ""
|
|
|
|
options = (
|
|
connection.connectionOptions.dict()
|
|
if connection.connectionOptions
|
|
else connection.connectionOptions
|
|
)
|
|
|
|
if options:
|
|
if not connection.databaseSchema:
|
|
url += "/"
|
|
url += "/"
|
|
params = "&".join(
|
|
f"{key}={quote_plus(value)}" for (key, value) in options.items() if value
|
|
)
|
|
url = f"{url}?{params}"
|
|
if connection.authOptions:
|
|
return f"{url};{connection.authOptions}"
|
|
return url
|
|
|
|
|
|
@get_connection_url.register
|
|
def _(connection: BigQueryConnection):
|
|
from google import auth
|
|
|
|
_, project_id = auth.default()
|
|
if not project_id and isinstance(connection.credentials.gcsConfig, GCSValues):
|
|
project_id = connection.credentials.gcsConfig.projectId
|
|
if project_id:
|
|
# Setting environment variable based on project id given by user / set in ADC
|
|
os.environ["GOOGLE_CLOUD_PROJECT"] = project_id
|
|
return f"{connection.scheme.value}://{project_id}"
|
|
return f"{connection.scheme.value}://"
|
|
|
|
|
|
@get_connection_url.register
|
|
def _(connection: AzureSQLConnection):
|
|
|
|
url = f"{connection.scheme.value}://"
|
|
|
|
if connection.username:
|
|
url += f"{connection.username}"
|
|
url += f":{connection.password.get_secret_value()}" if connection else ""
|
|
url += "@"
|
|
|
|
url += f"{connection.hostPort}"
|
|
url += f"/{quote_plus(connection.database)}" if connection.database else ""
|
|
url += f"?driver={quote_plus(connection.driver)}"
|
|
options = (
|
|
connection.connectionOptions.dict()
|
|
if connection.connectionOptions
|
|
else connection.connectionOptions
|
|
)
|
|
if options:
|
|
if not connection.database:
|
|
url += "/"
|
|
params = "&".join(
|
|
f"{key}={quote_plus(value)}" for (key, value) in options.items() if value
|
|
)
|
|
url = f"{url}?{params}"
|
|
|
|
return url
|
|
|
|
|
|
@get_connection_url.register
|
|
def _(connection: AthenaConnection):
|
|
url = f"{connection.scheme.value}://"
|
|
if connection.awsConfig.awsAccessKeyId:
|
|
url += connection.awsConfig.awsAccessKeyId
|
|
if connection.awsConfig.awsSecretAccessKey:
|
|
url += f":{connection.awsConfig.awsSecretAccessKey.get_secret_value()}"
|
|
else:
|
|
url += ":"
|
|
url += f"@athena.{connection.awsConfig.awsRegion}.amazonaws.com:443"
|
|
|
|
url += f"?s3_staging_dir={quote_plus(connection.s3StagingDir)}"
|
|
if connection.workgroup:
|
|
url += f"&work_group={connection.workgroup}"
|
|
if connection.awsConfig.awsSessionToken:
|
|
url += f"&aws_session_token={quote_plus(connection.awsConfig.awsSessionToken)}"
|
|
|
|
return url
|
|
|
|
|
|
@get_connection_url.register
|
|
def _(connection: DruidConnection):
|
|
url = get_connection_url_common(connection)
|
|
return f"{url}/druid/v2/sql"
|
|
|
|
|
|
@get_connection_url.register
|
|
def _(connection: PinotDBConnection):
|
|
url = get_connection_url_common(connection)
|
|
url += f"/query/sql?controller={connection.pinotControllerHost}"
|
|
return url
|