required-fields-updated-in-snowflake-and-athena (#5143)

required-fields-updated-in-snowflake-and-athena (#5143)
This commit is contained in:
Abhishek Pandey 2022-05-30 18:27:26 +05:30 committed by GitHub
parent f8f79ebb2d
commit 12e8a1fcf6
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
6 changed files with 23 additions and 22 deletions

View File

@ -41,11 +41,6 @@
"description": "Host and port of the Athena service.",
"type": "string"
},
"database": {
"title": "Database",
"description": "Database of the data source. This is optional parameter, if you would like to restrict the metadata reading to a single database. When left blank, OpenMetadata Ingestion attempts to scan all the databases.",
"type": "string"
},
"s3StagingDir": {
"title": "S3 Staging Directory",
"description": "S3 Staging Directory.",
@ -73,5 +68,6 @@
"$ref": "../connectionBasicType.json#/definitions/supportsProfiler"
}
},
"additionalProperties": false
"additionalProperties": false,
"required": ["s3StagingDir", "awsConfig", "workgroup"]
}

View File

@ -96,5 +96,5 @@
}
},
"additionalProperties": false,
"required": ["username", "account"]
"required": ["username", "account", "password", "warehouse"]
}

View File

@ -137,7 +137,7 @@ class SqlColumnHandler:
)
table_columns = []
columns = inspector.get_columns(
table, schema, db_name=self.service_connection.database
table, schema, db_name=self._get_database_name()
)
for column in columns:
try:

View File

@ -127,12 +127,17 @@ class SqlAlchemySource(Source, ABC):
Method to fetch tags associated with table
"""
def get_database_entity(self, database_name: Optional[str]) -> Database:
def _get_database_name(self) -> str:
if hasattr(self.service_connection, "database"):
return self.service_connection.database or "default"
return "default"
def get_database_entity(self) -> Database:
"""
Method to get database enetity from db name
"""
return Database(
name=database_name if database_name else "default",
name=self._get_database_name(),
service=EntityReference(
id=self.service.id, type=self.service_connection.type.value
),
@ -173,7 +178,7 @@ class SqlAlchemySource(Source, ABC):
self.metadata,
entity_type=DatabaseSchema,
service_name=self.config.serviceName,
database_name=self.service_connection.database,
database_name=self._get_database_name(),
schema_name=schema,
)
yield from self.delete_tables(schema_fqn)
@ -253,7 +258,7 @@ class SqlAlchemySource(Source, ABC):
schema, table_name, table_type, inspector
)
database = self.get_database_entity(self.service_connection.database)
database = self.get_database_entity()
# check if we have any model to associate with
table_entity.dataModel = self.get_data_model(
database.name.__root__, schema, table_name

View File

@ -97,7 +97,8 @@ def get_connection_url_common(connection):
url += "@"
url += connection.hostPort
url += f"/{connection.database}" if connection.database else ""
if hasattr(connection, "database"):
url += f"/{connection.database}" if connection.database else ""
options = (
connection.connectionOptions.dict()
@ -347,10 +348,10 @@ def _(connection: AthenaConnection):
else:
url += ":"
url += f"@athena.{connection.awsConfig.awsRegion}.amazonaws.com:443"
if connection.database:
url += f"/{connection.database}"
url += f"?s3_staging_dir={quote_plus(connection.s3StagingDir)}"
url += f"&work_group={connection.workgroup}"
if connection.workgroup:
url += f"&work_group={connection.workgroup}"
return url

View File

@ -625,9 +625,9 @@ class SouceConnectionTest(TestCase):
expected_args = {}
snowflake_conn_obj = SnowflakeConnection(
username="user",
password=None,
password="test-pwd",
database="tiny",
connectionArguments=None,
warehouse="COMPUTE_WH",
scheme=SnowflakeScheme.snowflake,
account="account.region_name.cloud_service",
)
@ -637,8 +637,9 @@ class SouceConnectionTest(TestCase):
expected_args = {"user": "user-to-be-impersonated"}
snowflake_conn_obj = SnowflakeConnection(
username="user",
password=None,
password="test-pwd",
database="tiny",
warehouse="COMPUTE_WH",
connectionArguments={"user": "user-to-be-impersonated"},
scheme=SnowflakeScheme.snowflake,
account="account.region_name.cloud_service",
@ -657,18 +658,16 @@ class SouceConnectionTest(TestCase):
s3StagingDir="s3athena-postgres",
workgroup="primary",
scheme=AthenaScheme.awsathena_rest,
database=None,
)
assert expected_url == get_connection_url(athena_conn_obj)
# connection arguments witho db
expected_url = "awsathena+rest://key:secret_key@athena.us-east-2.amazonaws.com:443/test?s3_staging_dir=s3athena-postgres&work_group=primary"
expected_url = "awsathena+rest://key:secret_key@athena.us-east-2.amazonaws.com:443?s3_staging_dir=s3athena-postgres&work_group=primary"
athena_conn_obj = AthenaConnection(
awsConfig=awsCreds,
s3StagingDir="s3athena-postgres",
workgroup="primary",
scheme=AthenaScheme.awsathena_rest,
database="test",
)
assert expected_url == get_connection_url(athena_conn_obj)