21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118 | class ArkUAPSIADBInstanceTarget(ArkCamelizedModel):
instance_name: Annotated[
str, Field(strict=True, min_length=1, max_length=UAP_SIA_DB_INSTANCE_NAME_LENGTH, description='The name of the database instance')
]
instance_type: Annotated[ArkSIADBDatabaseFamilyType, Field(description='The database type of the database instance')]
instance_id: Annotated[
str, Field(strict=True, min_length=1, max_length=UAP_SIA_DB_INSTANCE_ID_LENGTH, description='The id of the database instance')
]
authentication_method: Annotated[
ArkUAPSIADBAuthenticationMethod, Field(description='The authentication method corresponding to this profile')
]
# Profiles, only one of these will be set based on the authentication method.
# Note that the API has a single profiles field, but we separate them here for clarity and easier usage.
ldap_auth_profile: Annotated[
Optional[ArkUAPSIADBLDAPAuthProfile], Field(description='The LDAP authentication profile for this database instance')
] = None
db_auth_profile: Annotated[
Optional[ArkUAPSIADBLocalDBAuthProfile],
Field(description='The local database authentication profile for this database instance'),
] = None
oracle_db_auth_profile: Annotated[
ArkUAPSIADBOracleDBAuthProfile,
Field(description='The Oracle database authentication profile for this database instance'),
] = None
mongo_auth_profile: Annotated[
Optional[ArkUAPSIADBMongoAuthProfile],
Field(description='The MongoDB authentication profile for this database instance'),
] = None
sql_server_auth_profile: Annotated[
Optional[ArkUAPSIADBSqlServerAuthProfile],
Field(description='The SQL Server authentication profile for this database instance'),
] = None
rds_iam_user_auth_profile: Annotated[
Optional[ArkUAPSIADBRDSIAMUserAuthProfile],
Field(description='The RDS IAM User authentication profile for this database instance'),
] = None
def databases_count(self) -> int:
profile = self.profile_by_authentication_method()
if not profile:
raise ValueError(
f'No profile found for the given authentication method, instance: [{self.instance_name}], authentication method: [{self.authentication_method}]'
)
return profile.databases_count()
def profile_by_authentication_method(self) -> ArkUAPSIADBProfile:
"""
Returns the profile corresponding to the authentication method.
"""
if self.authentication_method == ArkUAPSIADBAuthenticationMethod.LDAP_AUTH:
return self.ldap_auth_profile
if self.authentication_method == ArkUAPSIADBAuthenticationMethod.DB_AUTH:
return self.db_auth_profile
if self.authentication_method == ArkUAPSIADBAuthenticationMethod.ORACLE_AUTH:
return self.oracle_db_auth_profile
if self.authentication_method == ArkUAPSIADBAuthenticationMethod.MONGO_AUTH:
return self.mongo_auth_profile
if self.authentication_method == ArkUAPSIADBAuthenticationMethod.SQLSERVER_AUTH:
return self.sql_server_auth_profile
if self.authentication_method == ArkUAPSIADBAuthenticationMethod.RDS_IAM_USER_AUTH:
return self.rds_iam_user_auth_profile
raise ValueError(f'Unsupported authentication method: {self.authentication_method}')
@model_validator(mode='before')
@classmethod
def move_profile_to_specific_field(cls, data: Dict[str, Any]) -> Dict[str, Any]:
if 'profile' not in data:
return data
auth_method = data.get('authentication_method') or data.get('authenticationMethod')
if not auth_method:
return data # This case will throw an error later since authentication method is defined as required,
if isinstance(auth_method, str):
auth_method = ArkUAPSIADBAuthenticationMethod(auth_method)
profile_field_name = cls.auth_method_to_profile_field_name().get(auth_method)
if not profile_field_name:
raise ValueError(f'Unsupported authentication method for profile: {auth_method}')
# Move 'profile' to the specific field.
data[profile_field_name] = data.pop('profile')
return data
@classmethod
def auth_method_to_profile_field_name(cls) -> Dict[ArkUAPSIADBAuthenticationMethod, str]:
return {
ArkUAPSIADBAuthenticationMethod.LDAP_AUTH: 'ldapAuthProfile',
ArkUAPSIADBAuthenticationMethod.DB_AUTH: 'dbAuthProfile',
ArkUAPSIADBAuthenticationMethod.ORACLE_AUTH: 'oracleDbAuthProfile',
ArkUAPSIADBAuthenticationMethod.MONGO_AUTH: 'mongoAuthProfile',
ArkUAPSIADBAuthenticationMethod.SQLSERVER_AUTH: 'sqlServerAuthProfile',
ArkUAPSIADBAuthenticationMethod.RDS_IAM_USER_AUTH: 'rdsIamUserAuthProfile',
}
|