Skip to content

Commit e1a8daa

Browse files
kanterovkanterov
authored
[Python] Add new task types (#2788)
## Changes Add new task types that have been recently added to the JSON schema. ## Why It ensures compatibility of Python code with the latest JSON schema
1 parent 288a449 commit e1a8daa

12 files changed

Lines changed: 511 additions & 6 deletions

File tree

‎experimental/python/codegen/codegen/jsonschema_patch.py‎

Lines changed: 0 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -3,11 +3,6 @@
33
from codegen.jsonschema import Schema
44

55
REMOVED_FIELDS = {
6-
# TODO remove as a follow-up
7-
"jobs.Task": {
8-
"dashboard_task",
9-
"power_bi_task",
10-
},
116
"compute.ClusterSpec": {
127
# doesn't work, openapi schema needs to be updated to be enum
138
"kind",

‎experimental/python/databricks/bundles/jobs/__init__.py‎

Lines changed: 57 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2,6 +2,8 @@
22
"Adlsgen2Info",
33
"Adlsgen2InfoDict",
44
"Adlsgen2InfoParam",
5+
"AuthenticationMethod",
6+
"AuthenticationMethodParam",
57
"AutoScale",
68
"AutoScaleDict",
79
"AutoScaleParam",
@@ -43,6 +45,9 @@
4345
"CronSchedule",
4446
"CronScheduleDict",
4547
"CronScheduleParam",
48+
"DashboardTask",
49+
"DashboardTaskDict",
50+
"DashboardTaskParam",
4651
"DataSecurityMode",
4752
"DataSecurityModeParam",
4853
"DbfsStorageInfo",
@@ -153,6 +158,15 @@
153158
"PipelineTask",
154159
"PipelineTaskDict",
155160
"PipelineTaskParam",
161+
"PowerBiModel",
162+
"PowerBiModelDict",
163+
"PowerBiModelParam",
164+
"PowerBiTable",
165+
"PowerBiTableDict",
166+
"PowerBiTableParam",
167+
"PowerBiTask",
168+
"PowerBiTaskDict",
169+
"PowerBiTaskParam",
156170
"PythonPyPiLibrary",
157171
"PythonPyPiLibraryDict",
158172
"PythonPyPiLibraryParam",
@@ -204,6 +218,14 @@
204218
"SqlTaskSubscription",
205219
"SqlTaskSubscriptionDict",
206220
"SqlTaskSubscriptionParam",
221+
"StorageMode",
222+
"StorageModeParam",
223+
"Subscription",
224+
"SubscriptionDict",
225+
"SubscriptionParam",
226+
"SubscriptionSubscriber",
227+
"SubscriptionSubscriberDict",
228+
"SubscriptionSubscriberParam",
207229
"TableUpdateTriggerConfiguration",
208230
"TableUpdateTriggerConfigurationDict",
209231
"TableUpdateTriggerConfigurationParam",
@@ -384,6 +406,10 @@
384406
WorkspaceStorageInfoDict,
385407
WorkspaceStorageInfoParam,
386408
)
409+
from databricks.bundles.jobs._models.authentication_method import (
410+
AuthenticationMethod,
411+
AuthenticationMethodParam,
412+
)
387413
from databricks.bundles.jobs._models.clean_rooms_notebook_task import (
388414
CleanRoomsNotebookTask,
389415
CleanRoomsNotebookTaskDict,
@@ -414,6 +440,11 @@
414440
CronScheduleDict,
415441
CronScheduleParam,
416442
)
443+
from databricks.bundles.jobs._models.dashboard_task import (
444+
DashboardTask,
445+
DashboardTaskDict,
446+
DashboardTaskParam,
447+
)
417448
from databricks.bundles.jobs._models.dbt_task import DbtTask, DbtTaskDict, DbtTaskParam
418449
from databricks.bundles.jobs._models.file_arrival_trigger_configuration import (
419450
FileArrivalTriggerConfiguration,
@@ -523,6 +554,21 @@
523554
PipelineTaskDict,
524555
PipelineTaskParam,
525556
)
557+
from databricks.bundles.jobs._models.power_bi_model import (
558+
PowerBiModel,
559+
PowerBiModelDict,
560+
PowerBiModelParam,
561+
)
562+
from databricks.bundles.jobs._models.power_bi_table import (
563+
PowerBiTable,
564+
PowerBiTableDict,
565+
PowerBiTableParam,
566+
)
567+
from databricks.bundles.jobs._models.power_bi_task import (
568+
PowerBiTask,
569+
PowerBiTaskDict,
570+
PowerBiTaskParam,
571+
)
526572
from databricks.bundles.jobs._models.python_wheel_task import (
527573
PythonWheelTask,
528574
PythonWheelTaskDict,
@@ -581,6 +627,17 @@
581627
SqlTaskSubscriptionDict,
582628
SqlTaskSubscriptionParam,
583629
)
630+
from databricks.bundles.jobs._models.storage_mode import StorageMode, StorageModeParam
631+
from databricks.bundles.jobs._models.subscription import (
632+
Subscription,
633+
SubscriptionDict,
634+
SubscriptionParam,
635+
)
636+
from databricks.bundles.jobs._models.subscription_subscriber import (
637+
SubscriptionSubscriber,
638+
SubscriptionSubscriberDict,
639+
SubscriptionSubscriberParam,
640+
)
584641
from databricks.bundles.jobs._models.table_update_trigger_configuration import (
585642
TableUpdateTriggerConfiguration,
586643
TableUpdateTriggerConfigurationDict,

‎experimental/python/databricks/bundles/jobs/_models/authentication_method.py‎

Lines changed: 14 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,14 @@
1+
from enum import Enum
2+
from typing import Literal
3+
4+
5+
class AuthenticationMethod(Enum):
6+
"""
7+
:meta private: [EXPERIMENTAL]
8+
"""
9+
10+
OAUTH = "OAUTH"
11+
PAT = "PAT"
12+
13+
14+
AuthenticationMethodParam = Literal["OAUTH", "PAT"] | AuthenticationMethod

‎experimental/python/databricks/bundles/jobs/_models/dashboard_task.py‎

Lines changed: 54 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,54 @@
1+
from dataclasses import dataclass
2+
from typing import TYPE_CHECKING, TypedDict
3+
4+
from databricks.bundles.core._transform import _transform
5+
from databricks.bundles.core._transform_to_json import _transform_to_json_value
6+
from databricks.bundles.core._variable import VariableOrOptional
7+
from databricks.bundles.jobs._models.subscription import (
8+
Subscription,
9+
SubscriptionParam,
10+
)
11+
12+
if TYPE_CHECKING:
13+
from typing_extensions import Self
14+
15+
16+
@dataclass(kw_only=True)
17+
class DashboardTask:
18+
"""
19+
Configures the Lakeview Dashboard job task type.
20+
"""
21+
22+
dashboard_id: VariableOrOptional[str] = None
23+
24+
subscription: VariableOrOptional[Subscription] = None
25+
26+
warehouse_id: VariableOrOptional[str] = None
27+
"""
28+
Optional: The warehouse id to execute the dashboard with for the schedule.
29+
If not specified, the default warehouse of the dashboard will be used.
30+
"""
31+
32+
@classmethod
33+
def from_dict(cls, value: "DashboardTaskDict") -> "Self":
34+
return _transform(cls, value)
35+
36+
def as_dict(self) -> "DashboardTaskDict":
37+
return _transform_to_json_value(self) # type:ignore
38+
39+
40+
class DashboardTaskDict(TypedDict, total=False):
41+
""""""
42+
43+
dashboard_id: VariableOrOptional[str]
44+
45+
subscription: VariableOrOptional[SubscriptionParam]
46+
47+
warehouse_id: VariableOrOptional[str]
48+
"""
49+
Optional: The warehouse id to execute the dashboard with for the schedule.
50+
If not specified, the default warehouse of the dashboard will be used.
51+
"""
52+
53+
54+
DashboardTaskParam = DashboardTaskDict | DashboardTask

‎experimental/python/databricks/bundles/jobs/_models/job.py‎

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -17,7 +17,10 @@
1717
CronSchedule,
1818
CronScheduleParam,
1919
)
20-
from databricks.bundles.jobs._models.git_source import GitSource, GitSourceParam
20+
from databricks.bundles.jobs._models.git_source import (
21+
GitSource,
22+
GitSourceParam,
23+
)
2124
from databricks.bundles.jobs._models.job_cluster import JobCluster, JobClusterParam
2225
from databricks.bundles.jobs._models.job_email_notifications import (
2326
JobEmailNotifications,

‎experimental/python/databricks/bundles/jobs/_models/power_bi_model.py‎

Lines changed: 85 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,85 @@
1+
from dataclasses import dataclass
2+
from typing import TYPE_CHECKING, TypedDict
3+
4+
from databricks.bundles.core._transform import _transform
5+
from databricks.bundles.core._transform_to_json import _transform_to_json_value
6+
from databricks.bundles.core._variable import VariableOrOptional
7+
from databricks.bundles.jobs._models.authentication_method import (
8+
AuthenticationMethod,
9+
AuthenticationMethodParam,
10+
)
11+
from databricks.bundles.jobs._models.storage_mode import StorageMode, StorageModeParam
12+
13+
if TYPE_CHECKING:
14+
from typing_extensions import Self
15+
16+
17+
@dataclass(kw_only=True)
18+
class PowerBiModel:
19+
"""
20+
:meta private: [EXPERIMENTAL]
21+
"""
22+
23+
authentication_method: VariableOrOptional[AuthenticationMethod] = None
24+
"""
25+
How the published Power BI model authenticates to Databricks
26+
"""
27+
28+
model_name: VariableOrOptional[str] = None
29+
"""
30+
The name of the Power BI model
31+
"""
32+
33+
overwrite_existing: VariableOrOptional[bool] = None
34+
"""
35+
Whether to overwrite existing Power BI models
36+
"""
37+
38+
storage_mode: VariableOrOptional[StorageMode] = None
39+
"""
40+
The default storage mode of the Power BI model
41+
"""
42+
43+
workspace_name: VariableOrOptional[str] = None
44+
"""
45+
The name of the Power BI workspace of the model
46+
"""
47+
48+
@classmethod
49+
def from_dict(cls, value: "PowerBiModelDict") -> "Self":
50+
return _transform(cls, value)
51+
52+
def as_dict(self) -> "PowerBiModelDict":
53+
return _transform_to_json_value(self) # type:ignore
54+
55+
56+
class PowerBiModelDict(TypedDict, total=False):
57+
""""""
58+
59+
authentication_method: VariableOrOptional[AuthenticationMethodParam]
60+
"""
61+
How the published Power BI model authenticates to Databricks
62+
"""
63+
64+
model_name: VariableOrOptional[str]
65+
"""
66+
The name of the Power BI model
67+
"""
68+
69+
overwrite_existing: VariableOrOptional[bool]
70+
"""
71+
Whether to overwrite existing Power BI models
72+
"""
73+
74+
storage_mode: VariableOrOptional[StorageModeParam]
75+
"""
76+
The default storage mode of the Power BI model
77+
"""
78+
79+
workspace_name: VariableOrOptional[str]
80+
"""
81+
The name of the Power BI workspace of the model
82+
"""
83+
84+
85+
PowerBiModelParam = PowerBiModelDict | PowerBiModel

‎experimental/python/databricks/bundles/jobs/_models/power_bi_table.py‎

Lines changed: 71 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,71 @@
1+
from dataclasses import dataclass
2+
from typing import TYPE_CHECKING, TypedDict
3+
4+
from databricks.bundles.core._transform import _transform
5+
from databricks.bundles.core._transform_to_json import _transform_to_json_value
6+
from databricks.bundles.core._variable import VariableOrOptional
7+
from databricks.bundles.jobs._models.storage_mode import StorageMode, StorageModeParam
8+
9+
if TYPE_CHECKING:
10+
from typing_extensions import Self
11+
12+
13+
@dataclass(kw_only=True)
14+
class PowerBiTable:
15+
"""
16+
:meta private: [EXPERIMENTAL]
17+
"""
18+
19+
catalog: VariableOrOptional[str] = None
20+
"""
21+
The catalog name in Databricks
22+
"""
23+
24+
name: VariableOrOptional[str] = None
25+
"""
26+
The table name in Databricks
27+
"""
28+
29+
schema: VariableOrOptional[str] = None
30+
"""
31+
The schema name in Databricks
32+
"""
33+
34+
storage_mode: VariableOrOptional[StorageMode] = None
35+
"""
36+
The Power BI storage mode of the table
37+
"""
38+
39+
@classmethod
40+
def from_dict(cls, value: "PowerBiTableDict") -> "Self":
41+
return _transform(cls, value)
42+
43+
def as_dict(self) -> "PowerBiTableDict":
44+
return _transform_to_json_value(self) # type:ignore
45+
46+
47+
class PowerBiTableDict(TypedDict, total=False):
48+
""""""
49+
50+
catalog: VariableOrOptional[str]
51+
"""
52+
The catalog name in Databricks
53+
"""
54+
55+
name: VariableOrOptional[str]
56+
"""
57+
The table name in Databricks
58+
"""
59+
60+
schema: VariableOrOptional[str]
61+
"""
62+
The schema name in Databricks
63+
"""
64+
65+
storage_mode: VariableOrOptional[StorageModeParam]
66+
"""
67+
The Power BI storage mode of the table
68+
"""
69+
70+
71+
PowerBiTableParam = PowerBiTableDict | PowerBiTable

0 commit comments

Comments
 (0)