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

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 **kwargs: Any, 

47) -> None: 

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

49 field_info = config_class.model_fields[field_name] 

50 if name is None: 50 ↛ 52line 50 didn't jump to line 52 because the condition on line 50 was always true

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

52 if help is None: 

53 help = field_info.description 

54 if "type" not in kwargs and field_info.annotation in (str, int, float): 

55 kwargs["type"] = field_info.annotation 

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

57 

58 

59# Options for fields in ApdbConfig. 

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

61 """Define common configuration options.""" 

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

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

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

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

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

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

68 

69 

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

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

72 """Define SQL backend configuration options.""" 

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

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

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

76 _option_from_pydantic_field( 

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

78 ) 

79 _option_from_pydantic_field(group, ApdbSqlPixelizationConfig, "htm_level") 

80 _option_from_pydantic_field(group, ApdbSqlPixelizationConfig, "htm_index_column") 

81 _option_from_pydantic_field( 

82 group, 

83 ApdbSqlConfig, 

84 "ra_dec_columns", 

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

86 metavar="RA_COLUMN,DEC_COLUMN", 

87 ) 

88 _option_from_pydantic_field(group, ApdbSqlConfig, "prefix") 

89 

90 

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

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

93 """Define Cassandra backend configuration options.""" 

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

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

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

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

98 _option_from_pydantic_field(group, ApdbCassandraConnectionConfig, "dbauth_alias", metavar="NAME") 

99 _option_from_pydantic_field(group, ApdbCassandraConfig, "prefix") 

100 group.add_argument( 

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

102 ) 

103 group.add_argument( 

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

105 ) 

106 _option_from_pydantic_field( 

107 group, 

108 ApdbCassandraConnectionConfig, 

109 "read_consistency", 

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

111 ) 

112 _option_from_pydantic_field( 

113 group, 

114 ApdbCassandraConnectionConfig, 

115 "write_consistency", 

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

117 ) 

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

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

120 _option_from_pydantic_field( 

121 group, 

122 ApdbCassandraConfig, 

123 "ra_dec_columns", 

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

125 metavar="RA_COLUMN,DEC_COLUMN", 

126 ) 

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

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

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

130 _option_from_pydantic_field( 

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

132 ) 

133 _option_from_pydantic_field( 

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

135 ) 

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