Repository: getmoto/moto (python)

Issue:
[athena] Object of type WorkGroup is not JSON serializable
Hi, latest docker image of motoserver (`motoserver/moto:latest`) uploaded on Dec 09th 2024 contains an issue related to type serialization. 
The issue is not present in previous version `5.0.22`.

When calling Athena endpoint `get_query_execution`, the following error is raised:
```
TypeError: Object of type WorkGroup is not JSON serializable
```

Here's the full traceback:
```
Traceback (most recent call last):
  File "/usr/local/lib/python3.11/site-packages/werkzeug/serving.py", line 370, in run_wsgi
    execute(self.server.app)
  File "/usr/local/lib/python3.11/site-packages/werkzeug/serving.py", line 331, in execute
    application_iter = app(environ, start_response)
                       ^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.11/site-packages/moto/moto_server/werkzeug_app.py", line 261, in __call__
    return backend_app(environ, start_response)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.11/site-packages/flask/app.py", line 1536, in __call__
    return self.wsgi_app(environ, start_response)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.11/site-packages/flask/app.py", line 1514, in wsgi_app
    response = self.handle_exception(e)
               ^^^^^^^^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.11/site-packages/flask_cors/extension.py", line 194, in wrapped_function
    return cors_after_request(app.make_response(f(*args, **kwargs)))
                                                ^^^^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.11/site-packages/flask/app.py", line 1511, in wsgi_app
    response = self.full_dispatch_request()
               ^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.11/site-packages/flask/app.py", line 919, in full_dispatch_request
    rv = self.handle_user_exception(e)
         ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.11/site-packages/flask_cors/extension.py", line 194, in wrapped_function
    return cors_after_request(app.make_response(f(*args, **kwargs)))
                                                ^^^^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.11/site-packages/flask/app.py", line 917, in full_dispatch_request
    rv = self.dispatch_request()
         ^^^^^^^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.11/site-packages/flask/app.py", line 902, in dispatch_request
    return self.ensure_sync(self.view_functions[rule.endpoint])(**view_args)  *** type: ignore[no-any-return]
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.11/site-packages/moto/core/utils.py", line 112, in __call__
    result = self.callback(request, request.url, dict(request.headers))
             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.11/site-packages/moto/core/responses.py", line [291](https://github.com/docker/saas-mega/actions/runs/12242202016/job/34149046626?pr=14868#step:21:292), in dispatch
    return cls()._dispatch(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.11/site-packages/moto/core/responses.py", line 503, in _dispatch
    return self.call_action()
           ^^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.11/site-packages/moto/core/responses.py", line 591, in call_action
    response = method()
               ^^^^^^^^
  File "/usr/local/lib/python3.11/site-packages/moto/athena/responses.py", line 99, in get_query_execution
    return json.dumps(result)
           ^^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.11/json/__init__.py", line 231, in dumps
    return _default_encoder.encode(obj)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.11/json/encoder.py", line 200, in encode
    chunks = self.iterencode(o, _one_shot=True)
             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.11/json/encoder.py", line 258, in iterencode
    return _iterencode(o, 0)
           ^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.11/json/encoder.py", line 180, in default
    raise TypeError(f'Object of type {o.__class__.__name__} '
TypeError: Object of type WorkGroup is not JSON serializable
```

Pull request description:
Athena: start_query_execution() should support WorkGroups
Fixes #8388 

Patch:
diff --git a/moto/athena/responses.py b/moto/athena/responses.py
index 23b074160..727711167 100644
--- a/moto/athena/responses.py
+++ b/moto/athena/responses.py
@@ -89,7 +89,7 @@ class AthenaResponse(BaseResponse):
                     "QueryPlanningTimeInMillis": 0,
                     "ServiceProcessingTimeInMillis": 0,
                 },
-                "WorkGroup": execution.workgroup,
+                "WorkGroup": execution.workgroup.name if execution.workgroup else None,
             }
         }
         if execution.execution_parameters is not None:
diff --git a/moto/es/models.py b/moto/es/models.py
index 6c7631752..4cc65d6fb 100644
--- a/moto/es/models.py
+++ b/moto/es/models.py
@@ -142,5 +142,14 @@ class ElasticsearchServiceBackend(BaseBackend):
         """
         return [{"DomainName": domain.domain_name} for domain in self.domains.values()]
 
+    def describe_elasticsearch_domains(
+        self, domain_names: List[str]
+    ) -> List[Dict[str, Any]]:
+        queried_domains = []
+        for domain_name in domain_names:
+            if domain_name in self.domains:
+                queried_domains.append(self.domains[domain_name].to_json())
+        return queried_domains
+
 
 es_backends = BackendDict(ElasticsearchServiceBackend, "es")
diff --git a/moto/es/responses.py b/moto/es/responses.py
index d8e7ac813..35293b33d 100644
--- a/moto/es/responses.py
+++ b/moto/es/responses.py
@@ -13,7 +13,7 @@ class ElasticsearchServiceResponse(BaseResponse):
     """Handler for ElasticsearchService requests and responses."""
 
     def __init__(self) -> None:
-        super().__init__(service_name="elasticsearch")
+        super().__init__(service_name="es")
 
     @property
     def es_backend(self) -> ElasticsearchServiceBackend:
@@ -98,3 +98,10 @@ class ElasticsearchServiceResponse(BaseResponse):
     def list_domain_names(self) -> TYPE_RESPONSE:
         domain_names = self.es_backend.list_domain_names()
         return 200, {}, json.dumps({"DomainNames": domain_names})
+
+    def describe_elasticsearch_domains(self) -> TYPE_RESPONSE:
+        domain_names = self._get_param("DomainNames")
+        domain_list = self.es_backend.describe_elasticsearch_domains(
+            domain_names=domain_names,
+        )
+        return 200, {}, json.dumps({"DomainStatusList": domain_list})
diff --git a/moto/es/urls.py b/moto/es/urls.py
index d125e9a09..fbee7fcb7 100644
--- a/moto/es/urls.py
+++ b/moto/es/urls.py
@@ -11,11 +11,13 @@ url_paths = {
     "{0}/2015-01-01/domain$": ElasticsearchServiceResponse.list_domains,
     "{0}/2015-01-01/es/domain$": ElasticsearchServiceResponse.domains,
     "{0}/2015-01-01/es/domain/(?P<domainname>[^/]+)": ElasticsearchServiceResponse.domain,
+    "{0}/2015-01-01/es/domain-info$": ElasticsearchServiceResponse.dispatch,
     "{0}/2021-01-01/domain$": OpenSearchServiceResponse.dispatch,
     "{0}/2021-01-01/opensearch/compatibleVersions": OpenSearchServiceResponse.dispatch,
     "{0}/2021-01-01/opensearch/domain": OpenSearchServiceResponse.dispatch,
     "{0}/2021-01-01/opensearch/domain/(?P<domainname>[^/]+)": OpenSearchServiceResponse.dispatch,
     "{0}/2021-01-01/opensearch/domain/(?P<domainname>[^/]+)/config": OpenSearchServiceResponse.dispatch,
+    "{0}/2021-01-01/opensearch/domain-info$": OpenSearchServiceResponse.dispatch,
     "{0}/2021-01-01/tags/?": OpenSearchServiceResponse.dispatch,
     "{0}/2021-01-01/tags-removal/": OpenSearchServiceResponse.dispatch,
 }
diff --git a/moto/opensearch/models.py b/moto/opensearch/models.py
index 31c91821d..2e7242983 100644
--- a/moto/opensearch/models.py
+++ b/moto/opensearch/models.py
@@ -375,5 +375,12 @@ class OpenSearchServiceBackend(BaseBackend):
                 )
         return domains
 
+    def describe_domains(self, domain_names: List[str]) -> List[OpenSearchDomain]:
+        queried_domains = []
+        for domain_name in domain_names:
+            if domain_name in self.domains:
+                queried_domains.append(self.domains[domain_name])
+        return queried_domains
+
 
 opensearch_backends = BackendDict(OpenSearchServiceBackend, "opensearch")
diff --git a/moto/opensearch/responses.py b/moto/opensearch/responses.py
index e8f102e14..4034a16b2 100644
--- a/moto/opensearch/responses.py
+++ b/moto/opensearch/responses.py
@@ -147,3 +147,11 @@ class OpenSearchServiceResponse(BaseResponse):
             engine_type=engine_type,
         )
         return json.dumps(dict(DomainNames=domain_names))
+
+    def describe_domains(self) -> str:
+        domain_names = self._get_param("DomainNames")
+        domains = self.opensearch_backend.describe_domains(
+            domain_names=domain_names,
+        )
+        domain_list = [domain.to_dict() for domain in domains]
+        return json.dumps({"DomainStatusList": domain_list})
diff --git a/moto/opensearch/urls.py b/moto/opensearch/urls.py
index 0d7d1a4cf..c62d8c04c 100644
--- a/moto/opensearch/urls.py
+++ b/moto/opensearch/urls.py
@@ -8,4 +8,6 @@ url_bases = [r"https?://es\.(.+)\.amazonaws\.com"]
 response = OpenSearchServiceResponse()
 
 
-url_paths = {"{0}/.*$": response.dispatch}
+url_paths = {
+    "{0}/.*$": response.dispatch,
+}
diff --git a/moto/rds/models.py b/moto/rds/models.py
index c4825b6a9..8aae7e86d 100644
--- a/moto/rds/models.py
+++ b/moto/rds/models.py
@@ -1890,6 +1890,10 @@ class RDSBackend(BaseBackend):
         tags: Optional[List[Dict[str, str]]] = None,
         copy_tags: bool = False,
     ) -> DBSnapshot:
+        if source_snapshot_identifier.startswith("arn:aws:rds:"):
+            source_snapshot_identifier = self.extract_snapshot_name_from_arn(
+                source_snapshot_identifier
+            )
         if source_snapshot_identifier not in self.database_snapshots:
             raise DBSnapshotNotFoundError(source_snapshot_identifier)
 
@@ -2010,9 +2014,28 @@ class RDSBackend(BaseBackend):
     def reboot_db_instance(self, db_instance_identifier: str) -> DBInstance:
         return self.describe_db_instances(db_instance_identifier)[0]
 
+    def extract_snapshot_name_from_arn(self, snapshot_arn: str) -> str:
+        arn_breakdown = snapshot_arn.split(":")
+        region_name, account_id, resource_type, snapshot_name = arn_breakdown[3:7]
+        if resource_type != "snapshot":
+            raise InvalidParameterValue(
+                "The parameter SourceDBSnapshotIdentifier is not a valid identifier. "
+                "Identifiers must begin with a letter; must contain only ASCII "
+                "letters, digits, and hyphens; and must not end with a hyphen or "
+                "contain two consecutive hyphens."
+            )
+        if region_name != self.region_name or account_id != self.account_id:
+            raise NotImplementedError(
+                "Cross account/region snapshot handling is not yet implemented in moto."
+            )
+        return snapshot_name
+
     def restore_db_instance_from_db_snapshot(
         self, from_snapshot_id: str, overrides: Dict[str, Any]
     ) -> DBInstance:
+        if from_snapshot_id.startswith("arn:aws:rds:"):
+            from_snapshot_id = self.extract_snapshot_name_from_arn(from_snapshot_id)
+
         snapshot = self.describe_db_snapshots(
             db_instance_identifier=None, db_snapshot_identifier=from_snapshot_id
         )[0]


Test patch:
diff --git a/tests/test_athena/test_athena.py b/tests/test_athena/test_athena.py
index e539c42b1..1e75ebc86 100644
--- a/tests/test_athena/test_athena.py
+++ b/tests/test_athena/test_athena.py
@@ -237,6 +237,29 @@ def test_stop_query_execution():
     assert details["Status"]["State"] == "CANCELLED"
 
 
+@mock_aws
+def test_start_execution_with_workgroup():
+    client = boto3.client("athena", region_name="us-east-1")
+
+    client.create_work_group(
+        Name="myworkgroup",
+        Description="Test work group",
+        Configuration={
+            "ResultConfiguration": {"OutputLocation": "s3://bucket-name/prefix/"}
+        },
+    )
+
+    exec_id = client.start_query_execution(
+        QueryString="SELECT stuff",
+        QueryExecutionContext={"Database": "database"},
+        ResultConfiguration={"OutputLocation": "s3://bucket-name/prefix/"},
+        WorkGroup="myworkgroup",
+    )["QueryExecutionId"]
+
+    execution = client.get_query_execution(QueryExecutionId=exec_id)["QueryExecution"]
+    assert execution["WorkGroup"] == "myworkgroup"
+
+
 @mock_aws
 def test_create_named_query():
     client = boto3.client("athena", region_name="us-east-1")
diff --git a/tests/test_es/test_es.py b/tests/test_es/test_es.py
index 3ee124388..259f444da 100644
--- a/tests/test_es/test_es.py
+++ b/tests/test_es/test_es.py
@@ -215,3 +215,28 @@ def test_list_domain_names_with_multiple_domains():
     assert len(resp["DomainNames"]) == 4
     for name in domain_names:
         assert {"DomainName": name} in resp["DomainNames"]
+
+
+@mock_aws
+def test_describe_elasticsearch_domains():
+    client = boto3.client("es", region_name="us-east-1")
+    domain_names = [f"env{i}" for i in range(1, 5)]
+    for name in domain_names:
+        client.create_elasticsearch_domain(
+            DomainName=name,
+            ElasticsearchVersion="7.10",
+            AdvancedOptions={"option": "value"},
+            AdvancedSecurityOptions={"Enabled": False},
+        )
+    resp = client.describe_elasticsearch_domains(DomainNames=domain_names)
+
+    assert len(resp["DomainStatusList"]) == 4
+    for domain in resp["DomainStatusList"]:
+        assert domain["DomainName"] in domain_names
+        assert domain["ElasticsearchVersion"] == "7.10"
+        assert "AdvancedSecurityOptions" in domain.keys()
+        assert "AdvancedOptions" in domain.keys()
+
+    # Test for invalid domain name
+    resp = client.describe_elasticsearch_domains(DomainNames=["invalid"])
+    assert len(resp["DomainStatusList"]) == 0
diff --git a/tests/test_opensearch/test_opensearch.py b/tests/test_opensearch/test_opensearch.py
index 1ce66fef1..22955de1b 100644
--- a/tests/test_opensearch/test_opensearch.py
+++ b/tests/test_opensearch/test_opensearch.py
@@ -217,3 +217,23 @@ def test_list_unknown_domain_names_engine_type():
     err = exc.value.response["Error"]
     assert err["Code"] == "EngineTypeNotFoundException"
     assert err["Message"] == "Engine Type not found: testdn"
+
+
+@mock_aws
+def test_describe_domains():
+    client = boto3.client("opensearch", region_name="us-east-1")
+    domain_names = [f"env{i}" for i in range(1, 5)]
+    opensearch_engine_version = "OpenSearch_1.0"
+    for name in domain_names:
+        client.create_domain(DomainName=name, EngineVersion=opensearch_engine_version)
+    resp = client.describe_domains(DomainNames=domain_names)
+
+    assert len(resp["DomainStatusList"]) == 4
+    for domain in resp["DomainStatusList"]:
+        assert domain["DomainName"] in domain_names
+        assert "AdvancedSecurityOptions" in domain.keys()
+        assert "AdvancedOptions" in domain.keys()
+
+    # Test for invalid domain name
+    resp = client.describe_domains(DomainNames=["invalid"])
+    assert len(resp["DomainStatusList"]) == 0
diff --git a/tests/test_rds/test_rds.py b/tests/test_rds/test_rds.py
index 2c7fc5fd1..36dddcadf 100644
--- a/tests/test_rds/test_rds.py
+++ b/tests/test_rds/test_rds.py
@@ -1141,8 +1141,13 @@ def test_create_db_snapshots_with_tags():
 
 
 @pytest.mark.parametrize("delete_db_instance", [True, False])
+@pytest.mark.parametrize(
+    "db_snapshot_identifier",
+    ("snapshot-1", f"arn:aws:rds:{DEFAULT_REGION}:123456789012:snapshot:snapshot-1"),
+    ids=("by_name", "by_arn"),
+)
 @mock_aws
-def test_copy_db_snapshots(delete_db_instance: bool):
+def test_copy_db_snapshots(delete_db_instance: bool, db_snapshot_identifier: str):
     conn = boto3.client("rds", region_name=DEFAULT_REGION)
 
     conn.create_db_instance(
@@ -1166,7 +1171,8 @@ def test_copy_db_snapshots(delete_db_instance: bool):
         conn.delete_db_instance(DBInstanceIdentifier="db-primary-1")
 
     target_snapshot = conn.copy_db_snapshot(
-        SourceDBSnapshotIdentifier="snapshot-1", TargetDBSnapshotIdentifier="snapshot-2"
+        SourceDBSnapshotIdentifier=db_snapshot_identifier,
+        TargetDBSnapshotIdentifier="snapshot-2",
     ).get("DBSnapshot")
 
     assert target_snapshot.get("Engine") == "postgres"
@@ -1176,6 +1182,21 @@ def test_copy_db_snapshots(delete_db_instance: bool):
     assert result["TagList"] == []
 
 
+@mock_aws
+def test_copy_db_snapshot_invalid_arns():
+    conn = boto3.client("rds", region_name=DEFAULT_REGION)
+
+    invalid_arn = (
+        f"arn:aws:rds:{DEFAULT_REGION}:123456789012:this-is-not-a-snapshot:snapshot-1"
+    )
+    with pytest.raises(ClientError) as ex:
+        conn.copy_db_snapshot(
+            SourceDBSnapshotIdentifier=invalid_arn,
+            TargetDBSnapshotIdentifier="snapshot-2",
+        )
+    assert "is not a valid identifier" in ex.value.response["Error"]["Message"]
+
+
 original_snapshot_tags = [{"Key": "original", "Value": "snapshot tags"}]
 new_snapshot_tags = [{"Key": "new", "Value": "tag"}]
 
@@ -1321,11 +1342,18 @@ def test_delete_db_snapshot():
         conn.describe_db_snapshots(DBSnapshotIdentifier="snapshot-1")
 
 
+@pytest.mark.parametrize(
+    "db_snapshot_identifier",
+    ("snapshot-1", f"arn:aws:rds:{DEFAULT_REGION}:123456789012:snapshot:snapshot-1"),
+    ids=("by_name", "by_arn"),
+)
 @pytest.mark.parametrize(
     "custom_db_subnet_group", [True, False], ids=("custom_subnet", "default_subnet")
 )
 @mock_aws
-def test_restore_db_instance_from_db_snapshot(custom_db_subnet_group: bool):
+def test_restore_db_instance_from_db_snapshot(
+    db_snapshot_identifier: str, custom_db_subnet_group: bool
+):
     conn = boto3.client("rds", region_name=DEFAULT_REGION)
     conn.create_db_instance(
         DBInstanceIdentifier="db-primary-1",
@@ -1353,9 +1381,12 @@ def test_restore_db_instance_from_db_snapshot(custom_db_subnet_group: bool):
         )
 
     # restore
+    new_instance = conn.restore_db_instance_from_db_snapshot(
+        DBInstanceIdentifier="db-restore-1", DBSnapshotIdentifier=db_snapshot_identifier
+    )["DBInstance"]
     kwargs = {
         "DBInstanceIdentifier": "db-restore-1",
-        "DBSnapshotIdentifier": "snapshot-1",
+        "DBSnapshotIdentifier": db_snapshot_identifier,
     }
     if custom_db_subnet_group:
         kwargs["DBSubnetGroupName"] = "custom-subnet-group"