Skip to content

fix: Add postgres compatibility in PrestoViewMetadataExtractor #1878

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 2 commits into from
Aug 2, 2022
Merged

fix: Add postgres compatibility in PrestoViewMetadataExtractor #1878

Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
f5fe84a
Provide postgres compatibility in presto view metadata extractor
May 28, 2022
5a8fae1
Fix ci UT: check the beginning of connection string to determine whet…
chonyy Jul 11, 2022
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
24 changes: 22 additions & 2 deletions databuilder/databuilder/extractor/presto_view_metadata_extractor.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -24,7 +24,7 @@ class PrestoViewMetadataExtractor(Extractor):
"""
# SQL statement to extract View metadata
# {where_clause_suffix} could be used to filter schemas
SQL_STATEMENT = """
DEFAULT_SQL_STATEMENT = """
SELECT t.TBL_ID, d.NAME as `schema`, t.TBL_NAME name, t.TBL_TYPE, t.VIEW_ORIGINAL_TEXT as view_original_text
FROM TBLS t
JOIN DBS d ON t.DB_ID = d.DB_ID
Expand All @@ -33,6 +33,19 @@ class PrestoViewMetadataExtractor(Extractor):
ORDER BY t.TBL_ID desc;
"""

DEFAULT_POSTGRES_SQL_STATEMENT = """
SELECT t."TBL_ID",
d."NAME" as "schema",
t."TBL_NAME" as name,
t."TBL_TYPE",
t."VIEW_ORIGINAL_TEXT" as view_original_text
FROM "TBLS" t
JOIN "DBS" d ON t."DB_ID" = d."DB_ID"
WHERE t."VIEW_EXPANDED_TEXT" = '/* Presto View */'
{where_clause_suffix}
ORDER BY t."TBL_ID" desc;
"""

# Presto View data prefix and suffix definition:
# https://github.com/prestodb/presto/blob/43bd519052ba4c56ff1f4fc807075637ab5f4f10/presto-hive/src/main/java/com/facebook/presto/hive/HiveUtil.java#L153-L154
PRESTO_VIEW_PREFIX = '/* Presto View: '
Expand All @@ -49,14 +62,21 @@ def init(self, conf: ConfigTree) -> None:
conf = conf.with_fallback(PrestoViewMetadataExtractor.DEFAULT_CONFIG)
self._cluster = conf.get_string(PrestoViewMetadataExtractor.CLUSTER_KEY)

self.sql_stmt = PrestoViewMetadataExtractor.SQL_STATEMENT.format(
self.sql_stmt = self._choose_default_sql_stm(conf).format(
where_clause_suffix=conf.get_string(PrestoViewMetadataExtractor.WHERE_CLAUSE_SUFFIX_KEY))

LOGGER.info('SQL for hive metastore: %s', self.sql_stmt)

self._alchemy_extractor = sql_alchemy_extractor.from_surrounding_config(conf, self.sql_stmt)
self._extract_iter: Union[None, Iterator] = None

def _choose_default_sql_stm(self, conf: ConfigTree) -> str:
conn_string = conf.get_string("extractor.sqlalchemy.conn_string")
if conn_string.startswith('postgres') or conn_string.startswith('postgresql'):
return self.DEFAULT_POSTGRES_SQL_STATEMENT
else:
return self.DEFAULT_SQL_STATEMENT

def close(self) -> None:
if getattr(self, '_alchemy_extractor', None) is not None:
self._alchemy_extractor.close()
Expand Down