diff --git a/gaiadmpsetup/gaiadmpsetup.py b/gaiadmpsetup/gaiadmpsetup.py index 8a5b733..c28871a 100644 --- a/gaiadmpsetup/gaiadmpsetup.py +++ b/gaiadmpsetup/gaiadmpsetup.py @@ -1,5 +1,6 @@ from pyspark.sql.types import * from pyspark.sql.session import SparkSession +from pyspark.sql.utils import AnalysisException from . import gaiaedr3_pyspark_schema_structures as edr3 from . import gaiadr3_pyspark_schema_structures as dr3 @@ -18,17 +19,24 @@ def __init__(self): @staticmethod def setup(): - def tablesExist(expected_tables): - actual_tables = [i.name for i in spark.catalog.listTables()] - check = all(item in actual_tables for item in expected_tables) + def tablesExist(expected_tables, database): + + check = False + + try: + + spark.sql("use " + database) + actual_tables = [i.name for i in spark.catalog.listTables()] + check = all(item in actual_tables for item in expected_tables) + + except AnalysisException: pass + return check # check EDR3 - if not tablesExist(edr3.table_dict.keys()): + database = "gaiaedr3" + if not tablesExist(edr3.table_dict.keys(), database): - # database name to create - database = "gaiaedr3" - # create the database and switch the current SQL database context to it (from default) spark.sql("create database if not exists " + database) spark.sql("use " + database) @@ -40,10 +48,10 @@ def tablesExist(expected_tables): reattachParquetFileResourceToSparkContext(table_key, data_store + folder_path, schemas) # check DR3 - if not tablesExist(dr3.table_dict.keys()): + database = "gaiadr3" + if not tablesExist(dr3.table_dict.keys(), database): # ... similarly for Gaia DR3 - database = "gaiadr3" spark.sql("create database if not exists " + database) spark.sql("use " + database)