Coverage for python / lsst / dax / apdb / cli / options.py: 20%

1# This file is part of dax_apdb 

2# 

3# Developed for the LSST Data Management System. 

4# This product includes software developed by the LSST Project 

5# (https://www.lsst.org). 

6# See the COPYRIGHT file at the top-level directory of this distribution 

7# for details of code ownership. 

8# 

9# This program is free software: you can redistribute it and/or modify 

10# it under the terms of the GNU General Public License as published by 

11# the Free Software Foundation, either version 3 of the License, or 

12# (at your option) any later version. 

13# 

14# This program is distributed in the hope that it will be useful, 

15# but WITHOUT ANY WARRANTY; without even the implied warranty of 

16# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the 

17# GNU General Public License for more details. 

18# 

19# You should have received a copy of the GNU General Public License 

20# along with this program. If not, see <https://www.gnu.org/licenses/>. 

21 

22from __future__ import annotations 

23 

24__all__ = ["cassandra_config_options", "common_apdb_options", "sql_config_options"] 

25 

26import argparse 

27from typing import TYPE_CHECKING, Any 

28 

29from ..apdb import ApdbConfig 

30from ..cassandra import ApdbCassandraConfig 

31from ..cassandra.config import ApdbCassandraConnectionConfig, ApdbCassandraPartitioningConfig 

32from ..sql import ApdbSqlConfig 

33from ..sql.config import ApdbSqlConnectionConfig, ApdbSqlPixelizationConfig 

34 

35if TYPE_CHECKING: 

36 from pydantic import BaseModel 

37 

38 

39def _option_from_pydantic_field( 

40 group: argparse._ArgumentGroup, 

41 config_class: type[BaseModel], 

42 field_name: str, 

43 *, 

44 name: str | None = None, 

45 help: str | None = None, 

46 type: type | None = None, 

47 **kwargs: Any, 

48) -> None: 

49 """Convert pydantic Field to argparse argument definition.""" 

50 field_info = config_class.model_fields[field_name] 

51 if name is None: 

52 name = "--" + field_name.replace("_", "-") 

53 if help is None: 

54 help = field_info.description 

55 

56 if type is None: 

57 type = field_info.annotation 

58 if type in (str, int, float): 

59 kwargs["type"] = type 

60 else: 

61 kwargs["type"] = type 

62 

63 group.add_argument(name, dest=field_name, help=help, **kwargs) 

64 

65 

66# Options for fields in ApdbConfig. 

67def common_apdb_options(parser: argparse.ArgumentParser) -> None: 

68 """Define common configuration options.""" 

69 group = parser.add_argument_group("common APDB options") 

70 _option_from_pydantic_field(group, ApdbConfig, "schema_file", metavar="URL") 

71 _option_from_pydantic_field(group, ApdbConfig, "ss_schema_file", metavar="URL") 

72 _option_from_pydantic_field(group, ApdbConfig, "read_sources_months", metavar="NUMBER") 

73 _option_from_pydantic_field(group, ApdbConfig, "read_forced_sources_months", metavar="NUMBER") 

74 _option_from_pydantic_field(group, ApdbConfig, "enable_replica", action="store_true", default=False) 

75 

76 

77# Options for fields in ApdbSqlConfig, db_url is not included. 

78def sql_config_options(parser: argparse.ArgumentParser) -> None: 

79 """Define SQL backend configuration options.""" 

80 group = parser.add_argument_group("SQL backend options") 

81 _option_from_pydantic_field(group, ApdbSqlConfig, "namespace", metavar="IDENTIFIER") 

82 _option_from_pydantic_field(group, ApdbSqlConnectionConfig, "connection_timeout", metavar="SECONDS") 

83 _option_from_pydantic_field( 

84 group, ApdbSqlConfig, "dia_object_index", choices=["baseline", "pix_id_iov", "last_object_table"] 

85 ) 

86 _option_from_pydantic_field(group, ApdbSqlPixelizationConfig, "htm_level") 

87 _option_from_pydantic_field(group, ApdbSqlPixelizationConfig, "htm_index_column") 

88 _option_from_pydantic_field( 

89 group, 

90 ApdbSqlConfig, 

91 "ra_dec_columns", 

92 help="Names of ra/dec columns in DiaObject table, comma-separated.", 

93 metavar="RA_COLUMN,DEC_COLUMN", 

94 ) 

95 _option_from_pydantic_field(group, ApdbSqlConfig, "prefix") 

96 

97 

98# Options for fields in ApdbCassandraConfig, contact_points is not included. 

99def cassandra_config_options(parser: argparse.ArgumentParser) -> None: 

100 """Define Cassandra backend configuration options.""" 

101 group = parser.add_argument_group("Cassandra backend options") 

102 _option_from_pydantic_field(group, ApdbCassandraConfig, "replica_skips_diaobjects", action="store_true") 

103 _option_from_pydantic_field(group, ApdbCassandraConnectionConfig, "port", metavar="PORT") 

104 _option_from_pydantic_field(group, ApdbCassandraConnectionConfig, "username", metavar="USER") 

105 _option_from_pydantic_field(group, ApdbCassandraConfig, "prefix") 

106 group.add_argument( 

107 "--replication-factor", help="Replication factor used when creating new keyspace.", type=int 

108 ) 

109 group.add_argument( 

110 "--table-options", help="Path or URI of YAML file containing table options.", metavar="URI" 

111 ) 

112 _option_from_pydantic_field( 

113 group, 

114 ApdbCassandraConnectionConfig, 

115 "read_consistency", 

116 choices=["ONE", "TWO", "THREE", "QUORUM", "ALL"], 

117 ) 

118 _option_from_pydantic_field( 

119 group, 

120 ApdbCassandraConnectionConfig, 

121 "write_consistency", 

122 choices=["ONE", "TWO", "THREE", "QUORUM", "ALL"], 

123 ) 

124 _option_from_pydantic_field(group, ApdbCassandraConnectionConfig, "read_timeout", metavar="SECONDS") 

125 _option_from_pydantic_field(group, ApdbCassandraConnectionConfig, "write_timeout", metavar="SECONDS") 

126 _option_from_pydantic_field( 

127 group, 

128 ApdbCassandraConfig, 

129 "ra_dec_columns", 

130 help="Names of ra/dec columns in DiaObject table, comma-separated.", 

131 metavar="RA_COLUMN,DEC_COLUMN", 

132 ) 

133 group = parser.add_argument_group("Cassandra partitioning options") 

134 _option_from_pydantic_field(group, ApdbCassandraPartitioningConfig, "part_pixelization", metavar="NAME") 

135 _option_from_pydantic_field(group, ApdbCassandraPartitioningConfig, "part_pix_level", metavar="LEVEL") 

136 _option_from_pydantic_field( 

137 group, ApdbCassandraPartitioningConfig, "time_partition_tables", action="store_true" 

138 ) 

139 _option_from_pydantic_field( 

140 group, ApdbCassandraPartitioningConfig, "time_partition_start", metavar="TIME" 

141 ) 

142 _option_from_pydantic_field(group, ApdbCassandraPartitioningConfig, "time_partition_end", metavar="TIME")