3434 ValidationInfo ,
3535 field_validator ,
3636 model_validator ,
37+ validation_data ,
3738 validation_error_message ,
3839 get_concrete_types_from_typehint ,
3940)
@@ -813,6 +814,8 @@ class DatabricksConnectionConfig(ConnectionConfig):
813814 DISPLAY_NAME : t .ClassVar [t .Literal ["Databricks" ]] = "Databricks"
814815 DISPLAY_ORDER : t .ClassVar [t .Literal [3 ]] = 3
815816
817+ shared_connection : t .ClassVar [bool ] = True
818+
816819 _concurrent_tasks_validator = concurrent_tasks_validator
817820 _http_headers_validator = http_headers_validator
818821
@@ -1081,7 +1084,7 @@ def validate_execution_project(
10811084 v : t .Optional [str ],
10821085 info : ValidationInfo ,
10831086 ) -> t .Optional [str ]:
1084- if v and not info . data .get ("project" ):
1087+ if v and not validation_data ( info ) .get ("project" ):
10851088 raise ConfigError (
10861089 "If the `execution_project` field is specified, you must also specify the `project` field to provide a default object location."
10871090 )
@@ -1093,7 +1096,7 @@ def validate_quota_project(
10931096 v : t .Optional [str ],
10941097 info : ValidationInfo ,
10951098 ) -> t .Optional [str ]:
1096- if v and not info . data .get ("project" ):
1099+ if v and not validation_data ( info ) .get ("project" ):
10971100 raise ConfigError (
10981101 "If the `quota_project` field is specified, you must also specify the `project` field to provide a default object location."
10991102 )
@@ -2341,35 +2344,28 @@ def init(cursor: t.Any) -> None:
23412344 return init
23422345
23432346
2347+ _CONNECTION_CONFIG_EXCLUDE : t .Set [t .Type [ConnectionConfig ]] = {
2348+ ConnectionConfig , # type: ignore[type-abstract]
2349+ BaseDuckDBConnectionConfig , # type: ignore[type-abstract]
2350+ }
2351+
23442352CONNECTION_CONFIG_TO_TYPE = {
23452353 # Map all subclasses of ConnectionConfig to the value of their `type_` field.
23462354 tpe .all_field_infos ()["type_" ].default : tpe
2347- for tpe in subclasses (
2348- __name__ ,
2349- ConnectionConfig ,
2350- exclude = {ConnectionConfig , BaseDuckDBConnectionConfig },
2351- )
2355+ for tpe in subclasses (__name__ , ConnectionConfig , exclude = _CONNECTION_CONFIG_EXCLUDE )
23522356}
23532357
23542358DIALECT_TO_TYPE = {
23552359 tpe .all_field_infos ()["type_" ].default : tpe .DIALECT
2356- for tpe in subclasses (
2357- __name__ ,
2358- ConnectionConfig ,
2359- exclude = {ConnectionConfig , BaseDuckDBConnectionConfig },
2360- )
2360+ for tpe in subclasses (__name__ , ConnectionConfig , exclude = _CONNECTION_CONFIG_EXCLUDE )
23612361}
23622362
23632363INIT_DISPLAY_INFO_TO_TYPE = {
23642364 tpe .all_field_infos ()["type_" ].default : (
23652365 tpe .DISPLAY_ORDER ,
23662366 tpe .DISPLAY_NAME ,
23672367 )
2368- for tpe in subclasses (
2369- __name__ ,
2370- ConnectionConfig ,
2371- exclude = {ConnectionConfig , BaseDuckDBConnectionConfig },
2372- )
2368+ for tpe in subclasses (__name__ , ConnectionConfig , exclude = _CONNECTION_CONFIG_EXCLUDE )
23732369}
23742370
23752371
0 commit comments