1919
2020BIGQUERY_TYPE = "BIGQUERY"
2121DATABRICKS_TYPE = "DATABRICKS"
22+ SNOWFLAKE_TYPE = "SNOWFLAKE"
23+ PRIVATE_KEY = "private_key"
24+ PRIVATE_KEY_PASSPHRASE = "private_key_passphrase"
2225LAYOUT_DATA_SOURCES_DIR = "data_sources"
2326
2427
@@ -30,25 +33,45 @@ def _inject_base(self, credentials: dict[str, Any]) -> DeclarativeDataSources:
3033 data_sources = []
3134 client_class = self .client_class ()
3235 for data_source in self .data_sources :
33- if data_source .id in credentials :
36+ data_source_credentials = credentials .get (data_source .id )
37+ if data_source_credentials is None :
38+ data_sources .append (data_source .to_api ())
39+ else :
3440 if data_source .type == BIGQUERY_TYPE :
35- token = TokenCredentialsFromFile .token_from_file (credentials [ data_source . id ] )
41+ token = TokenCredentialsFromFile .token_from_file (data_source_credentials )
3642 data_sources .append (data_source .to_api (token = token ))
3743 elif data_source .type == DATABRICKS_TYPE :
3844 if data_source .client_id and data_source .client_id .strip ():
39- client_secret = ClientSecretCredentialsFromFile .client_secret_from_file (
40- credentials [data_source .id ]
41- )
45+ client_secret = ClientSecretCredentialsFromFile .client_secret_from_file (data_source_credentials )
4246 data_sources .append (data_source .to_api (client_secret = client_secret ))
4347 else :
4448 token = TokenCredentialsFromFile .token_from_file (
45- file_path = credentials [ data_source . id ] , base64_encode = False
49+ file_path = data_source_credentials , base64_encode = False
4650 )
4751 data_sources .append (data_source .to_api (token = token ))
52+ elif data_source .type == SNOWFLAKE_TYPE :
53+ if isinstance (data_source_credentials , str ):
54+ data_sources .append (data_source .to_api (password = data_source_credentials ))
55+ elif isinstance (data_source_credentials , dict ):
56+ private_key = data_source_credentials .get (PRIVATE_KEY )
57+ private_key_passphrase = data_source_credentials .get (PRIVATE_KEY_PASSPHRASE )
58+ if private_key is None :
59+ raise ValueError (
60+ f"Credentials for data source { data_source .id } should contain { PRIVATE_KEY } but it is missing."
61+ )
62+ else :
63+ data_sources .append (
64+ data_source .to_api (
65+ private_key = private_key , private_key_passphrase = private_key_passphrase
66+ )
67+ )
68+ else :
69+ raise ValueError (
70+ f"Credentials for data source { data_source .id } should be a string or a dictionary, "
71+ f"but got { type (data_source_credentials )} ."
72+ )
4873 else :
49- data_sources .append (data_source .to_api (password = credentials [data_source .id ]))
50- else :
51- data_sources .append (data_source .to_api ())
74+ data_sources .append (data_source .to_api (password = data_source_credentials ))
5275 return client_class (data_sources = data_sources )
5376
5477 def _inject_credentials_legacy (self , credentials : dict [str , Any ]) -> DeclarativeDataSources :
@@ -117,6 +140,7 @@ class CatalogDeclarativeDataSource(Base):
117140 decoded_parameters : Optional [list [CatalogParameter ]] = None
118141 permissions : list [CatalogDeclarativeDataSourcePermission ] = attr .field (factory = list )
119142 client_id : Optional [str ] = None
143+ authentication_type : Optional [str ] = None
120144
121145 def to_test_request (
122146 self ,
0 commit comments