mohamedawnallah commented on code in PR #34398: URL: https://github.com/apache/beam/pull/34398#discussion_r2112486286
########## sdks/python/apache_beam/transforms/enrichment_handlers/cloudsql.py: ########## @@ -31,43 +34,75 @@ ConditionValueFn = Callable[[beam.Row], list[Any]] -def _validate_cloudsql_metadata( - table_id, - where_clause_template, - where_clause_fields, - where_clause_value_fn, - query_fn): - if query_fn: - if any([table_id, - where_clause_template, - where_clause_fields, - where_clause_value_fn]): +@dataclass +class CustomQueryConfig: + """Configuration for using a custom query function.""" + query_fn: QueryFn + + +@dataclass +class TableFieldsQueryConfig: + """Configuration for using table name, where clause, and field names.""" + table_id: str + where_clause_template: str + where_clause_fields: List[str] + + +@dataclass +class TableFunctionQueryConfig: + """Configuration for using table name, where clause, and a value function.""" + table_id: str + where_clause_template: str + where_clause_value_fn: ConditionValueFn + + +QueryConfig = Union[CustomQueryConfig, + TableFieldsQueryConfig, + TableFunctionQueryConfig] + + +def _validate_query_config(query_config: QueryConfig): + """Validates the provided query configuration.""" + if isinstance(query_config, CustomQueryConfig): Review Comment: Yeah I see they are unnecessary. I don't know that dataclasses are that powerful! Thanks ########## sdks/python/apache_beam/transforms/enrichment_handlers/cloudsql.py: ########## @@ -31,43 +34,75 @@ ConditionValueFn = Callable[[beam.Row], list[Any]] -def _validate_cloudsql_metadata( - table_id, - where_clause_template, - where_clause_fields, - where_clause_value_fn, - query_fn): - if query_fn: - if any([table_id, - where_clause_template, - where_clause_fields, - where_clause_value_fn]): +@dataclass +class CustomQueryConfig: + """Configuration for using a custom query function.""" + query_fn: QueryFn + + +@dataclass +class TableFieldsQueryConfig: + """Configuration for using table name, where clause, and field names.""" + table_id: str + where_clause_template: str + where_clause_fields: List[str] + + +@dataclass +class TableFunctionQueryConfig: + """Configuration for using table name, where clause, and a value function.""" + table_id: str + where_clause_template: str + where_clause_value_fn: ConditionValueFn + + +QueryConfig = Union[CustomQueryConfig, + TableFieldsQueryConfig, + TableFunctionQueryConfig] + + +def _validate_query_config(query_config: QueryConfig): + """Validates the provided query configuration.""" + if isinstance(query_config, CustomQueryConfig): + if not query_config.query_fn: + raise ValueError("CustomQueryConfig must provide a valid query_fn") + elif isinstance(query_config, + (TableFieldsQueryConfig, TableFunctionQueryConfig)): + if not query_config.table_id or not query_config.where_clause_template: raise ValueError( - "Please provide either `query_fn` or the parameters `table_id`, " - "`where_clause_template`, and " - "`where_clause_fields/where_clause_value_fn` together.") + "TableFieldsQueryConfig and " + + "TableFunctionQueryConfig must provide table_id " + + "and where_clause_template") + + is_table_fields = isinstance(query_config, TableFieldsQueryConfig) + if is_table_fields: + table_fields_config = cast(TableFieldsQueryConfig, query_config) + if not table_fields_config.where_clause_fields: + raise ValueError( + "TableFieldsQueryConfig must provide non-empty " + + "where_clause_fields") + + is_table_function = isinstance(query_config, TableFunctionQueryConfig) + if is_table_function: + table_function_config = cast(TableFunctionQueryConfig, query_config) + if not table_function_config.where_clause_value_fn: + raise ValueError( + "TableFunctionQueryConfig must provide " + "where_clause_value_fn") else: - if not (table_id and where_clause_template): - raise ValueError( - "Please provide either `query_fn` or the parameters " - "`table_id` and `where_clause_template` together.") - if (bool(where_clause_fields) == bool(where_clause_value_fn)): - raise ValueError( - "Please provide exactly one of `where_clause_fields` or " - "`where_clause_value_fn`.") + raise ValueError("Invalid query_config type provided") class DatabaseTypeAdapter(Enum): POSTGRESQL = "psycopg2" MYSQL = "pymysql" - SQLSERVER = "pytds" + SQLSERVER = "pymysql" Review Comment: Good catch! Updated it with `pymssql` also I found I need to be explicit with connection config when connecting to the test containers ########## sdks/python/apache_beam/transforms/enrichment_handlers/cloudsql.py: ########## @@ -31,43 +34,75 @@ ConditionValueFn = Callable[[beam.Row], list[Any]] -def _validate_cloudsql_metadata( - table_id, - where_clause_template, - where_clause_fields, - where_clause_value_fn, - query_fn): - if query_fn: - if any([table_id, - where_clause_template, - where_clause_fields, - where_clause_value_fn]): +@dataclass +class CustomQueryConfig: + """Configuration for using a custom query function.""" + query_fn: QueryFn + + +@dataclass +class TableFieldsQueryConfig: + """Configuration for using table name, where clause, and field names.""" + table_id: str + where_clause_template: str + where_clause_fields: List[str] + + +@dataclass +class TableFunctionQueryConfig: + """Configuration for using table name, where clause, and a value function.""" + table_id: str + where_clause_template: str + where_clause_value_fn: ConditionValueFn + + +QueryConfig = Union[CustomQueryConfig, + TableFieldsQueryConfig, + TableFunctionQueryConfig] + + +def _validate_query_config(query_config: QueryConfig): + """Validates the provided query configuration.""" + if isinstance(query_config, CustomQueryConfig): + if not query_config.query_fn: + raise ValueError("CustomQueryConfig must provide a valid query_fn") + elif isinstance(query_config, + (TableFieldsQueryConfig, TableFunctionQueryConfig)): + if not query_config.table_id or not query_config.where_clause_template: raise ValueError( - "Please provide either `query_fn` or the parameters `table_id`, " - "`where_clause_template`, and " - "`where_clause_fields/where_clause_value_fn` together.") + "TableFieldsQueryConfig and " + + "TableFunctionQueryConfig must provide table_id " + + "and where_clause_template") + + is_table_fields = isinstance(query_config, TableFieldsQueryConfig) + if is_table_fields: + table_fields_config = cast(TableFieldsQueryConfig, query_config) + if not table_fields_config.where_clause_fields: + raise ValueError( + "TableFieldsQueryConfig must provide non-empty " + + "where_clause_fields") + + is_table_function = isinstance(query_config, TableFunctionQueryConfig) + if is_table_function: + table_function_config = cast(TableFunctionQueryConfig, query_config) Review Comment: > Why are we casting the query_config to a TableFunctionQueryConfig? I think I did this to fix linting issues in Beam CI but they are not needed 👍 -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: github-unsubscr...@beam.apache.org For queries about this service, please contact Infrastructure at: us...@infra.apache.org